Add "Update Asset Group Selectors" endpoint (#138)

* move endpoint to bhce, fix bug in response format, update docs * removed formatting update * wrap swagger request definition in array * updated endpoint description in docs
SpecterOps · Oct 16, 2023 · 026ca83 · 026ca83
1 parent 4ecb306
commit 026ca83
Show file tree

Hide file tree

Showing 7 changed files with 285 additions and 61 deletions.
diff --git a/cmd/api/src/api/agi_test.go b/cmd/api/src/api/agi_test.go
@@ -1,29 +1,44 @@
 // Copyright 2023 Specter Ops, Inc.
-// 
+//
 // Licensed under the Apache License, Version 2.0
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
-// 
+//
 //     http://www.apache.org/licenses/LICENSE-2.0
-// 
+//
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-// 
+//
 // SPDX-License-Identifier: Apache-2.0
 
 package api_test
 
 import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
 	"testing"
 
+	"github.com/gorilla/mux"
+	"github.com/specterops/bloodhound/graphschema/ad"
+	"github.com/specterops/bloodhound/graphschema/azure"
+	"github.com/specterops/bloodhound/headers"
+	"github.com/specterops/bloodhound/mediatypes"
 	"github.com/specterops/bloodhound/src/api"
+	v2 "github.com/specterops/bloodhound/src/api/v2"
+	"github.com/specterops/bloodhound/src/auth"
+	"github.com/specterops/bloodhound/src/ctx"
+	datapipeMocks "github.com/specterops/bloodhound/src/daemons/datapipe/mocks"
+	dbMocks "github.com/specterops/bloodhound/src/database/mocks"
 	"github.com/specterops/bloodhound/src/model"
+	"github.com/specterops/bloodhound/src/test/must"
 	"github.com/stretchr/testify/require"
-	"github.com/specterops/bloodhound/graphschema/ad"
-	"github.com/specterops/bloodhound/graphschema/azure"
+	"go.uber.org/mock/gomock"
 )
 
 func TestAssetGroupMembers_SortBy(t *testing.T) {
@@ -381,3 +396,150 @@ func TestAssetGroupMembers_BuildFilteringConditional_Error(t *testing.T) {
 	require.Error(t, err)
 	require.Contains(t, err.Error(), api.ErrorResponseDetailsBadQueryParameterFilters)
 }
+
+func TestResources_UpdateAssetGroupSelectors_GetAssetGroupError(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	payload := []model.AssetGroupSelectorSpec{
+		{
+			SelectorName:   "test",
+			EntityObjectID: "1",
+		},
+	}
+
+	req, err := http.NewRequest("POST", "/api/v2/asset-groups/1/selectors", must.MarshalJSONReader(payload))
+	require.Nil(t, err)
+
+	req.Header.Set(headers.ContentType.String(), mediatypes.ApplicationJson.String())
+	req = mux.SetURLVars(req, map[string]string{api.URIPathVariableAssetGroupID: "1"})
+
+	mockDB := dbMocks.NewMockDatabase(mockCtrl)
+	mockDB.EXPECT().GetAssetGroup(gomock.Any()).Return(model.AssetGroup{}, fmt.Errorf("test error"))
+	handlers := v2.Resources{DB: mockDB}
+
+	response := httptest.NewRecorder()
+	handler := http.HandlerFunc(handlers.UpdateAssetGroupSelectors)
+
+	handler.ServeHTTP(response, req)
+	require.Equal(t, http.StatusInternalServerError, response.Code)
+	require.Contains(t, response.Body.String(), api.ErrorResponseDetailsInternalServerError)
+}
+
+func TestResources_UpdateAssetGroupSelectors_PayloadError(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	payload := "INVALID PAYLOAD"
+
+	req, err := http.NewRequest("POST", "/api/v2/asset-groups/1/selectors", must.MarshalJSONReader(payload))
+	require.Nil(t, err)
+
+	req.Header.Set(headers.ContentType.String(), mediatypes.ApplicationJson.String())
+	req = mux.SetURLVars(req, map[string]string{api.URIPathVariableAssetGroupID: "1"})
+
+	assetGroup := model.AssetGroup{
+		Name:        "test group",
+		Tag:         "test tag",
+		SystemGroup: false,
+	}
+
+	mockDB := dbMocks.NewMockDatabase(mockCtrl)
+	mockDB.EXPECT().GetAssetGroup(gomock.Any()).Return(assetGroup, nil)
+	handlers := v2.Resources{DB: mockDB}
+
+	response := httptest.NewRecorder()
+	handler := http.HandlerFunc(handlers.UpdateAssetGroupSelectors)
+
+	handler.ServeHTTP(response, req)
+	require.Equal(t, http.StatusBadRequest, response.Code)
+	require.Contains(t, response.Body.String(), api.ErrorResponsePayloadUnmarshalError)
+}
+
+func TestResources_UpdateAssetGroupSelectors_Success(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	payload := []model.AssetGroupSelectorSpec{
+		{
+			SelectorName:   "test",
+			EntityObjectID: "1",
+			Action:         model.SelectorSpecActionAdd,
+		},
+		{
+			SelectorName:   "test2",
+			EntityObjectID: "2",
+			Action:         model.SelectorSpecActionRemove,
+		},
+	}
+
+	req, err := http.NewRequest("POST", "/api/v2/asset-groups/1/selectors", must.MarshalJSONReader(payload))
+	require.Nil(t, err)
+
+	req.Header.Set(headers.ContentType.String(), mediatypes.ApplicationJson.String())
+
+	bheCtx := ctx.Context{
+		RequestID: "requestID",
+		AuthCtx: auth.Context{
+			Owner:   model.User{},
+			Session: model.UserSession{},
+		},
+	}
+	req = req.WithContext(context.WithValue(context.Background(), ctx.ValueKey, bheCtx.WithRequestID("requestID")))
+	req = mux.SetURLVars(req, map[string]string{api.URIPathVariableAssetGroupID: "1"})
+
+	assetGroup := model.AssetGroup{
+		Name:        "test group",
+		Tag:         "test tag",
+		SystemGroup: false,
+	}
+
+	expectedResult := map[string]model.AssetGroupSelectors{
+		"added_selectors": {
+			model.AssetGroupSelector{
+				AssetGroupID: assetGroup.ID,
+				Name:         payload[0].SelectorName,
+				Selector:     payload[0].EntityObjectID,
+			},
+		},
+		"removed_selectors": {
+			model.AssetGroupSelector{
+				AssetGroupID: assetGroup.ID,
+				Name:         payload[1].SelectorName,
+				Selector:     payload[1].EntityObjectID,
+			},
+		},
+	}
+
+	mockDB := dbMocks.NewMockDatabase(mockCtrl)
+	mockDB.EXPECT().GetAssetGroup(gomock.Any()).Return(assetGroup, nil)
+	mockDB.EXPECT().UpdateAssetGroupSelectors(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(expectedResult, nil)
+
+	mockTasker := datapipeMocks.NewMockTasker(mockCtrl)
+	mockTasker.EXPECT().RequestAnalysis()
+
+	handlers := v2.Resources{
+		DB:           mockDB,
+		TaskNotifier: mockTasker,
+	}
+
+	response := httptest.NewRecorder()
+	handler := http.HandlerFunc(handlers.UpdateAssetGroupSelectors)
+
+	handler.ServeHTTP(response, req)
+	require.Equal(t, http.StatusOK, response.Code)
+
+	resp := api.ResponseWrapper{}
+	err = json.Unmarshal(response.Body.Bytes(), &resp)
+	require.Nil(t, err)
+
+	dataJSON, err := json.Marshal(resp.Data)
+	require.Nil(t, err)
+
+	data := make(map[string][]model.AssetGroupSelector, 0)
+	err = json.Unmarshal(dataJSON, &data)
+	require.Nil(t, err)
+
+	require.Equal(t, expectedResult["added_selectors"][0].Name, data["added_selectors"][0].Name)
+	require.Equal(t, expectedResult["removed_selectors"][0].Name, data["removed_selectors"][0].Name)
+}
diff --git a/cmd/api/src/api/registration/v2.go b/cmd/api/src/api/registration/v2.go
@@ -137,6 +137,7 @@ func NewV2API(cfg config.Configuration, resources v2.Resources, routerInst *rout
 		routerInst.GET(fmt.Sprintf("/api/v2/asset-groups/{%s}/custom-selectors", api.URIPathVariableAssetGroupID), resources.GetAssetGroupCustomMemberCount).RequirePermissions(permissions.GraphDBRead),
 		routerInst.DELETE(fmt.Sprintf("/api/v2/asset-groups/{%s}", api.URIPathVariableAssetGroupID), resources.DeleteAssetGroup).RequirePermissions(permissions.GraphDBWrite),
 		routerInst.PUT(fmt.Sprintf("/api/v2/asset-groups/{%s}", api.URIPathVariableAssetGroupID), resources.UpdateAssetGroup).RequirePermissions(permissions.GraphDBWrite),
+		routerInst.POST(fmt.Sprintf("/api/v2/asset-groups/{%s}/selectors", api.URIPathVariableAssetGroupID), resources.UpdateAssetGroupSelectors).RequirePermissions(permissions.GraphDBWrite),
 		routerInst.DELETE(fmt.Sprintf("/api/v2/asset-groups/{%s}/selectors/{%s}", api.URIPathVariableAssetGroupID, api.URIPathVariableAssetGroupSelectorID), resources.DeleteAssetGroupSelector).RequirePermissions(permissions.GraphDBWrite),
 		routerInst.GET(fmt.Sprintf("/api/v2/asset-groups/{%s}/collections", api.URIPathVariableAssetGroupID), resources.ListAssetGroupCollections).RequirePermissions(permissions.GraphDBRead),
 		routerInst.GET(fmt.Sprintf("/api/v2/asset-groups/{%s}/members", api.URIPathVariableAssetGroupID), resources.ListAssetGroupMembers).RequirePermissions(permissions.GraphDBRead),

diff --git a/cmd/api/src/api/v2/agi.go b/cmd/api/src/api/v2/agi.go
@@ -1,17 +1,17 @@
 // Copyright 2023 Specter Ops, Inc.
-// 
+//
 // Licensed under the Apache License, Version 2.0
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
-// 
+//
 //     http://www.apache.org/licenses/LICENSE-2.0
-// 
+//
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-// 
+//
 // SPDX-License-Identifier: Apache-2.0
 
 package v2
@@ -23,10 +23,6 @@ import (
 	"strconv"
 	"strings"
 
-	"github.com/specterops/bloodhound/src/api"
-	"github.com/specterops/bloodhound/src/ctx"
-	"github.com/specterops/bloodhound/src/model"
-	"github.com/specterops/bloodhound/src/utils"
 	"github.com/gorilla/mux"
 	"github.com/specterops/bloodhound/analysis"
 	"github.com/specterops/bloodhound/dawgs/graph"
@@ -35,6 +31,10 @@ import (
 	"github.com/specterops/bloodhound/graphschema/common"
 	"github.com/specterops/bloodhound/headers"
 	"github.com/specterops/bloodhound/slices"
+	"github.com/specterops/bloodhound/src/api"
+	"github.com/specterops/bloodhound/src/ctx"
+	"github.com/specterops/bloodhound/src/model"
+	"github.com/specterops/bloodhound/src/utils"
 )
 
 // CreateAssetGroupRequest holds data required to create an asset group
@@ -234,6 +234,37 @@ func (s Resources) DeleteAssetGroup(response http.ResponseWriter, request *http.
 	}
 }
 
+func (s Resources) UpdateAssetGroupSelectors(response http.ResponseWriter, request *http.Request) {
+	var (
+		pathVars        = mux.Vars(request)
+		rawAssetGroupID = pathVars[api.URIPathVariableAssetGroupID]
+		selectorSpecs   []model.AssetGroupSelectorSpec
+	)
+
+	if assetGroupID, err := strconv.Atoi(rawAssetGroupID); err != nil {
+		api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, api.ErrorResponseDetailsIDMalformed, request), response)
+	} else if assetGroup, err := s.DB.GetAssetGroup(int32(assetGroupID)); err != nil {
+		api.HandleDatabaseError(request, response, err)
+	} else if err := api.ReadJSONRequestPayloadLimited(&selectorSpecs, request); err != nil {
+		api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, api.ErrorResponsePayloadUnmarshalError, request), response)
+	} else {
+		for _, selectorSpec := range selectorSpecs {
+			if err := selectorSpec.Validate(); err != nil {
+				api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, err.Error(), request), response)
+			}
+		}
+
+		if result, err := s.DB.UpdateAssetGroupSelectors(*ctx.FromRequest(request), assetGroup, selectorSpecs, false); err != nil {
+			api.HandleDatabaseError(request, response, err)
+		} else {
+			// When asset group selectors are modified we must trigger analysis
+			s.TaskNotifier.RequestAnalysis()
+
+			api.WriteBasicResponse(request.Context(), result, http.StatusOK, response)
+		}
+	}
+}
+
 func (s Resources) DeleteAssetGroupSelector(response http.ResponseWriter, request *http.Request) {
 	var (
 		pathVars                = mux.Vars(request)

diff --git a/cmd/api/src/database/agi.go b/cmd/api/src/database/agi.go
@@ -1,17 +1,17 @@
 // Copyright 2023 Specter Ops, Inc.
-// 
+//
 // Licensed under the Apache License, Version 2.0
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
-// 
+//
 //     http://www.apache.org/licenses/LICENSE-2.0
-// 
+//
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-// 
+//
 // SPDX-License-Identifier: Apache-2.0
 
 package database
@@ -189,47 +189,49 @@ func (s *BloodhoundDB) UpdateAssetGroupSelectors(ctx ctx.Context, assetGroup mod
 		removedSelectors = make([]model.AssetGroupSelector, 0)
 	)
 
-	return map[string]model.AssetGroupSelectors{
-			"added_selectors":   addedSelectors,
-			"removed_selectors": removedSelectors,
-		}, s.db.Transaction(func(tx *gorm.DB) error {
-			for _, selectorSpec := range selectorSpecs {
-				switch selectorSpec.Action {
-				case model.SelectorSpecActionAdd:
-					assetGroupSelector := model.AssetGroupSelector{
-						AssetGroupID:   assetGroup.ID,
-						Name:           selectorSpec.SelectorName,
-						Selector:       selectorSpec.EntityObjectID,
-						SystemSelector: systemSelector,
-					}
-
-					if result := tx.Create(&assetGroupSelector); result.Error != nil {
-						return CheckError(result)
-					} else {
-						addedSelectors = append(addedSelectors, assetGroupSelector)
-					}
-
-				case model.SelectorSpecActionRemove:
-					if result := tx.Where("asset_group_id=? AND name=?", assetGroup.ID, selectorSpec.SelectorName).Delete(&model.AssetGroupSelector{}); result.Error != nil {
-						return CheckError(result)
-					} else {
-						removedSelectors = append(removedSelectors, model.AssetGroupSelector{
-							AssetGroupID: assetGroup.ID,
-							Name:         selectorSpec.SelectorName,
-							Selector:     selectorSpec.EntityObjectID,
-						})
-					}
+	err := s.db.Transaction(func(tx *gorm.DB) error {
+		for _, selectorSpec := range selectorSpecs {
+			switch selectorSpec.Action {
+			case model.SelectorSpecActionAdd:
+				assetGroupSelector := model.AssetGroupSelector{
+					AssetGroupID:   assetGroup.ID,
+					Name:           selectorSpec.SelectorName,
+					Selector:       selectorSpec.EntityObjectID,
+					SystemSelector: systemSelector,
+				}
+
+				if result := tx.Create(&assetGroupSelector); result.Error != nil {
+					return CheckError(result)
+				} else {
+					addedSelectors = append(addedSelectors, assetGroupSelector)
 				}
 
-				if auditLog, err := newAuditLog(ctx, "UpdateAssetGroupSelectors", assetGroup.AuditData().MergeLeft(selectorSpec), s.idResolver); err != nil {
-					return err
-				} else if result := tx.Create(&auditLog); result.Error != nil {
-					return result.Error
+			case model.SelectorSpecActionRemove:
+				if result := tx.Where("asset_group_id=? AND name=?", assetGroup.ID, selectorSpec.SelectorName).Delete(&model.AssetGroupSelector{}); result.Error != nil {
+					return CheckError(result)
+				} else {
+					removedSelectors = append(removedSelectors, model.AssetGroupSelector{
+						AssetGroupID: assetGroup.ID,
+						Name:         selectorSpec.SelectorName,
+						Selector:     selectorSpec.EntityObjectID,
+					})
 				}
 			}
 
-			return nil
-		})
+			if auditLog, err := newAuditLog(ctx, "UpdateAssetGroupSelectors", assetGroup.AuditData().MergeLeft(selectorSpec), s.idResolver); err != nil {
+				return err
+			} else if result := tx.Create(&auditLog); result.Error != nil {
+				return result.Error
+			}
+		}
+
+		return nil
+	})
+
+	return map[string]model.AssetGroupSelectors{
+		"added_selectors":   addedSelectors,
+		"removed_selectors": removedSelectors,
+	}, err
 }
 
 func (s *BloodhoundDB) GetAllAssetGroupSelectors() (model.AssetGroupSelectors, error) {