BED-4299 -- Display members of custom asset groups on Group Management (

#623) * rough draft, working on tests * tests passing * removing unnecessary change * updated to avoid matching substrings
SpecterOps · May 22, 2024 · 2c3eaac · 2c3eaac
1 parent 96380dd
commit 2c3eaac
Show file tree

Hide file tree

Showing 3 changed files with 83 additions and 2 deletions.
diff --git a/cmd/api/src/queries/graph.go b/cmd/api/src/queries/graph.go
@@ -24,6 +24,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"slices"
 	"sort"
 	"strconv"
 	"strings"
@@ -233,15 +234,29 @@ func (s *GraphQuery) GetAssetGroupNodes(ctx context.Context, assetGroupTag strin
 		if assetGroupNodes, err = ops.FetchNodeSet(tx.Nodes().Filterf(func() graph.Criteria {
 			filters := []graph.Criteria{
 				query.KindIn(query.Node(), azure.Entity, ad.Entity),
-				query.StringContains(query.NodeProperty(common.SystemTags.String()), assetGroupTag),
+				query.Or(
+					query.StringContains(query.NodeProperty(common.SystemTags.String()), assetGroupTag),
+					query.StringContains(query.NodeProperty(common.UserTags.String()), assetGroupTag),
+				),
 			}
 
 			return query.And(filters...)
 		})); err != nil {
 			return err
 		} else {
 			for _, node := range assetGroupNodes {
-				node.Properties.Set("type", analysis.GetNodeKindDisplayLabel(node))
+				// 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
 		}

diff --git a/cmd/api/src/queries/graph_integration_test.go b/cmd/api/src/queries/graph_integration_test.go
@@ -287,6 +287,35 @@ func TestGetAssetGroupComboNode(t *testing.T) {
 	})
 }
 
+func TestGetAssetGroupNodes(t *testing.T) {
+	testContext := integration.NewGraphTestContext(t, schema.DefaultGraphSchema())
+	testContext.DatabaseTestWithSetup(func(harness *integration.HarnessDetails) error {
+		harness.AssetGroupNodesHarness.Setup(testContext)
+		return nil
+	}, func(harness integration.HarnessDetails, db graph.Database) {
+		graphQuery := queries.NewGraphQuery(db, cache.Cache{}, config.Configuration{})
+
+		tierZeroNodes, err := graphQuery.GetAssetGroupNodes(context.Background(), harness.AssetGroupNodesHarness.TierZeroTag)
+		require.Nil(t, err)
+
+		customGroup1Nodes, err := graphQuery.GetAssetGroupNodes(context.Background(), harness.AssetGroupNodesHarness.CustomTag1)
+		require.Nil(t, err)
+
+		customGroup2Nodes, err := graphQuery.GetAssetGroupNodes(context.Background(), harness.AssetGroupNodesHarness.CustomTag2)
+		require.Nil(t, err)
+
+		require.True(t, tierZeroNodes.Contains(harness.AssetGroupNodesHarness.GroupB))
+		require.True(t, tierZeroNodes.Contains(harness.AssetGroupNodesHarness.GroupC))
+		require.Equal(t, 2, len(tierZeroNodes))
+
+		require.True(t, customGroup1Nodes.Contains(harness.AssetGroupNodesHarness.GroupD))
+		require.Equal(t, 1, len(customGroup1Nodes))
+
+		require.True(t, customGroup2Nodes.Contains(harness.AssetGroupNodesHarness.GroupE))
+		require.Equal(t, 1, len(customGroup2Nodes))
+	})
+}
+
 func TestGraphQuery_GetAllShortestPaths(t *testing.T) {
 	testContext := integration.NewGraphTestContext(t, schema.DefaultGraphSchema())
 	testContext.DatabaseTestWithSetup(

diff --git a/cmd/api/src/test/integration/harnesses.go b/cmd/api/src/test/integration/harnesses.go
@@ -330,6 +330,42 @@ func (s *AssetGroupComboNodeHarness) Setup(testCtx *GraphTestContext) {
 	testCtx.NewRelationship(s.GroupA, s.GroupB, ad.MemberOf)
 }
 
+type AssetGroupNodesHarness struct {
+	GroupA      *graph.Node
+	GroupB      *graph.Node
+	GroupC      *graph.Node
+	GroupD      *graph.Node
+	GroupE      *graph.Node
+	TierZeroTag string
+	CustomTag1  string
+	CustomTag2  string
+}
+
+func (s *AssetGroupNodesHarness) Setup(testCtx *GraphTestContext) {
+	domainSID := RandomDomainSID()
+
+	// use one tag value that contains the other as a substring to test that we only match exactly
+	s.TierZeroTag = ad.AdminTierZero
+	s.CustomTag1 = "custom_tag"
+	s.CustomTag2 = "another_custom_tag"
+
+	s.GroupA = testCtx.NewActiveDirectoryGroup("GroupA", domainSID)
+	s.GroupB = testCtx.NewActiveDirectoryGroup("GroupB", domainSID)
+	s.GroupC = testCtx.NewActiveDirectoryGroup("GroupC", domainSID)
+	s.GroupD = testCtx.NewActiveDirectoryGroup("GroupD", domainSID)
+	s.GroupE = testCtx.NewActiveDirectoryGroup("GroupE", domainSID)
+
+	s.GroupB.Properties.Set(common.SystemTags.String(), s.TierZeroTag)
+	s.GroupC.Properties.Set(common.SystemTags.String(), s.TierZeroTag)
+	s.GroupD.Properties.Set(common.UserTags.String(), s.CustomTag1)
+	s.GroupE.Properties.Set(common.UserTags.String(), s.CustomTag2)
+
+	testCtx.UpdateNode(s.GroupB)
+	testCtx.UpdateNode(s.GroupC)
+	testCtx.UpdateNode(s.GroupD)
+	testCtx.UpdateNode(s.GroupE)
+}
+
 type InboundControlHarness struct {
 	ControlledUser  *graph.Node
 	ControlledGroup *graph.Node
@@ -6833,6 +6869,7 @@ type HarnessDetails struct {
 	OutboundControl                                 OutboundControlHarness
 	InboundControl                                  InboundControlHarness
 	AssetGroupComboNodeHarness                      AssetGroupComboNodeHarness
+	AssetGroupNodesHarness                          AssetGroupNodesHarness
 	OUHarness                                       OUContainedHarness
 	MembershipHarness                               MembershipHarness
 	ForeignHarness                                  ForeignDomainHarness