Bed-4168: fix: Inaccurate asset group counts attempt 2(#794)

cmd/api/src/api/v2/agi.go

@@ -431,7 +431,7 @@ func (s Resources) getAssetGroupMembers(response http.ResponseWriter, request *h
 		} else if assetGroup, err := s.DB.GetAssetGroup(request.Context(), int32(assetGroupID)); err != nil {
 			api.HandleDatabaseError(request, response, err)
 			return agMembers, err
-		} else if assetGroupNodes, err := s.GraphQuery.GetAssetGroupNodes(request.Context(), assetGroup.Tag); err != nil {
+		} else if assetGroupNodes, err := s.GraphQuery.GetAssetGroupNodes(request.Context(), assetGroup.Tag, assetGroup.SystemGroup); err != nil {
 			api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusInternalServerError, fmt.Sprintf("Graph error fetching nodes for asset group ID %v: %v", assetGroup.ID, err), request), response)
 			return agMembers, err
 		} else if agMembers, err = parseAGMembersFromNodes(assetGroupNodes, assetGroup.Selectors, int(assetGroup.ID)).SortBy(sortByColumns); err != nil {
@@ -520,16 +520,14 @@ func parseAGMembersFromNodes(nodes graph.NodeSet, selectors model.AssetGroupSele
 
 		if node.Kinds.ContainsOneOf(azure.Entity) {
 			if tenantID, err := node.Properties.Get(azure.TenantID.String()).String(); err != nil {
-				log.Warnf("%s is missing for node %d, skipping AG Membership...", azure.TenantID.String(), node.ID)
-				continue
+				log.Warnf("%s is missing for node %d", azure.TenantID.String(), node.ID)
 			} else {
 				agMember.EnvironmentKind = azure.Tenant.String()
 				agMember.EnvironmentID = tenantID
 			}
 		} else if node.Kinds.ContainsOneOf(ad.Entity) {
 			if domainSID, err := node.Properties.Get(ad.DomainSID.String()).String(); err != nil {
-				log.Warnf("%s is missing for node %d, skipping AG Membership...", ad.DomainSID.String(), node.ID)
-				continue
+				log.Warnf("%s is missing for node %d", ad.DomainSID.String(), node.ID)
 			} else {
 				agMember.EnvironmentKind = ad.Domain.String()
 				agMember.EnvironmentID = domainSID

cmd/api/src/api/v2/agi_internal_test.go

@@ -100,7 +100,7 @@ func TestParseAGMembersFromNodes_(t *testing.T) {
 
 func TestParseAGMembersFromNodes_MissingNodeProperties(t *testing.T) {
 	nodes := graph.NodeSet{
-		// the parse fn should handle nodes with missing name and missing properties with warnings and no output
+		// the parse fn should handle nodes with missing name and missing properties with warnings
 		1: &graph.Node{
 			ID:    1,
 			Kinds: graph.Kinds{ad.Entity, ad.Domain},
@@ -125,5 +125,5 @@ func TestParseAGMembersFromNodes_MissingNodeProperties(t *testing.T) {
 			SystemSelector: false,
 		}}, 1)
 
-	require.Equal(t, 0, len(members))
+	require.Equal(t, 2, len(members))
 }

cmd/api/src/api/v2/agi_test.go

@@ -1211,7 +1211,7 @@ func TestResources_ListAssetGroupMembers(t *testing.T) {
 						GetAssetGroup(gomock.Any(), gomock.Any()).
 						Return(assetGroup, nil)
 					mockGraph.EXPECT().
-						GetAssetGroupNodes(gomock.Any(), gomock.Any()).
+						GetAssetGroupNodes(gomock.Any(), gomock.Any(), gomock.Any()).
 						Return(graph.NodeSet{}, fmt.Errorf("GetAssetGroupNodes fail"))
 				},
 				Test: func(output apitest.Output) {
@@ -1228,7 +1228,7 @@ func TestResources_ListAssetGroupMembers(t *testing.T) {
 						GetAssetGroup(gomock.Any(), gomock.Any()).
 						Return(assetGroup, nil)
 					mockGraph.EXPECT().
-						GetAssetGroupNodes(gomock.Any(), gomock.Any()).
+						GetAssetGroupNodes(gomock.Any(), gomock.Any(), gomock.Any()).
 						Return(graph.NodeSet{
 							1: &graph.Node{
 								ID:    1,
@@ -1278,7 +1278,7 @@ func TestResources_ListAssetGroupMembers(t *testing.T) {
 						GetAssetGroup(gomock.Any(), gomock.Any()).
 						Return(assetGroup, nil)
 					mockGraph.EXPECT().
-						GetAssetGroupNodes(gomock.Any(), gomock.Any()).
+						GetAssetGroupNodes(gomock.Any(), gomock.Any(), gomock.Any()).
 						Return(graph.NodeSet{
 							1: &graph.Node{
 								ID:    1,
@@ -1314,7 +1314,7 @@ func TestResources_ListAssetGroupMembers(t *testing.T) {
 						GetAssetGroup(gomock.Any(), gomock.Any()).
 						Return(assetGroup, nil)
 					mockGraph.EXPECT().
-						GetAssetGroupNodes(gomock.Any(), gomock.Any()).
+						GetAssetGroupNodes(gomock.Any(), gomock.Any(), gomock.Any()).
 						Return(graph.NodeSet{
 							1: &graph.Node{
 								ID:    1,
@@ -1355,7 +1355,7 @@ func TestResources_ListAssetGroupMembers(t *testing.T) {
 						GetAssetGroup(gomock.Any(), gomock.Any()).
 						Return(assetGroup, nil)
 					mockGraph.EXPECT().
-						GetAssetGroupNodes(gomock.Any(), gomock.Any()).
+						GetAssetGroupNodes(gomock.Any(), gomock.Any(), gomock.Any()).
 						Return(graph.NodeSet{
 							1: &graph.Node{
 								ID:    1,
@@ -1398,7 +1398,7 @@ func TestResources_ListAssetGroupMembers(t *testing.T) {
 						GetAssetGroup(gomock.Any(), gomock.Any()).
 						Return(assetGroup, nil)
 					mockGraph.EXPECT().
-						GetAssetGroupNodes(gomock.Any(), gomock.Any()).
+						GetAssetGroupNodes(gomock.Any(), gomock.Any(), gomock.Any()).
 						Return(graph.NodeSet{
 							1: &graph.Node{
 								ID:    1,
@@ -1433,7 +1433,7 @@ func TestResources_ListAssetGroupMembers(t *testing.T) {
 						GetAssetGroup(gomock.Any(), gomock.Any()).
 						Return(assetGroup, nil)
 					mockGraph.EXPECT().
-						GetAssetGroupNodes(gomock.Any(), gomock.Any()).
+						GetAssetGroupNodes(gomock.Any(), gomock.Any(), gomock.Any()).
 						Return(graph.NodeSet{
 							1: &graph.Node{
 								ID:    1,
@@ -1574,7 +1574,7 @@ func TestResources_ListAssetGroupMembersCount(t *testing.T) {
 						GetAssetGroup(gomock.Any(), gomock.Any()).
 						Return(assetGroup, nil)
 					mockGraph.EXPECT().
-						GetAssetGroupNodes(gomock.Any(), gomock.Any()).
+						GetAssetGroupNodes(gomock.Any(), gomock.Any(), gomock.Any()).
 						Return(graph.NodeSet{}, fmt.Errorf("GetAssetGroupNodes fail"))
 				},
 				Test: func(output apitest.Output) {
@@ -1591,7 +1591,7 @@ func TestResources_ListAssetGroupMembersCount(t *testing.T) {
 						GetAssetGroup(gomock.Any(), gomock.Any()).
 						Return(assetGroup, nil)
 					mockGraph.EXPECT().
-						GetAssetGroupNodes(gomock.Any(), gomock.Any()).
+						GetAssetGroupNodes(gomock.Any(), gomock.Any(), gomock.Any()).
 						Return(graph.NodeSet{
 							1: &graph.Node{
 								ID:    1,
@@ -1627,7 +1627,7 @@ func TestResources_ListAssetGroupMembersCount(t *testing.T) {
 						GetAssetGroup(gomock.Any(), gomock.Any()).
 						Return(assetGroup, nil)
 					mockGraph.EXPECT().
-						GetAssetGroupNodes(gomock.Any(), gomock.Any()).
+						GetAssetGroupNodes(gomock.Any(), gomock.Any(), gomock.Any()).
 						Return(graph.NodeSet{
 							1: &graph.Node{
 								ID:    1,

cmd/api/src/daemons/datapipe/analysis.go

@@ -21,7 +21,6 @@ import (
 	"errors"
 	"fmt"
 
-	"github.com/specterops/bloodhound/analysis"
 	adAnalysis "github.com/specterops/bloodhound/analysis/ad"
 	"github.com/specterops/bloodhound/dawgs/graph"
 	"github.com/specterops/bloodhound/log"
@@ -92,7 +91,7 @@ func RunAnalysisOperations(ctx context.Context, db database.Database, graphDB gr
 		stats.LogStats()
 	}
 
-	if err := agi.RunAssetGroupIsolationCollections(ctx, db, graphDB, analysis.GetNodeKindDisplayLabel); err != nil {
+	if err := agi.RunAssetGroupIsolationCollections(ctx, db, graphDB); err != nil {
 		collectedErrors = append(collectedErrors, fmt.Errorf("asset group isolation collection failed: %w", err))
 		agiFailed = true
 	}

cmd/api/src/queries/graph.go

@@ -24,7 +24,6 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
-	"slices"
 	"sort"
 	"strconv"
 	"strings"
@@ -132,7 +131,7 @@ func BuildEntityQueryParams(request *http.Request, queryName string, pathDelegat
 
 type Graph interface {
 	GetAssetGroupComboNode(ctx context.Context, owningObjectID string, assetGroupTag string) (map[string]any, error)
-	GetAssetGroupNodes(ctx context.Context, assetGroupTag string) (graph.NodeSet, error)
+	GetAssetGroupNodes(ctx context.Context, assetGroupTag string, isSystemGroup bool) (graph.NodeSet, error)
 	GetAllShortestPaths(ctx context.Context, startNodeID string, endNodeID string, filter graph.Criteria) (graph.PathSet, error)
 	SearchNodesByName(ctx context.Context, nodeKinds graph.Kinds, nameQuery string, skip int, limit int) ([]model.SearchResult, error)
 	SearchByNameOrObjectID(ctx context.Context, searchValue string, searchType string) (graph.NodeSet, error)
@@ -154,9 +153,9 @@ type GraphQuery struct {
 	Cache                        cache.Cache
 	SlowQueryThreshold           int64 // Threshold in milliseconds
 	DisableCypherComplexityLimit bool
-	EnableCypherMutations bool
-	cypherEmitter         format.Emitter
-	strippedCypherEmitter format.Emitter
+	EnableCypherMutations        bool
+	cypherEmitter                format.Emitter
+	strippedCypherEmitter        format.Emitter
 }
 
 func NewGraphQuery(graphDB graph.Database, cache cache.Cache, cfg config.Configuration) *GraphQuery {
@@ -223,42 +222,20 @@ func (s *GraphQuery) GetAssetGroupComboNode(ctx context.Context, owningObjectID
 	})
 }
 
-func (s *GraphQuery) GetAssetGroupNodes(ctx context.Context, assetGroupTag string) (graph.NodeSet, error) {
+func (s *GraphQuery) GetAssetGroupNodes(ctx context.Context, assetGroupTag string, isSystemGroup bool) (graph.NodeSet, error) {
 	var (
 		assetGroupNodes graph.NodeSet
 		err             error
 	)
-	return assetGroupNodes, s.Graph.ReadTransaction(ctx, func(tx graph.Transaction) error {
-		if assetGroupNodes, err = ops.FetchNodeSet(tx.Nodes().Filterf(func() graph.Criteria {
-			filters := []graph.Criteria{
-				query.KindIn(query.Node(), azure.Entity, ad.Entity),
-				query.Or(
-					query.StringContains(query.NodeProperty(common.SystemTags.String()), assetGroupTag),
-					query.StringContains(query.NodeProperty(common.UserTags.String()), assetGroupTag),
-				),
-			}
 
-			return query.And(filters...)
-		})); err != nil {
+	err = s.Graph.ReadTransaction(ctx, func(tx graph.Transaction) error {
+		if assetGroupNodes, err = agi.FetchAssetGroupNodes(tx, assetGroupTag, isSystemGroup); err != nil {
 			return err
-		} else {
-			for _, node := range assetGroupNodes {
-				// We need to filter out nodes that do not contain an exact tag match
-				var (
-					systemTags, _ = node.Properties.Get(common.SystemTags.String()).String()
-					userTags, _   = node.Properties.Get(common.UserTags.String()).String()
-					allTags       = append(strings.Split(systemTags, " "), strings.Split(userTags, " ")...)
-				)
-
-				if !slices.Contains(allTags, assetGroupTag) {
-					assetGroupNodes.Remove(node.ID)
-				} else {
-					node.Properties.Set("type", analysis.GetNodeKindDisplayLabel(node))
-				}
-			}
-			return nil
 		}
+		return nil
 	})
+
+	return assetGroupNodes, err
 }
 
 func (s *GraphQuery) GetAllShortestPaths(ctx context.Context, startNodeID string, endNodeID string, filter graph.Criteria) (graph.PathSet, error) {

cmd/api/src/queries/graph_integration_test.go

@@ -295,13 +295,13 @@ func TestGetAssetGroupNodes(t *testing.T) {
 	}, func(harness integration.HarnessDetails, db graph.Database) {
 		graphQuery := queries.NewGraphQuery(db, cache.Cache{}, config.Configuration{})
 
-		tierZeroNodes, err := graphQuery.GetAssetGroupNodes(context.Background(), harness.AssetGroupNodesHarness.TierZeroTag)
+		tierZeroNodes, err := graphQuery.GetAssetGroupNodes(context.Background(), harness.AssetGroupNodesHarness.TierZeroTag, true)
 		require.Nil(t, err)
 
-		customGroup1Nodes, err := graphQuery.GetAssetGroupNodes(context.Background(), harness.AssetGroupNodesHarness.CustomTag1)
+		customGroup1Nodes, err := graphQuery.GetAssetGroupNodes(context.Background(), harness.AssetGroupNodesHarness.CustomTag1, false)
 		require.Nil(t, err)
 
-		customGroup2Nodes, err := graphQuery.GetAssetGroupNodes(context.Background(), harness.AssetGroupNodesHarness.CustomTag2)
+		customGroup2Nodes, err := graphQuery.GetAssetGroupNodes(context.Background(), harness.AssetGroupNodesHarness.CustomTag2, false)
 		require.Nil(t, err)
 
 		require.True(t, tierZeroNodes.Contains(harness.AssetGroupNodesHarness.GroupB))

cmd/api/src/services/agi/agi.go

@@ -19,7 +19,10 @@ package agi
 
 import (
 	"context"
+	"slices"
+	"strings"
 
+	"github.com/specterops/bloodhound/analysis"
 	"github.com/specterops/bloodhound/dawgs/graph"
 	"github.com/specterops/bloodhound/dawgs/ops"
 	"github.com/specterops/bloodhound/dawgs/query"
@@ -36,26 +39,46 @@ type AgiData interface {
 	CreateAssetGroupCollection(ctx context.Context, collection model.AssetGroupCollection, entries model.AssetGroupCollectionEntries) error
 }
 
-func RunAssetGroupIsolationCollections(ctx context.Context, db AgiData, graphDB graph.Database, kindGetter func(*graph.Node) string) error {
+func FetchAssetGroupNodes(tx graph.Transaction, assetGroupTag string, isSystemGroup bool) (graph.NodeSet, error) {
+	var (
+		assetGroupNodes graph.NodeSet
+		tagPropertyStr  = common.SystemTags.String()
+		err             error
+	)
+
+	if !isSystemGroup {
+		tagPropertyStr = common.UserTags.String()
+	}
+
+	if assetGroupNodes, err = ops.FetchNodeSet(tx.Nodes().Filterf(func() graph.Criteria {
+		return query.And(
+			query.KindIn(query.Node(), ad.Entity, azure.Entity),
+			query.StringContains(query.NodeProperty(tagPropertyStr), assetGroupTag),
+		)
+	})); err != nil {
+		return graph.NodeSet{}, err
+	} else {
+		// tags are space seperated, so we have to loop and remove any that are not exact matches
+		for _, node := range assetGroupNodes {
+			tags, _ := node.Properties.Get(tagPropertyStr).String()
+			if !slices.Contains(strings.Split(tags, " "), assetGroupTag) {
+				assetGroupNodes.Remove(node.ID)
+			}
+		}
+	}
+
+	return assetGroupNodes, err
+}
+
+func RunAssetGroupIsolationCollections(ctx context.Context, db AgiData, graphDB graph.Database) error {
 	defer log.Measure(log.LevelInfo, "Asset Group Isolation Collections")()
 
 	if assetGroups, err := db.GetAllAssetGroups(ctx, "", model.SQLFilter{}); err != nil {
 		return err
 	} else {
 		return graphDB.WriteTransaction(ctx, func(tx graph.Transaction) error {
 			for _, assetGroup := range assetGroups {
-				if assetGroupNodes, err := ops.FetchNodes(tx.Nodes().Filterf(func() graph.Criteria {
-					tagPropertyStr := common.SystemTags.String()
-
-					if !assetGroup.SystemGroup {
-						tagPropertyStr = common.UserTags.String()
-					}
-
-					return query.And(
-						query.KindIn(query.Node(), ad.Entity, azure.Entity),
-						query.StringContains(query.NodeProperty(tagPropertyStr), assetGroup.Tag),
-					)
-				})); err != nil {
+				if assetGroupNodes, err := FetchAssetGroupNodes(tx, assetGroup.Tag, assetGroup.SystemGroup); err != nil {
 					return err
 				} else {
 					var (
@@ -65,16 +88,18 @@ func RunAssetGroupIsolationCollections(ctx context.Context, db AgiData, graphDB
 						}
 					)
 
-					for idx, node := range assetGroupNodes {
+					idx := 0
+					for _, node := range assetGroupNodes {
 						if objectID, err := node.Properties.Get(common.ObjectID.String()).String(); err != nil {
 							log.Errorf("Node %d that does not have valid %s property", node.ID, common.ObjectID)
 						} else {
 							entries[idx] = model.AssetGroupCollectionEntry{
 								ObjectID:   objectID,
-								NodeLabel:  kindGetter(node),
+								NodeLabel:  analysis.GetNodeKindDisplayLabel(node),
 								Properties: node.Properties.Map,
 							}
 						}
+						idx++
 					}
 
 					// Enter a collection, even if it's empty to signal that we did do a tagging/collection run

cmd/ui/src/utils.ts

@@ -25,7 +25,7 @@ import { isLink, isNode } from 'src/ducks/graph/utils';
 import { Glyph } from 'src/rendering/programs/node.glyphs';
 import { store } from 'src/store';
 
-const IGNORE_401_LOGOUT = ['/api/v2/login', '/api/v2/logout', '/api/v2/features']
+const IGNORE_401_LOGOUT = ['/api/v2/login', '/api/v2/logout', '/api/v2/features'];
 
 export const getDatesInRange = (startDate: Date, endDate: Date) => {
     const date = new Date(startDate.getTime());

packages/go/analysis/go.mod

@@ -21,7 +21,7 @@ go 1.21
 require (
 	github.com/RoaringBitmap/roaring v1.3.0
 	github.com/bloodhoundad/azurehound/v2 v2.0.1
-	github.com/stretchr/testify v1.8.4
+	github.com/stretchr/testify v1.9.0
 	go.uber.org/mock v0.2.0
 )
 

packages/go/analysis/go.sum

@@ -7,7 +7,7 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/mschoch/smat v0.2.0 h1:8imxQsjDm8yFEAVBe7azKmKSgzSkZXDuKkSq9374khM=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ=
-github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 go.uber.org/mock v0.2.0 h1:TaP3xedm7JaAgScZO7tlvlKrqT0p7I6OsdGB5YNSMDU=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=