Skip to content

Commit

Permalink
refactor: move in-memory analysis request state to postgres (#631)
Browse files Browse the repository at this point in the history
  • Loading branch information
mistahj67 authored Jun 7, 2024
1 parent f37f4c0 commit dffec24
Show file tree
Hide file tree
Showing 20 changed files with 549 additions and 125 deletions.
1 change: 1 addition & 0 deletions cmd/api/src/api/registration/v2.go
Original file line number Diff line number Diff line change
Expand Up @@ -286,6 +286,7 @@ func NewV2API(cfg config.Configuration, resources v2.Resources, routerInst *rout
// Datapipe API
routerInst.GET("/api/v2/datapipe/status", resources.GetDatapipeStatus).RequireAuth(),
//TODO: Update the permission on this once we get something more concrete
routerInst.GET("/api/v2/analysis/status", resources.GetAnalysisRequest).RequirePermissions(permissions.GraphDBRead),
routerInst.PUT("/api/v2/analysis", resources.RequestAnalysis).RequirePermissions(permissions.GraphDBWrite),
)
}
14 changes: 13 additions & 1 deletion cmd/api/src/api/v2/agi.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ import (
"github.com/specterops/bloodhound/headers"
"github.com/specterops/bloodhound/log"
"github.com/specterops/bloodhound/src/api"
"github.com/specterops/bloodhound/src/auth"
"github.com/specterops/bloodhound/src/ctx"
"github.com/specterops/bloodhound/src/model"
"github.com/specterops/bloodhound/src/utils"
Expand Down Expand Up @@ -265,7 +266,18 @@ func (s Resources) UpdateAssetGroupSelectors(response http.ResponseWriter, reque

if assetGroup.Tag == model.TierZeroAssetGroupTag {
// When T0 asset group selectors are modified, entire analysis must be re-run
s.TaskNotifier.RequestAnalysis()
var userId string
if user, isUser := auth.GetUserFromAuthCtx(ctx.FromRequest(request).AuthCtx); !isUser {
log.Warnf("encountered request analysis for unknown user, this shouldn't happen")
userId = "unknown-user-update-asset-group-selectors"
} else {
userId = user.ID.String()
}

if err := s.DB.RequestAnalysis(request.Context(), userId); err != nil {
api.HandleDatabaseError(request, response, err)
return
}
}

api.WriteBasicResponse(request.Context(), result, http.StatusCreated, response)
Expand Down
32 changes: 11 additions & 21 deletions cmd/api/src/api/v2/agi_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,25 +25,24 @@ import (
"net/url"
"testing"

"github.com/specterops/bloodhound/headers"
"github.com/specterops/bloodhound/mediatypes"
"github.com/specterops/bloodhound/src/auth"
"github.com/specterops/bloodhound/src/test/must"

"github.com/gofrs/uuid"
"github.com/gorilla/mux"
"github.com/specterops/bloodhound/dawgs/graph"
"github.com/specterops/bloodhound/errors"
"github.com/specterops/bloodhound/graphschema/ad"
"github.com/specterops/bloodhound/graphschema/azure"
"github.com/specterops/bloodhound/graphschema/common"
"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/api/v2/apitest"
"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"
queriesMocks "github.com/specterops/bloodhound/src/queries/mocks"
"github.com/specterops/bloodhound/src/test/must"
"github.com/specterops/bloodhound/src/utils/test"
"github.com/stretchr/testify/require"
"go.uber.org/mock/gomock"
Expand Down Expand Up @@ -592,16 +591,11 @@ func TestResources_UpdateAssetGroupSelectors_SuccessT0(t *testing.T) {
mockDB.EXPECT().UpdateAssetGroupSelectors(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(expectedResult, nil)
mockGraph.EXPECT().UpdateSelectorTags(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil)

mockTasker := datapipeMocks.NewMockTasker(mockCtrl)
// MockTasker should receive a call to RequestAnalysis() since this is a Tier Zero Asset group.
// Should receive a call to RequestAnalysis() since this is a Tier Zero Asset group.
// Analysis must be run upon updating a T0 AG
mockTasker.EXPECT().RequestAnalysis()
mockDB.EXPECT().RequestAnalysis(gomock.Any(), uuid.UUID{}.String())

handlers := v2.Resources{
DB: mockDB,
TaskNotifier: mockTasker,
GraphQuery: mockGraph,
}
handlers := v2.Resources{DB: mockDB, GraphQuery: mockGraph}

response := httptest.NewRecorder()
handler := http.HandlerFunc(handlers.UpdateAssetGroupSelectors)
Expand Down Expand Up @@ -687,15 +681,11 @@ func TestResources_UpdateAssetGroupSelectors_SuccessOwned(t *testing.T) {

mockGraph.EXPECT().UpdateSelectorTags(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil)

mockTasker := datapipeMocks.NewMockTasker(mockCtrl)
// NOTE MockTasker should NOT receive a call to RequestAnalysis() since this is not a Tier Zero Asset group.
// NOTE should NOT receive a call to RequestAnalysis() since this is not a Tier Zero Asset group.
// Analysis should not be re-run when a non T0 AG is updated
mockDB.EXPECT().RequestAnalysis(gomock.Any(), uuid.UUID{}.String()).Times(0)

handlers := v2.Resources{
DB: mockDB,
TaskNotifier: mockTasker,
GraphQuery: mockGraph,
}
handlers := v2.Resources{DB: mockDB, GraphQuery: mockGraph}

response := httptest.NewRecorder()
handler := http.HandlerFunc(handlers.UpdateAssetGroupSelectors)
Expand Down
57 changes: 57 additions & 0 deletions cmd/api/src/api/v2/analysisrequest.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
// Copyright 2024 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

import (
"database/sql"
"net/http"

"github.com/specterops/bloodhound/errors"
"github.com/specterops/bloodhound/log"
"github.com/specterops/bloodhound/src/api"
"github.com/specterops/bloodhound/src/auth"
"github.com/specterops/bloodhound/src/ctx"
)

func (s Resources) GetAnalysisRequest(response http.ResponseWriter, request *http.Request) {
if analRequest, err := s.DB.GetAnalysisRequest(request.Context()); err != nil && !errors.Is(err, sql.ErrNoRows) {
api.HandleDatabaseError(request, response, err)
} else if errors.Is(err, sql.ErrNoRows) {
api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusNotFound, api.ErrorResponseDetailsResourceNotFound, request), response)
} else {
api.WriteBasicResponse(request.Context(), analRequest, http.StatusOK, response)
}
}

func (s Resources) RequestAnalysis(response http.ResponseWriter, request *http.Request) {
defer log.Measure(log.LevelDebug, "Requesting analysis")()

var userId string
if user, isUser := auth.GetUserFromAuthCtx(ctx.FromRequest(request).AuthCtx); !isUser {
log.Warnf("encountered request analysis for unknown user, this shouldn't happen")
userId = "unknown-user"
} else {
userId = user.ID.String()
}

if err := s.DB.RequestAnalysis(request.Context(), userId); err != nil {
api.HandleDatabaseError(request, response, err)
return
}

response.WriteHeader(http.StatusAccepted)
}
36 changes: 36 additions & 0 deletions cmd/api/src/api/v2/analysisrequest_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
// Copyright 2024 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_test

import (
"testing"

"github.com/specterops/bloodhound/src/api/v2/integration"
"github.com/specterops/bloodhound/src/model"
"github.com/stretchr/testify/require"
)

func TestRequestAnalysis(t *testing.T) {
testCtx := integration.NewFOSSContext(t)

err := testCtx.AdminClient().RequestAnalysis()
require.Nil(t, err)

analReq, err := testCtx.AdminClient().GetAnalysisRequest()
require.Nil(t, err)
require.Equal(t, analReq.RequestType, model.AnalysisRequestAnalysis)
}
54 changes: 54 additions & 0 deletions cmd/api/src/api/v2/apiclient/analysisrequest.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
// Copyright 2024 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 apiclient

import (
"net/http"

"github.com/specterops/bloodhound/src/api"
"github.com/specterops/bloodhound/src/model"
)

func (s Client) RequestAnalysis() error {
if response, err := s.Request(http.MethodPut, "api/v2/analysis", nil, nil); err != nil {
return err
} else {
defer response.Body.Close()

if api.IsErrorResponse(response) {
return ReadAPIError(response)
}

return nil
}
}

func (s Client) GetAnalysisRequest() (model.AnalysisRequest, error) {
var analReq model.AnalysisRequest

if response, err := s.Request(http.MethodGet, "api/v2/analysis/status", nil, nil); err != nil {
return analReq, err
} else {
defer response.Body.Close()

if api.IsErrorResponse(response) {
return analReq, ReadAPIError(response)
}

return analReq, api.ReadAPIV2ResponsePayload(&analReq, response)
}
}
30 changes: 27 additions & 3 deletions cmd/api/src/api/v2/database_wipe.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ import (

"github.com/specterops/bloodhound/log"
"github.com/specterops/bloodhound/src/api"
"github.com/specterops/bloodhound/src/auth"
"github.com/specterops/bloodhound/src/ctx"
"github.com/specterops/bloodhound/src/model"
"github.com/specterops/bloodhound/src/model/appcfg"
)
Expand Down Expand Up @@ -108,7 +110,18 @@ func (s Resources) HandleDatabaseWipe(response http.ResponseWriter, request *htt
)
return
} else {
s.TaskNotifier.RequestDeletion()
var userId string
if user, isUser := auth.GetUserFromAuthCtx(ctx.FromRequest(request).AuthCtx); !isUser {
log.Warnf("encountered request analysis for unknown user, this shouldn't happen")
userId = "unknown-user-database-wipe"
} else {
userId = user.ID.String()
}

if err := s.DB.RequestCollectedGraphDataDeletion(request.Context(), userId); err != nil {
api.HandleDatabaseError(request, response, err)
return
}
s.handleAuditLogForDatabaseWipe(request.Context(), &auditEntry, true, "collected graph data")
}

Expand All @@ -125,7 +138,18 @@ func (s Resources) HandleDatabaseWipe(response http.ResponseWriter, request *htt

// if deleting `nodes` or deleting `asset group selectors` is successful, kickoff an analysis
if kickoffAnalysis {
s.TaskNotifier.RequestAnalysis()
var userId string
if user, isUser := auth.GetUserFromAuthCtx(ctx.FromRequest(request).AuthCtx); !isUser {
log.Warnf("encountered request analysis for unknown user, this shouldn't happen")
userId = "unknown-user-database-wipe"
} else {
userId = user.ID.String()
}

if err := s.DB.RequestAnalysis(request.Context(), userId); err != nil {
api.HandleDatabaseError(request, response, err)
return
}
}

// delete file ingest history
Expand All @@ -142,7 +166,7 @@ func (s Resources) HandleDatabaseWipe(response http.ResponseWriter, request *htt
}
}

// return a user friendly error message indicating what operations failed
// return a user-friendly error message indicating what operations failed
if len(errors) > 0 {
api.WriteErrorResponse(
request.Context(),
Expand Down
Loading

0 comments on commit dffec24

Please sign in to comment.