refactor: introduce map and filter functions to reduce code duplication

SpecterOps · Sep 22, 2022 · a4bf0fb · a4bf0fb
1 parent 5a9cd7b
commit a4bf0fb
Show file tree

Hide file tree

Showing 37 changed files with 448 additions and 542 deletions.
diff --git a/cmd/list-app-owners.go b/cmd/list-app-owners.go
@@ -91,11 +91,11 @@ func listAppOwners(ctx context.Context, client client.AzureClient, apps <-chan i
 			for id := range stream {
 				var (
 					data = models.AppOwners{
-						AppId: id.(string),
+						AppId: id,
 					}
 					count = 0
 				)
-				for item := range client.ListAzureADAppOwners(ctx, id.(string), "", "", "", nil) {
+				for item := range client.ListAzureADAppOwners(ctx, id, "", "", "", nil) {
 					if item.Error != nil {
 						log.Error(item.Error, "unable to continue processing owners for this app", "appId", id)
 					} else {

diff --git a/cmd/list-azure-rm.go b/cmd/list-azure-rm.go
@@ -26,6 +26,7 @@ import (
 
 	"github.com/bloodhoundad/azurehound/client"
 	"github.com/bloodhoundad/azurehound/enums"
+	"github.com/bloodhoundad/azurehound/models"
 	"github.com/bloodhoundad/azurehound/pipeline"
 	"github.com/spf13/cobra"
 )
@@ -70,7 +71,11 @@ func listAllRM(ctx context.Context, client client.AzureClient) <-chan interface{
 		keyVaults  = make(chan interface{})
 		keyVaults2 = make(chan interface{})
 		keyVaults3 = make(chan interface{})
-		keyVaults4 = make(chan interface{})
+
+		keyVaultRoleAssignments1 = make(chan azureWrapper[models.KeyVaultRoleAssignments])
+		keyVaultRoleAssignments2 = make(chan azureWrapper[models.KeyVaultRoleAssignments])
+		keyVaultRoleAssignments3 = make(chan azureWrapper[models.KeyVaultRoleAssignments])
+		keyVaultRoleAssignments4 = make(chan azureWrapper[models.KeyVaultRoleAssignments])
 
 		mgmtGroups  = make(chan interface{})
 		mgmtGroups2 = make(chan interface{})
@@ -102,10 +107,14 @@ func listAllRM(ctx context.Context, client client.AzureClient) <-chan interface{
 	subscriptionUserAccessAdmins := listSubscriptionUserAccessAdmins(ctx, client, subscriptions6)
 
 	// Enumerate KeyVaults, KeyVaultOwners, KeyVaultAccessPolicies and KeyVaultUserAccessAdmins
-	pipeline.Tee(ctx.Done(), listKeyVaults(ctx, client, subscriptions2), keyVaults, keyVaults2, keyVaults3, keyVaults4)
-	keyVaultOwners := listKeyVaultOwners(ctx, client, keyVaults2)
+	pipeline.Tee(ctx.Done(), listKeyVaults(ctx, client, subscriptions2), keyVaults, keyVaults2, keyVaults3)
+	pipeline.Tee(ctx.Done(), listKeyVaultRoleAssignments(ctx, client, keyVaults2), keyVaultRoleAssignments1, keyVaultRoleAssignments2, keyVaultRoleAssignments3, keyVaultRoleAssignments4)
 	keyVaultAccessPolicies := listKeyVaultAccessPolicies(ctx, client, keyVaults3, []enums.KeyVaultAccessType{enums.GetCerts, enums.GetKeys, enums.GetCerts})
-	keyVaultUserAccessAdmins := listKeyVaultUserAccessAdmins(ctx, client, keyVaults4)
+
+	keyVaultOwners := listKeyVaultOwners(ctx, keyVaultRoleAssignments1)
+	keyVaultUserAccessAdmins := listKeyVaultUserAccessAdmins(ctx, keyVaultRoleAssignments2)
+	keyVaultContributors := listKeyVaultContributors(ctx, keyVaultRoleAssignments3)
+	keyVaultKVContributors := listKeyVaultKVContributors(ctx, keyVaultRoleAssignments4)
 
 	// Enumerate ManagementGroups, ManagementGroupOwners and ManagementGroupDescendants
 	pipeline.Tee(ctx.Done(), listManagementGroups(ctx, client), mgmtGroups, mgmtGroups2, mgmtGroups3, mgmtGroups4)
@@ -129,6 +138,8 @@ func listAllRM(ctx context.Context, client client.AzureClient) <-chan interface{
 
 	return pipeline.Mux(ctx.Done(),
 		keyVaultAccessPolicies,
+		keyVaultContributors,
+		keyVaultKVContributors,
 		keyVaultOwners,
 		keyVaultUserAccessAdmins,
 		keyVaults,

diff --git a/cmd/list-device-owners.go b/cmd/list-device-owners.go
@@ -90,11 +90,11 @@ func listDeviceOwners(ctx context.Context, client client.AzureClient, devices <-
 			for id := range stream {
 				var (
 					data = models.DeviceOwners{
-						DeviceId: id.(string),
+						DeviceId: id,
 					}
 					count = 0
 				)
-				for item := range client.ListAzureDeviceRegisteredOwners(ctx, id.(string), false) {
+				for item := range client.ListAzureDeviceRegisteredOwners(ctx, id, false) {
 					if item.Error != nil {
 						log.Error(item.Error, "unable to continue processing owners for this device", "deviceId", id)
 					} else {

diff --git a/cmd/list-group-members.go b/cmd/list-group-members.go
@@ -91,11 +91,11 @@ func listGroupMembers(ctx context.Context, client client.AzureClient, groups <-c
 			for id := range stream {
 				var (
 					data = models.GroupMembers{
-						GroupId: id.(string),
+						GroupId: id,
 					}
 					count = 0
 				)
-				for item := range client.ListAzureADGroupMembers(ctx, id.(string), "", "", "", nil) {
+				for item := range client.ListAzureADGroupMembers(ctx, id, "", "", "", nil) {
 					if item.Error != nil {
 						log.Error(item.Error, "unable to continue processing members for this group", "groupId", id)
 					} else {

diff --git a/cmd/list-group-owners.go b/cmd/list-group-owners.go
@@ -91,11 +91,11 @@ func listGroupOwners(ctx context.Context, client client.AzureClient, groups <-ch
 			for id := range stream {
 				var (
 					groupOwners = models.GroupOwners{
-						GroupId: id.(string),
+						GroupId: id,
 					}
 					count = 0
 				)
-				for item := range client.ListAzureADGroupOwners(ctx, id.(string), "", "", "", nil) {
+				for item := range client.ListAzureADGroupOwners(ctx, id, "", "", "", nil) {
 					if item.Error != nil {
 						log.Error(item.Error, "unable to continue processing owners for this group", "groupId", id)
 					} else {

diff --git a/cmd/list-key-vault-contributors.go b/cmd/list-key-vault-contributors.go
@@ -19,15 +19,13 @@ package cmd
 
 import (
 	"context"
-	"fmt"
 	"os"
 	"os/signal"
-	"path"
 	"time"
 
-	"github.com/bloodhoundad/azurehound/client"
 	"github.com/bloodhoundad/azurehound/constants"
 	"github.com/bloodhoundad/azurehound/enums"
+	"github.com/bloodhoundad/azurehound/internal"
 	"github.com/bloodhoundad/azurehound/models"
 	"github.com/bloodhoundad/azurehound/pipeline"
 	"github.com/spf13/cobra"
@@ -59,52 +57,30 @@ func listKeyVaultContributorsCmdImpl(cmd *cobra.Command, args []string) {
 		subscriptions := listSubscriptions(ctx, azClient)
 		keyVaults := listKeyVaults(ctx, azClient, subscriptions)
 		kvRoleAssignments := listKeyVaultRoleAssignments(ctx, azClient, keyVaults)
-		stream := listKeyVaultContributors(ctx, azClient, kvRoleAssignments)
+		stream := listKeyVaultContributors(ctx, kvRoleAssignments)
 		outputStream(ctx, stream)
 		duration := time.Since(start)
 		log.Info("collection completed", "duration", duration.String())
 	}
 }
 
-func listKeyVaultContributors(ctx context.Context, client client.AzureClient, vmRoleAssignments <-chan interface{}) <-chan interface{} {
-	out := make(chan interface{})
+func listKeyVaultContributors(
+	ctx context.Context,
+	kvRoleAssignments <-chan azureWrapper[models.KeyVaultRoleAssignments],
+) <-chan any {
+	return pipeline.Map(ctx.Done(), kvRoleAssignments, func(ra azureWrapper[models.KeyVaultRoleAssignments]) any {
+		filteredAssignments := internal.Filter(ra.Data.RoleAssignments, kvRoleAssignmentFilter(constants.ContributorRoleID))
 
-	go func() {
-		defer close(out)
-
-		for result := range pipeline.OrDone(ctx.Done(), vmRoleAssignments) {
-			if roleAssignments, ok := result.(AzureWrapper).Data.(models.KeyVaultRoleAssignments); !ok {
-				log.Error(fmt.Errorf("failed type assertion"), "unable to continue enumerating key vault contributors", "result", result)
-				return
-			} else {
-				var (
-					keyVaultContributors = models.KeyVaultContributors{
-						KeyVaultId: roleAssignments.KeyVaultId,
-					}
-					count = 0
-				)
-				for _, item := range roleAssignments.RoleAssignments {
-					roleDefinitionId := path.Base(item.RoleAssignment.Properties.RoleDefinitionId)
-
-					if roleDefinitionId == constants.ContributorRoleID {
-						keyVaultContributor := models.KeyVaultContributor{
-							Contributor: item.RoleAssignment,
-							KeyVaultId:  item.KeyVaultId,
-						}
-						log.V(2).Info("found key vault contributor", "keyVaultContributor", keyVaultContributor)
-						count++
-						keyVaultContributors.Contributors = append(keyVaultContributors.Contributors, keyVaultContributor)
-					}
-				}
-				out <- AzureWrapper{
-					Kind: enums.KindAZVMContributor,
-					Data: keyVaultContributors,
-				}
-				log.V(1).Info("finished listing key vault contributors", "keyVaultId", roleAssignments.KeyVaultId, "count", count)
+		contributors := internal.Map(filteredAssignments, func(ra models.KeyVaultRoleAssignment) models.KeyVaultContributor {
+			return models.KeyVaultContributor{
+				ra.RoleAssignment,
+				ra.KeyVaultId,
 			}
-		}
-		log.Info("finished listing all key vault contributors")
-	}()
+		})
 
-	return out
+		return NewAzureWrapper(enums.KindAZKeyVaultContributor, models.KeyVaultContributors{
+			KeyVaultId:   ra.Data.KeyVaultId,
+			Contributors: contributors,
+		})
+	})
 }
diff --git a/cmd/list-key-vault-contributors_test.go b/cmd/list-key-vault-contributors_test.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/bloodhoundad/azurehound/client/mocks"
 	"github.com/bloodhoundad/azurehound/constants"
+	"github.com/bloodhoundad/azurehound/enums"
 	"github.com/bloodhoundad/azurehound/models"
 	"github.com/bloodhoundad/azurehound/models/azure"
 	"github.com/golang/mock/gomock"
@@ -39,37 +40,31 @@ func TestListKeyVaultContributors(t *testing.T) {
 
 	mockClient := mocks.NewMockAzureClient(ctrl)
 
-	mockRoleAssignmentsChannel := make(chan interface{})
+	mockRoleAssignmentsChannel := make(chan azureWrapper[models.KeyVaultRoleAssignments])
 	mockTenant := azure.Tenant{}
 	mockClient.EXPECT().TenantInfo().Return(mockTenant).AnyTimes()
 	channel := listKeyVaultContributors(ctx, mockClient, mockRoleAssignmentsChannel)
 
 	go func() {
 		defer close(mockRoleAssignmentsChannel)
 
-		mockRoleAssignmentsChannel <- AzureWrapper{
-			Data: models.KeyVaultRoleAssignments{
-				KeyVaultId: "foo",
-				RoleAssignments: []models.KeyVaultRoleAssignment{
-					{
-						RoleAssignment: azure.RoleAssignment{
-							Name: constants.ContributorRoleID,
-							Properties: azure.RoleAssignmentPropertiesWithScope{
-								RoleDefinitionId: constants.ContributorRoleID,
-							},
+		mockRoleAssignmentsChannel <- NewAzureWrapper(enums.KindAZKeyVaultRoleAssignment, models.KeyVaultRoleAssignments{
+			KeyVaultId: "foo",
+			RoleAssignments: []models.KeyVaultRoleAssignment{
+				{
+					RoleAssignment: azure.RoleAssignment{
+						Name: constants.ContributorRoleID,
+						Properties: azure.RoleAssignmentPropertiesWithScope{
+							RoleDefinitionId: constants.ContributorRoleID,
 						},
 					},
 				},
 			},
-		}
+		})
 	}()
 
-	if result, ok := <-channel; !ok {
+	if _, ok := <-channel; !ok {
 		t.Fatalf("failed to receive from channel")
-	} else if wrapper, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	} else if _, ok := wrapper.Data.(models.KeyVaultContributors); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.KeyVaultContributors{})
 	}
 
 	if _, ok := <-channel; ok {

diff --git a/cmd/list-key-vault-kvcontributors.go b/cmd/list-key-vault-kvcontributors.go
@@ -19,15 +19,13 @@ package cmd
 
 import (
 	"context"
-	"fmt"
 	"os"
 	"os/signal"
-	"path"
 	"time"
 
-	"github.com/bloodhoundad/azurehound/client"
 	"github.com/bloodhoundad/azurehound/constants"
 	"github.com/bloodhoundad/azurehound/enums"
+	"github.com/bloodhoundad/azurehound/internal"
 	"github.com/bloodhoundad/azurehound/models"
 	"github.com/bloodhoundad/azurehound/pipeline"
 	"github.com/spf13/cobra"
@@ -59,52 +57,30 @@ func listKeyVaultKVContributorsCmdImpl(cmd *cobra.Command, args []string) {
 		subscriptions := listSubscriptions(ctx, azClient)
 		keyVaults := listKeyVaults(ctx, azClient, subscriptions)
 		kvRoleAssignments := listKeyVaultRoleAssignments(ctx, azClient, keyVaults)
-		stream := listKeyVaultKVContributors(ctx, azClient, kvRoleAssignments)
+		stream := listKeyVaultKVContributors(ctx, kvRoleAssignments)
 		outputStream(ctx, stream)
 		duration := time.Since(start)
 		log.Info("collection completed", "duration", duration.String())
 	}
 }
 
-func listKeyVaultKVContributors(ctx context.Context, client client.AzureClient, vmRoleAssignments <-chan interface{}) <-chan interface{} {
-	out := make(chan interface{})
+func listKeyVaultKVContributors(
+	ctx context.Context,
+	kvRoleAssignments <-chan azureWrapper[models.KeyVaultRoleAssignments],
+) <-chan any {
+	return pipeline.Map(ctx.Done(), kvRoleAssignments, func(ra azureWrapper[models.KeyVaultRoleAssignments]) any {
+		filteredAssignments := internal.Filter(ra.Data.RoleAssignments, kvRoleAssignmentFilter(constants.KeyVaultContributorRoleID))
 
-	go func() {
-		defer close(out)
-
-		for result := range pipeline.OrDone(ctx.Done(), vmRoleAssignments) {
-			if roleAssignments, ok := result.(AzureWrapper).Data.(models.KeyVaultRoleAssignments); !ok {
-				log.Error(fmt.Errorf("failed type assertion"), "unable to continue enumerating key vault kvContributors", "result", result)
-				return
-			} else {
-				var (
-					keyVaultKVContributors = models.KeyVaultKVContributors{
-						KeyVaultId: roleAssignments.KeyVaultId,
-					}
-					count = 0
-				)
-				for _, item := range roleAssignments.RoleAssignments {
-					roleDefinitionId := path.Base(item.RoleAssignment.Properties.RoleDefinitionId)
-
-					if roleDefinitionId == constants.KeyVaultContributorRoleID {
-						keyVaultKVContributor := models.KeyVaultKVContributor{
-							KVContributor: item.RoleAssignment,
-							KeyVaultId:    item.KeyVaultId,
-						}
-						log.V(2).Info("found key vault kvContributor", "keyVaultKVContributor", keyVaultKVContributor)
-						count++
-						keyVaultKVContributors.KVContributors = append(keyVaultKVContributors.KVContributors, keyVaultKVContributor)
-					}
-				}
-				out <- AzureWrapper{
-					Kind: enums.KindAZKeyVaultContributor,
-					Data: keyVaultKVContributors,
-				}
-				log.V(1).Info("finished listing key vault kvContributors", "keyVaultId", roleAssignments.KeyVaultId, "count", count)
+		kvContributors := internal.Map(filteredAssignments, func(ra models.KeyVaultRoleAssignment) models.KeyVaultKVContributor {
+			return models.KeyVaultKVContributor{
+				KVContributor: ra.RoleAssignment,
+				KeyVaultId:    ra.KeyVaultId,
 			}
-		}
-		log.Info("finished listing all key vault kvContributors")
-	}()
+		})
 
-	return out
+		return NewAzureWrapper(enums.KindAZKeyVaultKVContributor, models.KeyVaultKVContributors{
+			KeyVaultId:     ra.Data.KeyVaultId,
+			KVContributors: kvContributors,
+		})
+	})
 }
diff --git a/cmd/list-key-vault-kvcontributors_test.go b/cmd/list-key-vault-kvcontributors_test.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/bloodhoundad/azurehound/client/mocks"
 	"github.com/bloodhoundad/azurehound/constants"
+	"github.com/bloodhoundad/azurehound/enums"
 	"github.com/bloodhoundad/azurehound/models"
 	"github.com/bloodhoundad/azurehound/models/azure"
 	"github.com/golang/mock/gomock"
@@ -39,16 +40,17 @@ func TestListKeyVaultKVContributors(t *testing.T) {
 
 	mockClient := mocks.NewMockAzureClient(ctrl)
 
-	mockRoleAssignmentsChannel := make(chan interface{})
+	mockRoleAssignmentsChannel := make(chan azureWrapper[models.KeyVaultRoleAssignments])
 	mockTenant := azure.Tenant{}
 	mockClient.EXPECT().TenantInfo().Return(mockTenant).AnyTimes()
-	channel := listKeyVaultKVContributors(ctx, mockClient, mockRoleAssignmentsChannel)
+	channel := listKeyVaultKVContributors(ctx, mockRoleAssignmentsChannel)
 
 	go func() {
 		defer close(mockRoleAssignmentsChannel)
 
-		mockRoleAssignmentsChannel <- AzureWrapper{
-			Data: models.KeyVaultRoleAssignments{
+		mockRoleAssignmentsChannel <- NewAzureWrapper(
+			enums.KindAZKeyVaultRoleAssignment,
+			models.KeyVaultRoleAssignments{
 				KeyVaultId: "foo",
 				RoleAssignments: []models.KeyVaultRoleAssignment{
 					{
@@ -61,15 +63,11 @@ func TestListKeyVaultKVContributors(t *testing.T) {
 					},
 				},
 			},
-		}
+		)
 	}()
 
-	if result, ok := <-channel; !ok {
+	if _, ok := <-channel; !ok {
 		t.Fatalf("failed to receive from channel")
-	} else if wrapper, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	} else if _, ok := wrapper.Data.(models.KeyVaultKVContributors); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.KeyVaultKVContributors{})
 	}
 
 	if _, ok := <-channel; ok {