BED-5198 - Merge stage/v6.3.1

cmd/api/src/analysis/azure/queries_test.go

@@ -33,6 +33,7 @@ import (
 
 func TestAnalysisAzure_GraphStats(t *testing.T) {
 	testCtx := integration.NewGraphTestContext(t, schema.DefaultGraphSchema())
+	testCtx.SetupAzure()
 	testCtx.DatabaseTest(func(harness integration.HarnessDetails, db graph.Database) {
 
 		_, agg, err := azure2.GraphStats(context.TODO(), testCtx.Graph.Database)

cmd/api/src/api/tools/pg.go

@@ -82,10 +82,8 @@ func migrateTypes(ctx context.Context, neoDB, pgDB graph.Database) error {
 		return err
 	}
 
-	return pgDB.WriteTransaction(ctx, func(tx graph.Transaction) error {
-		_, err := pgDB.(*pg.Driver).KindMapper().AssertKinds(tx, append(neoNodeKinds, neoEdgeKinds...))
-		return err
-	})
+	_, err := pgDB.(*pg.Driver).KindMapper().AssertKinds(ctx, append(neoNodeKinds, neoEdgeKinds...))
+	return err
 }
 
 func convertNeo4jProperties(properties *graph.Properties) error {

cmd/api/src/api/v2/integration/ingest.go

@@ -21,6 +21,8 @@ import (
 	"strings"
 	"time"
 
+	"github.com/specterops/bloodhound/graphschema"
+
 	"github.com/specterops/bloodhound/dawgs/graph"
 	"github.com/specterops/bloodhound/src/model"
 	"github.com/specterops/bloodhound/src/model/appcfg"
@@ -198,7 +200,7 @@ func (s *Context) WaitForDatapipeAnalysis(timeout time.Duration, originalWrapper
 type IngestAssertion func(testCtrl test.Controller, tx graph.Transaction)
 
 func (s *Context) AssertIngest(assertion IngestAssertion) {
-	graphDB := integration.OpenGraphDB(s.TestCtrl)
+	graphDB := integration.OpenGraphDB(s.TestCtrl, graphschema.DefaultGraphSchema())
 	defer graphDB.Close(s.ctx)
 
 	require.Nil(s.TestCtrl, graphDB.ReadTransaction(s.ctx, func(tx graph.Transaction) error {
@@ -208,7 +210,7 @@ func (s *Context) AssertIngest(assertion IngestAssertion) {
 }
 
 func (s *Context) AssertIngestProperties(assertion IngestAssertion) {
-	graphDB := integration.OpenGraphDB(s.TestCtrl)
+	graphDB := integration.OpenGraphDB(s.TestCtrl, graphschema.DefaultGraphSchema())
 	defer graphDB.Close(s.ctx)
 
 	require.Nil(s.TestCtrl, graphDB.ReadTransaction(s.ctx, func(tx graph.Transaction) error {

cmd/api/src/api/v2/integration/reconciliation.go

@@ -18,6 +18,7 @@ package integration
 
 import (
 	"github.com/specterops/bloodhound/dawgs/graph"
+	"github.com/specterops/bloodhound/graphschema"
 	"github.com/specterops/bloodhound/src/test"
 	"github.com/specterops/bloodhound/src/test/integration"
 	"github.com/stretchr/testify/require"
@@ -26,7 +27,7 @@ import (
 type ReconciliationAssertion func(testCtrl test.Controller, tx graph.Transaction)
 
 func (s *Context) AssertReconciliation(assertion ReconciliationAssertion) {
-	graphDB := integration.OpenGraphDB(s.TestCtrl)
+	graphDB := integration.OpenGraphDB(s.TestCtrl, graphschema.DefaultGraphSchema())
 	defer graphDB.Close(s.ctx)
 
 	require.Nil(s.TestCtrl, graphDB.ReadTransaction(s.ctx, func(tx graph.Transaction) error {

cmd/api/src/queries/graph_integration_test.go

@@ -165,6 +165,7 @@ func TestGetEntityResults(t *testing.T) {
 	queryCache, err := cache.NewCache(cache.Config{MaxSize: 1})
 	require.Nil(t, err)
 
+	testContext.SetupActiveDirectory()
 	testContext.DatabaseTest(func(harness integration.HarnessDetails, db graph.Database) {
 		objectID, err := harness.InboundControl.ControlledUser.Properties.Get(common.ObjectID.String()).String()
 		require.Nil(t, err)
@@ -197,6 +198,7 @@ func TestGetEntityResults_QueryShorterThanSlowQueryThreshold(t *testing.T) {
 	queryCache, err := cache.NewCache(cache.Config{MaxSize: 1})
 	require.Nil(t, err)
 
+	testContext.SetupActiveDirectory()
 	testContext.DatabaseTest(func(harness integration.HarnessDetails, db graph.Database) {
 		objectID, err := harness.InboundControl.ControlledUser.Properties.Get(common.ObjectID.String()).String()
 		require.Nil(t, err)
@@ -230,6 +232,7 @@ func TestGetEntityResults_Cache(t *testing.T) {
 	queryCache, err := cache.NewCache(cache.Config{MaxSize: 2})
 	require.Nil(t, err)
 
+	testContext.SetupActiveDirectory()
 	testContext.DatabaseTest(func(harness integration.HarnessDetails, db graph.Database) {
 		objectID, err := harness.InboundControl.ControlledUser.Properties.Get(common.ObjectID.String()).String()
 		require.Nil(t, err)
@@ -270,6 +273,7 @@ func TestGetEntityResults_Cache(t *testing.T) {
 
 func TestGetAssetGroupComboNode(t *testing.T) {
 	testContext := integration.NewGraphTestContext(t, schema.DefaultGraphSchema())
+	testContext.SetupActiveDirectory()
 	testContext.DatabaseTest(func(harness integration.HarnessDetails, db graph.Database) {
 		graphQuery := queries.NewGraphQuery(db, cache.Cache{}, config.Configuration{})
 		comboNode, err := graphQuery.GetAssetGroupComboNode(context.Background(), "", ad.AdminTierZero)

cmd/api/src/test/integration/context.go

@@ -67,7 +67,7 @@ func (s *GraphContext) End(t test.Context) {
 func NewGraphContext(ctx test.Context, schema graph.Schema) *GraphContext {
 	graphContext := &GraphContext{
 		schema:   schema,
-		Database: OpenGraphDB(ctx),
+		Database: OpenGraphDB(ctx, schema),
 	}
 
 	// Initialize the graph context

cmd/api/src/test/integration/dawgs.go

@@ -23,7 +23,6 @@ import (
 	"github.com/specterops/bloodhound/dawgs/drivers/neo4j"
 	"github.com/specterops/bloodhound/dawgs/drivers/pg"
 	"github.com/specterops/bloodhound/dawgs/graph"
-	schema "github.com/specterops/bloodhound/graphschema"
 	"github.com/specterops/bloodhound/src/config"
 	"github.com/specterops/bloodhound/src/test"
 	"github.com/specterops/bloodhound/src/test/integration/utils"
@@ -39,7 +38,7 @@ func LoadConfiguration(testCtrl test.Controller) config.Configuration {
 	return cfg
 }
 
-func OpenGraphDB(testCtrl test.Controller) graph.Database {
+func OpenGraphDB(testCtrl test.Controller, schema graph.Schema) graph.Database {
 	var (
 		cfg           = LoadConfiguration(testCtrl)
 		graphDatabase graph.Database
@@ -62,7 +61,7 @@ func OpenGraphDB(testCtrl test.Controller) graph.Database {
 	}
 
 	test.RequireNilErrf(testCtrl, err, "Failed connecting to graph database: %v", err)
-	test.RequireNilErr(testCtrl, graphDatabase.AssertSchema(context.Background(), schema.DefaultGraphSchema()))
+	test.RequireNilErr(testCtrl, graphDatabase.AssertSchema(context.Background(), schema))
 
 	return graphDatabase
 }

cmd/api/src/test/integration/graph.go

@@ -82,10 +82,14 @@ func (s *GraphTestContext) UpdateNode(node *graph.Node) {
 	})
 }
 
-func (s *GraphTestContext) DatabaseTest(dbDelegate func(harness HarnessDetails, db graph.Database)) {
-	s.setupActiveDirectory()
-	s.setupAzure()
+func (s *GraphTestContext) InitializeHarness(harness GraphTestHarness) {
+	s.Graph.WriteTransaction(s.testCtx, func(tx graph.Transaction) error {
+		harness.Setup(s)
+		return nil
+	})
+}
 
+func (s *GraphTestContext) DatabaseTest(dbDelegate func(harness HarnessDetails, db graph.Database)) {
 	dbDelegate(s.Harness, s.Graph.Database)
 }
 
@@ -109,8 +113,7 @@ func (s *GraphTestContext) DatabaseTestWithSetup(setup func(harness *HarnessDeta
 }
 
 func (s *GraphTestContext) BatchTest(batchDelegate func(harness HarnessDetails, batch graph.Batch), assertionDelegate func(details HarnessDetails, tx graph.Transaction)) {
-	s.setupActiveDirectory()
-	s.setupAzure()
+	s.SetupAzureAndActiveDirectory()
 
 	s.Graph.BatchOperation(s.testCtx, func(batch graph.Batch) error {
 		batchDelegate(s.Harness, batch)
@@ -124,8 +127,7 @@ func (s *GraphTestContext) BatchTest(batchDelegate func(harness HarnessDetails,
 }
 
 func (s *GraphTestContext) TransactionalTest(txDelegate func(harness HarnessDetails, tx graph.Transaction)) {
-	s.setupActiveDirectory()
-	s.setupAzure()
+	s.SetupAzureAndActiveDirectory()
 
 	s.Graph.WriteTransaction(s.testCtx, func(tx graph.Transaction) error {
 		txDelegate(s.Harness, tx)
@@ -319,6 +321,7 @@ func (s *GraphTestContext) NewAzureTenant(tenantID string) *graph.Node {
 	return s.NewNode(graph.AsProperties(graph.PropertyMap{
 		common.Name:     "New Tenant",
 		common.ObjectID: tenantID,
+		azure.TenantID:  tenantID,
 		azure.License:   "license",
 	}), azure.Entity, azure.Tenant)
 }
@@ -345,12 +348,27 @@ func (s *GraphTestContext) NewActiveDirectoryComputer(name, domainSID string) *g
 	}), ad.Entity, ad.Computer)
 }
 
-func (s *GraphTestContext) NewActiveDirectoryUser(name, domainSID string, isTierZero ...bool) *graph.Node {
+func (s *GraphTestContext) NewActiveDirectoryContainer(name, domainSID string) *graph.Node {
 	return s.NewNode(graph.AsProperties(graph.PropertyMap{
+		common.Name:     name,
+		common.ObjectID: must.NewUUIDv4().String(),
+		ad.DomainSID:    domainSID,
+	}), ad.Entity, ad.Container)
+}
+
+func (s *GraphTestContext) NewActiveDirectoryUser(name, domainSID string, isTierZero ...bool) *graph.Node {
+
+	propertyMap := graph.PropertyMap{
 		common.Name:     name,
 		common.ObjectID: strings.ToUpper(must.NewUUIDv4().String()),
 		ad.DomainSID:    domainSID,
-	}), ad.Entity, ad.User)
+	}
+
+	if isTierZero != nil && isTierZero[0] {
+		propertyMap[common.SystemTags] = ad.AdminTierZero
+	}
+
+	return s.NewNode(graph.AsProperties(propertyMap), ad.Entity, ad.User)
 }
 
 func (s *GraphTestContext) NewCustomActiveDirectoryUser(properties *graph.Properties) *graph.Node {
@@ -515,7 +533,12 @@ type CertTemplateData struct {
 	CertificatePolicy             []string
 }
 
-func (s *GraphTestContext) setupAzure() {
+func (s *GraphTestContext) SetupAzureAndActiveDirectory() {
+	s.SetupAzure()
+	s.SetupActiveDirectory()
+}
+
+func (s *GraphTestContext) SetupAzure() {
 	s.Harness.AZBaseHarness.Setup(s)
 	s.Harness.AZGroupMembership.Setup(s)
 	s.Harness.AZEntityPanelHarness.Setup(s)
@@ -530,7 +553,7 @@ func (s *GraphTestContext) setupAzure() {
 	s.Harness.AZManagementGroup.Setup(s)
 }
 
-func (s *GraphTestContext) setupActiveDirectory() {
+func (s *GraphTestContext) SetupActiveDirectory() {
 	// startServer a host of Tier Zero tagged assets
 	s.Harness.RootADHarness.Setup(s)