Merge pull request #27 from BloodHoundAD/fix-app-owners-appid

fix: use appid instead of objectid for listing app owners

cmd/list-app-owners.go

@@ -19,7 +19,6 @@ package cmd
 
 import (
 	"context"
-	"fmt"
 	"os"
 	"os/signal"
 	"sync"
@@ -57,42 +56,28 @@ func listAppOwnersCmdImpl(cmd *cobra.Command, args []string) {
 	log.Info("collection completed", "duration", duration.String())
 }
 
-func listAppOwners(ctx context.Context, client client.AzureClient, apps <-chan interface{}) <-chan interface{} {
+func listAppOwners(ctx context.Context, client client.AzureClient, apps <-chan azureWrapper[models.App]) <-chan azureWrapper[models.AppOwners] {
 	var (
-		out     = make(chan interface{})
-		ids     = make(chan string)
-		streams = pipeline.Demux(ctx.Done(), ids, 25)
+		out     = make(chan azureWrapper[models.AppOwners])
+		streams = pipeline.Demux(ctx.Done(), apps, 25)
 		wg      sync.WaitGroup
 	)
 
-	go func() {
-		defer close(ids)
-
-		for result := range pipeline.OrDone(ctx.Done(), apps) {
-			if app, ok := result.(AzureWrapper).Data.(models.App); !ok {
-				log.Error(fmt.Errorf("failed type assertion"), "unable to continue enumerating app owners", "result", result)
-				return
-			} else {
-				ids <- app.Id
-			}
-		}
-	}()
-
 	wg.Add(len(streams))
 	for i := range streams {
 		stream := streams[i]
 		go func() {
 			defer wg.Done()
-			for id := range stream {
+			for app := range stream {
 				var (
 					data = models.AppOwners{
-						AppId: id,
+						AppId: app.Data.AppId,
 					}
 					count = 0
 				)
-				for item := range client.ListAzureADAppOwners(ctx, id, "", "", "", nil) {
+				for item := range client.ListAzureADAppOwners(ctx, app.Data.Id, "", "", "", nil) {
 					if item.Error != nil {
-						log.Error(item.Error, "unable to continue processing owners for this app", "appId", id)
+						log.Error(item.Error, "unable to continue processing owners for this app", "appId", app.Data.AppId)
 					} else {
 						appOwner := models.AppOwner{
 							Owner: item.Ok,
@@ -104,11 +89,11 @@ func listAppOwners(ctx context.Context, client client.AzureClient, apps <-chan i
 					}
 				}
 
-				out <- AzureWrapper{
-					Kind: enums.KindAZAppOwner,
-					Data: data,
-				}
-				log.V(1).Info("finished listing app owners", "appId", id, "count", count)
+				out <- NewAzureWrapper(
+					enums.KindAZAppOwner,
+					data,
+				)
+				log.V(1).Info("finished listing app owners", "appId", app.Data.AppId, "count", count)
 			}
 		}()
 	}

cmd/list-app-owners_test.go

@@ -24,6 +24,7 @@ import (
 	"testing"
 
 	"github.com/bloodhoundad/azurehound/client/mocks"
+	"github.com/bloodhoundad/azurehound/enums"
 	"github.com/bloodhoundad/azurehound/models"
 	"github.com/bloodhoundad/azurehound/models/azure"
 	"github.com/golang/mock/gomock"
@@ -40,7 +41,7 @@ func TestListAppOwners(t *testing.T) {
 
 	mockClient := mocks.NewMockAzureClient(ctrl)
 
-	mockAppsChannel := make(chan interface{})
+	mockAppsChannel := make(chan azureWrapper[models.App])
 	mockAppOwnerChannel := make(chan azure.AppOwnerResult)
 	mockAppOwnerChannel2 := make(chan azure.AppOwnerResult)
 
@@ -53,12 +54,8 @@ func TestListAppOwners(t *testing.T) {
 
 	go func() {
 		defer close(mockAppsChannel)
-		mockAppsChannel <- AzureWrapper{
-			Data: models.App{},
-		}
-		mockAppsChannel <- AzureWrapper{
-			Data: models.App{},
-		}
+		mockAppsChannel <- NewAzureWrapper(enums.KindAZApp, models.App{})
+		mockAppsChannel <- NewAzureWrapper(enums.KindAZApp, models.App{})
 	}()
 	go func() {
 		defer close(mockAppOwnerChannel)
@@ -81,21 +78,13 @@ func TestListAppOwners(t *testing.T) {
 
 	if result, 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 data, ok := wrapper.Data.(models.AppOwners); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.AppOwners{})
-	} else if len(data.Owners) != 2 {
-		t.Errorf("got %v, want %v", len(data.Owners), 2)
+	} else if len(result.Data.Owners) != 2 {
+		t.Errorf("got %v, want %v", len(result.Data.Owners), 2)
 	}
 
 	if result, 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 data, ok := wrapper.Data.(models.AppOwners); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.AppOwners{})
-	} else if len(data.Owners) != 1 {
-		t.Errorf("got %v, want %v", len(data.Owners), 2)
+	} else if len(result.Data.Owners) != 1 {
+		t.Errorf("got %v, want %v", len(result.Data.Owners), 2)
 	}
 }

cmd/list-apps.go

@@ -54,8 +54,8 @@ func listAppsCmdImpl(cmd *cobra.Command, args []string) {
 	log.Info("collection completed", "duration", duration.String())
 }
 
-func listApps(ctx context.Context, client client.AzureClient) <-chan interface{} {
-	out := make(chan interface{})
+func listApps(ctx context.Context, client client.AzureClient) <-chan azureWrapper[models.App] {
+	out := make(chan azureWrapper[models.App])
 
 	go func() {
 		defer close(out)
@@ -67,14 +67,14 @@ func listApps(ctx context.Context, client client.AzureClient) <-chan interface{}
 			} else {
 				log.V(2).Info("found application", "app", item)
 				count++
-				out <- AzureWrapper{
-					Kind: enums.KindAZApp,
-					Data: models.App{
+				out <- NewAzureWrapper(
+					enums.KindAZApp,
+					models.App{
 						Application: item.Ok,
 						TenantId:    client.TenantInfo().TenantId,
 						TenantName:  client.TenantInfo().DisplayName,
 					},
-				}
+				)
 			}
 		}
 		log.Info("finished listing all apps", "count", count)

cmd/list-apps_test.go

@@ -57,11 +57,7 @@ func TestListApps(t *testing.T) {
 	}()
 
 	channel := listApps(ctx, mockClient)
-	result := <-channel
-	if _, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	}
-
+	<-channel
 	if _, ok := <-channel; ok {
 		t.Error("expected channel to close from an error result but it did not")
 	}

cmd/list-azure-ad.go

@@ -61,9 +61,6 @@ func listAzureADCmdImpl(cmd *cobra.Command, args []string) {
 
 func listAllAD(ctx context.Context, client client.AzureClient) <-chan interface{} {
 	var (
-		apps  = make(chan interface{})
-		apps2 = make(chan interface{})
-
 		devices  = make(chan interface{})
 		devices2 = make(chan interface{})
 
@@ -82,8 +79,9 @@ func listAllAD(ctx context.Context, client client.AzureClient) <-chan interface{
 	)
 
 	// Enumerate Apps, AppOwners and AppMembers
-	pipeline.Tee(ctx.Done(), listApps(ctx, client), apps, apps2)
-	appOwners := listAppOwners(ctx, client, apps2)
+	appChans := pipeline.TeeFixed(ctx.Done(), listApps(ctx, client), 2)
+	apps := pipeline.ToAny(ctx.Done(), appChans[0])
+	appOwners := pipeline.ToAny(ctx.Done(), listAppOwners(ctx, client, appChans[1]))
 
 	// Enumerate Devices and DeviceOwners
 	pipeline.Tee(ctx.Done(), listDevices(ctx, client), devices, devices2)

pipeline/pipeline.go

@@ -118,6 +118,12 @@ func Demux[D, T any](done <-chan D, in <-chan T, size int) []<-chan T {
 	return internal.Map(outputs, func(out chan T) <-chan T { return out })
 }
 
+func ToAny[D, T any](done <-chan D, in <-chan T) <-chan any {
+	return Map(done, in, func(t T) any {
+		return any(t)
+	})
+}
+
 func Map[D, T, U any](done <-chan D, in <-chan T, fn func(T) U) <-chan U {
 	out := make(chan U)
 	go func() {
@@ -143,7 +149,7 @@ func Filter[D, T any](done <-chan D, in <-chan T, fn func(T) bool) <-chan T {
 }
 
 // Tee copies the stream of data from a single channel to zero or more channels
-func Tee[D, T any](done <-chan D, in <-chan T, outputs ...chan<- T) {
+func Tee[D, T any](done <-chan D, in <-chan T, outputs ...chan T) {
 	go func() {
 		// Need to close outputs when goroutine exits to ensure we avoid deadlock
 		defer func() {
@@ -163,6 +169,16 @@ func Tee[D, T any](done <-chan D, in <-chan T, outputs ...chan<- T) {
 	}()
 }
 
+func TeeFixed[D, T any](done <-chan D, in <-chan T, size int) []<-chan T {
+	out := internal.Map(make([]any, size), func(_ any) chan T {
+		return make(chan T)
+	})
+	Tee(done, in, out...)
+	return internal.Map(out, func(c chan T) <-chan T {
+		return c
+	})
+}
+
 func Batch[D, T any](done <-chan D, in <-chan T, maxItems int, maxTimeout time.Duration) <-chan []T {
 	out := make(chan []T)