Skip to content

Commit

Permalink
Merge pull request #58 from BloodHoundAD/close-idle-conns
Browse files Browse the repository at this point in the history 
chore: close idle connections after a successful collection
  • Loading branch information
@superlinkx
superlinkx committed Oct 24, 2023
2 parents 207d9ad + c239f88 commit 2ffd8e9
Show file tree
Hide file tree
Showing 48 changed files with 102 additions and 59 deletions.
12 changes: 8 additions & 4 deletions client/client.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,12 @@

package client

//go:generate go run github.com/golang/mock/mockgen -destination=./mocks/client.go -package=mocks . AzureClient
//go:generate go run go.uber.org/mock/mockgen -destination=./mocks/client.go -package=mocks . AzureClient

import (
"context"
"fmt"

"encoding/json"
"fmt"

"github.com/bloodhoundad/azurehound/v2/client/config"
"github.com/bloodhoundad/azurehound/v2/client/rest"
Expand All @@ -36,7 +35,6 @@ func NewClient(config config.Config) (AzureClient, error) {
} else if resourceManager, err := rest.NewRestClient(config.ResourceManagerUrl(), config); err != nil {
return nil, err
} else {

if config.JWT != "" {
if aud, err := rest.ParseAud(config.JWT); err != nil {
return nil, err
Expand Down Expand Up @@ -98,6 +96,11 @@ func (s azureClient) TenantInfo() azure.Tenant {
return s.tenant
}

func (s azureClient) CloseIdleConnections() {
s.msgraph.CloseIdleConnections()
s.resourceManager.CloseIdleConnections()
}

type AzureClient interface {
GetAzureADApp(ctx context.Context, objectId string, selectCols []string) (*azure.Application, error)
GetAzureADApps(ctx context.Context, filter, search, orderBy, expand string, selectCols []string, top int32, count bool) (azure.ApplicationList, error)
Expand Down Expand Up @@ -165,4 +168,5 @@ type AzureClient interface {
ListRoleAssignmentsForResource(ctx context.Context, resourceId string, filter string) <-chan azure.RoleAssignmentResult
ListAzureADAppRoleAssignments(ctx context.Context, servicePrincipal, filter, search, orderBy, expand string, selectCols []string) <-chan azure.AppRoleAssignmentResult
TenantInfo() azure.Tenant
CloseIdleConnections()
}
14 changes: 13 additions & 1 deletion client/mocks/client.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

7 changes: 6 additions & 1 deletion client/rest/client.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

package rest

//go:generate go run github.com/golang/mock/mockgen -destination=./mocks/client.go -package=mocks . RestClient
//go:generate go run go.uber.org/mock/mockgen -destination=./mocks/client.go -package=mocks . RestClient

import (
"bytes"
Expand All @@ -44,6 +44,7 @@ type RestClient interface {
Post(ctx context.Context, path string, body interface{}, params, headers map[string]string) (*http.Response, error)
Put(ctx context.Context, path string, body interface{}, params, headers map[string]string) (*http.Response, error)
Send(req *http.Request) (*http.Response, error)
CloseIdleConnections()
}

func NewRestClient(apiUrl string, config config.Config) (RestClient, error) {
Expand Down Expand Up @@ -289,3 +290,7 @@ func (s *restClient) send(req *http.Request) (*http.Response, error) {
return nil, fmt.Errorf("unable to complete the request after %d attempts: %w", maxRetries, err)
}
}

func (s *restClient) CloseIdleConnections() {
s.http.CloseIdleConnections()
}
14 changes: 13 additions & 1 deletion client/rest/mocks/client.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion cmd/list-app-owners_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/enums"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-apps_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import (

"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-device-owners_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-devices_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import (

"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-group-members_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-group-owners_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-groups_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import (

"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-key-vault-access-policies_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-key-vault-contributors_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/enums"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-key-vault-kvcontributors_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/enums"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-key-vault-owners_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/enums"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-key-vault-role-assignments_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/constants"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-key-vault-user-access-admins_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/enums"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-key-vaults_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-management-group-descendants_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-management-group-owners_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/enums"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-management-group-role-assignments_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/constants"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-management-group-user-access-admins_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/enums"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-management-groups_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import (

"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-resource-group-owners_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/enums"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-resource-group-role-assignments_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/constants"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-resource-group-user-access-admins_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/enums"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-resource-groups_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-role-assignments_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ package cmd
import (
"testing"

"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-roles_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import (

"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-service-principal-owners_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-service-principals_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import (

"github.com/bloodhoundad/azurehound/v2/client/mocks"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-subscription-owners_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/constants"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-subscription-role-assignments_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/constants"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
2 changes: 1 addition & 1 deletion cmd/list-subscription-user-access-admins_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ import (
"github.com/bloodhoundad/azurehound/v2/constants"
"github.com/bloodhoundad/azurehound/v2/models"
"github.com/bloodhoundad/azurehound/v2/models/azure"
"github.com/golang/mock/gomock"
"go.uber.org/mock/gomock"
)

func init() {
Expand Down
Loading

0 comments on commit 2ffd8e9

Please sign in to comment.