chore: cleanup

SpecterOps · Aug 23, 2024 · a038e15 · a038e15
1 parent 5f7a1f5
commit a038e15
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 39 deletions.
diff --git a/cmd/api/src/api/v2/saved_queries.go b/cmd/api/src/api/v2/saved_queries.go
@@ -234,7 +234,7 @@ func (s Resources) DeleteSavedQuery(response http.ResponseWriter, request *http.
 			if _, isAdmin := user.Roles.FindByName(auth.RoleAdministrator); !isAdmin {
 				api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusForbidden, "User does not have permission to delete this query", request), response)
 				return
-			} else if isPublicQuery, err := s.DB.IsSavedQueryPublic(request.Context(), int64(savedQueryID)); err != nil {
+			} else if isPublicQuery, err := s.DB.IsSavedQueryPublic(request.Context(), savedQueryID); err != nil {
 				api.HandleDatabaseError(request, response, err)
 				return
 			} else if !isPublicQuery {

diff --git a/cmd/api/src/database/mocks/db.go b/cmd/api/src/database/mocks/db.go
diff --git a/cmd/api/src/database/saved_queries.go b/cmd/api/src/database/saved_queries.go
@@ -33,8 +33,6 @@ type SavedQueriesData interface {
 	SavedQueryBelongsToUser(ctx context.Context, userID uuid.UUID, savedQueryID int64) (bool, error)
 	GetSharedSavedQueries(ctx context.Context, userID uuid.UUID) (model.SavedQueries, error)
 	GetPublicSavedQueries(ctx context.Context) (model.SavedQueries, error)
-	IsSavedQueryPublic(ctx context.Context, savedQueryID int64) (bool, error)
-	IsSavedQuerySharedToUser(ctx context.Context, queryID int64, userID uuid.UUID) (bool, error)
 }
 
 func (s *BloodhoundDB) GetSavedQuery(ctx context.Context, savedQueryID int64) (model.SavedQuery, error) {
@@ -118,24 +116,3 @@ func (s *BloodhoundDB) GetPublicSavedQueries(ctx context.Context) (model.SavedQu
 
 	return savedQueries, CheckError(result)
 }
-
-func (s *BloodhoundDB) IsSavedQueryPublic(ctx context.Context, savedQueryID int64) (bool, error) {
-	if publicQueries, err := s.GetPublicSavedQueries(ctx); err != nil {
-		return false, err
-	} else {
-		for _, publicQuery := range publicQueries {
-			if publicQuery.ID == savedQueryID {
-				return true, nil
-			}
-		}
-		return false, nil
-	}
-}
-
-// IsSavedQuerySharedToUser returns true or false whether a provided saved query is shared with a provided user
-func (s *BloodhoundDB) IsSavedQuerySharedToUser(ctx context.Context, queryID int64, userID uuid.UUID) (bool, error) {
-	rows := int64(0)
-	result := s.db.WithContext(ctx).Table("saved_queries_permissions").Where("query_id = ? AND shared_to_user_id = ?", queryID, userID).Count(&rows)
-
-	return rows > 0, CheckError(result)
-}
diff --git a/cmd/api/src/database/savedqueriespermissions.go b/cmd/api/src/database/savedqueriespermissions.go
@@ -31,11 +31,29 @@ type SavedQueriesPermissionsData interface {
 	CreateSavedQueryPermissionsToUsers(ctx context.Context, queryID int64, userIDs ...uuid.UUID) ([]model.SavedQueriesPermissions, error)
 	GetScopeForSavedQuery(ctx context.Context, queryID int64, userID uuid.UUID) (SavedQueryScopeMap, error)
 	DeleteSavedQueryPermissionsForUsers(ctx context.Context, queryID int64, userIDs ...uuid.UUID) error
+	IsSavedQueryPublic(ctx context.Context, savedQueryID int64) (bool, error)
+	IsSavedQuerySharedToUser(ctx context.Context, queryID int64, userID uuid.UUID) (bool, error)
 }
 
 // SavedQueryScopeMap holds the information of a saved query's scope [IE: owned, shared, public]
 type SavedQueryScopeMap map[model.SavedQueryScope]bool
 
+// IsSavedQueryPublic returns true or false whether a provided saved query is public
+func (s *BloodhoundDB) IsSavedQueryPublic(ctx context.Context, queryID int64) (bool, error) {
+	rows := int64(0)
+	result := s.db.WithContext(ctx).Select("saved_queries_permissions.*").Where("public = true AND query_id=", queryID).Count(&rows)
+
+ 	return rows > 0, CheckError(result)
+}
+
+// IsSavedQuerySharedToUser returns true or false whether a provided saved query is shared with a provided user
+func (s *BloodhoundDB) IsSavedQuerySharedToUser(ctx context.Context, queryID int64, userID uuid.UUID) (bool, error) {
+	rows := int64(0)
+	result := s.db.WithContext(ctx).Table("saved_queries_permissions").Where("query_id = ? AND shared_to_user_id = ?", queryID, userID).Count(&rows)
+
+	return rows > 0, CheckError(result)
+}
+
 // CreateSavedQueryPermissionToUser creates a new entry to the SavedQueriesPermissions table granting a provided user id to access a provided query
 func (s *BloodhoundDB) CreateSavedQueryPermissionToUser(ctx context.Context, queryID int64, userID uuid.UUID) (model.SavedQueriesPermissions, error) {
 	permission := model.SavedQueriesPermissions{