chore: cleanup

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

Hide file tree

Showing 4 changed files with 28 additions and 59 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
@@ -29,24 +29,15 @@ import (
 type SavedQueriesPermissionsData interface {
 	CreateSavedQueryPermissionToPublic(ctx context.Context, queryID int64) (model.SavedQueriesPermissions, error)
 	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
+	GetScopeForSavedQuery(ctx context.Context, queryID int64, userID uuid.UUID) (SavedQueryScopeMap, 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
 
-// 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{
-		QueryID:        queryID,
-		SharedToUserID: NullUUID(userID),
-		Public:         false,
-	}
-
-	return permission, CheckError(s.db.WithContext(ctx).Create(&permission))
-}
-
 // CreateSavedQueryPermissionToPublic creates a new entry to the SavedQueriesPermissions table granting public read permissions to all users
 func (s *BloodhoundDB) CreateSavedQueryPermissionToPublic(ctx context.Context, queryID int64) (model.SavedQueriesPermissions, error) {
 	permission := model.SavedQueriesPermissions{
@@ -85,6 +76,17 @@ func (s *BloodhoundDB) CreateSavedQueryPermissionsToUsers(ctx context.Context, q
 	return newPermissions, CheckError(result)
 }
 
+// DeleteSavedQueryPermissionsForUsers batch deletes permissions associated a query id and a list of users
+// If no user ids are supplied, all records for query id are deleted
+func (s *BloodhoundDB) DeleteSavedQueryPermissionsForUsers(ctx context.Context, queryID int64, userIDs ...uuid.UUID) error {
+	result := s.db.WithContext(ctx).Table("saved_queries_permissions").Where("query_id = ?", queryID)
+	if len(userIDs) > 0 {
+		result = result.Where("shared_to_user_id IN ?", userIDs)
+	}
+
+	return CheckError(result.Delete(&model.SavedQueriesPermissions{}))
+}
+
 // GetScopeForSavedQuery will return a map of the possible scopes given a query id and a user id
 func (s *BloodhoundDB) GetScopeForSavedQuery(ctx context.Context, queryID int64, userID uuid.UUID) (SavedQueryScopeMap, error) {
 	scopes := SavedQueryScopeMap{
@@ -120,13 +122,18 @@ func (s *BloodhoundDB) GetScopeForSavedQuery(ctx context.Context, queryID int64,
 	return scopes, nil
 }
 
-// DeleteSavedQueryPermissionsForUsers batch deletes permissions associated a query id and a list of users
-// If no user ids are supplied, all records for query id are deleted
-func (s *BloodhoundDB) DeleteSavedQueryPermissionsForUsers(ctx context.Context, queryID int64, userIDs ...uuid.UUID) error {
-	result := s.db.WithContext(ctx).Table("saved_queries_permissions").Where("query_id = ?", queryID)
-	if len(userIDs) > 0 {
-		result = result.Where("shared_to_user_id IN ?", userIDs)
-	}
+// 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 CheckError(result.Delete(&model.SavedQueriesPermissions{}))
+	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)
 }