rough draft, working on tests

SpecterOps · May 16, 2024 · 15c2fc2 · 15c2fc2
1 parent d9e67f9
commit 15c2fc2
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 4 deletions.
diff --git a/cmd/api/src/queries/graph.go b/cmd/api/src/queries/graph.go
@@ -179,7 +179,10 @@ func (s *GraphQuery) GetAssetGroupComboNode(ctx context.Context, owningObjectID
 		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),
+				),
 			}
 
 			if owningObjectID != "" {
@@ -233,7 +236,10 @@ 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...)

diff --git a/cmd/api/src/queries/graph_integration_test.go b/cmd/api/src/queries/graph_integration_test.go
@@ -287,6 +287,27 @@ func TestGetAssetGroupComboNode(t *testing.T) {
 	})
 }
 
+func TestGetAssetGroupNodes(t *testing.T) {
+	testContext := integration.NewGraphTestContext(t, schema.DefaultGraphSchema())
+	testContext.DatabaseTest(func(harness integration.HarnessDetails, db graph.Database) {
+		graphQuery := queries.NewGraphQuery(db, cache.Cache{}, config.Configuration{})
+
+		tierZeroNodes, err := graphQuery.GetAssetGroupNodes(context.Background(), ad.AdminTierZero)
+		require.Nil(t, err)
+
+		customGroupNodes, err := graphQuery.GetAssetGroupNodes(context.Background(), "custom_tag")
+		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.Contains(t, customGroupNodes, harness.AssetGroupNodesHarness.GroupD)
+		require.Contains(t, customGroupNodes, harness.AssetGroupNodesHarness.GroupE)
+		require.Equal(t, 2, len(customGroupNodes))
+	})
+}
+
 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,29 @@ 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
+}
+
+func (s *AssetGroupNodesHarness) Setup(testCtx *GraphTestContext) {
+	domainSID := RandomDomainSID()
+
+	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("GroupD", domainSID)
+
+	s.GroupB.Properties.Set(common.SystemTags.String(), ad.AdminTierZero)
+	s.GroupC.Properties.Set(common.SystemTags.String(), ad.AdminTierZero)
+	s.GroupD.Properties.Set(common.UserTags.String(), "custom_tag")
+	s.GroupE.Properties.Set(common.UserTags.String(), "custom_tag another_tag")
+}
+
 type InboundControlHarness struct {
 	ControlledUser  *graph.Node
 	ControlledGroup *graph.Node
@@ -6833,6 +6856,7 @@ type HarnessDetails struct {
 	OutboundControl                                 OutboundControlHarness
 	InboundControl                                  InboundControlHarness
 	AssetGroupComboNodeHarness                      AssetGroupComboNodeHarness
+	AssetGroupNodesHarness                          AssetGroupNodesHarness
 	OUHarness                                       OUContainedHarness
 	MembershipHarness                               MembershipHarness
 	ForeignHarness                                  ForeignDomainHarness

diff --git a/packages/go/dawgs/graph/node.go b/packages/go/dawgs/graph/node.go
@@ -18,11 +18,13 @@ package graph
 
 import (
 	"encoding/json"
+	"fmt"
+	"math"
+	"sync"
+
 	"github.com/RoaringBitmap/roaring"
 	"github.com/RoaringBitmap/roaring/roaring64"
 	"github.com/specterops/bloodhound/dawgs/util/size"
-	"math"
-	"sync"
 )
 
 const (
@@ -183,6 +185,7 @@ func (s NodeSet) KindSet() NodeKindSet {
 
 // Contains returns true if the ID of the given Node is stored within this NodeSet.
 func (s NodeSet) Contains(node *Node) bool {
+	fmt.Printf("THIS IS THE NODE %+v", node)
 	return s.ContainsID(node.ID)
 }