fix: lint issues and refactor

deepfence · Jun 11, 2024 · a40c687 · a40c687
1 parent a34fae6
commit a40c687
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 30 deletions.
diff --git a/deepfence_server/handler/notification.go b/deepfence_server/handler/notification.go
@@ -32,8 +32,7 @@ func (h *Handler) GetScansHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// respond with scans
-	err = httpext.JSON(w, http.StatusOK, scans)
-	return
+	httpext.JSON(w, http.StatusOK, scans)
 }
 
 func (h *Handler) MarkScansReadHandler(w http.ResponseWriter, r *http.Request) {
@@ -57,8 +56,7 @@ func (h *Handler) MarkScansReadHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// respond with success
-	err = httpext.JSON(w, http.StatusOK, nil)
-	return
+	httpext.JSON(w, http.StatusOK, nil)
 }
 
 /* Registry Sync Handlers */
@@ -76,8 +74,7 @@ func (h *Handler) GetRegistrySyncHandler(w http.ResponseWriter, r *http.Request)
 	}
 
 	// respond with registries
-	err = httpext.JSON(w, http.StatusOK, registries)
-	return
+	httpext.JSON(w, http.StatusOK, registries)
 }
 
 /* Integration Handlers */
@@ -95,6 +92,5 @@ func (h *Handler) GetIntegrationFailuresHandler(w http.ResponseWriter, r *http.R
 	}
 
 	// respond with integrations
-	err = httpext.JSON(w, http.StatusOK, integrations)
-	return
+	httpext.JSON(w, http.StatusOK, integrations)
 }
diff --git a/deepfence_server/model/notification.go b/deepfence_server/model/notification.go
@@ -1,10 +1,11 @@
 package model
 
 type NotificationGetScanResponse struct {
-	VulnerabilityScan []Scan `json:"vulnerability_scan"`
-	SecretScan        []Scan `json:"secret_scan"`
-	MalwareScan       []Scan `json:"malware_scan"`
-	PostureScan       []Scan `json:"posture_scan"`
+	VulnerabilityScan   []Scan `json:"vulnerability_scan"`
+	SecretScan          []Scan `json:"secret_scan"`
+	MalwareScan         []Scan `json:"malware_scan"`
+	ComplianceScan      []Scan `json:"compliance_scan"`
+	CloudComplianceScan []Scan `json:"cloud_compliance_scan"`
 }
 
 type Scan struct {

diff --git a/deepfence_server/reporters/notification/scan.go b/deepfence_server/reporters/notification/scan.go
@@ -8,59 +8,70 @@ import (
 	"github.com/deepfence/ThreatMapper/deepfence_server/model"
 	"github.com/deepfence/ThreatMapper/deepfence_utils/directory"
 	"github.com/deepfence/ThreatMapper/deepfence_utils/log"
+	"github.com/deepfence/ThreatMapper/deepfence_utils/utils"
 	"github.com/neo4j/neo4j-go-driver/v5/neo4j"
 )
 
 func GetScans(ctx context.Context, scanTypes []string, statues []string) (model.NotificationGetScanResponse, error) {
 	response := model.NotificationGetScanResponse{}
 	var err error
-	for _, scanType := range scanTypes {
+	for _, st := range scanTypes {
+		scanType := utils.DetectedNodeScanType[st]
 		switch scanType {
-		case "vulnerability":
-			response.VulnerabilityScan, err = GetScansFor(ctx, "VulnerabilityScan", statues)
+		case utils.NEO4JVulnerabilityScan:
+			response.VulnerabilityScan, err = GetScansFor(ctx, scanType, statues)
 			if err != nil {
 				return response, err
 			}
-		case "secret":
-			response.SecretScan, err = GetScansFor(ctx, "SecretScan", statues)
+		case utils.NEO4JSecretScan:
+			response.SecretScan, err = GetScansFor(ctx, scanType, statues)
 			if err != nil {
 				return response, err
 			}
-		case "malware":
-			response.MalwareScan, err = GetScansFor(ctx, "MalwareScan", statues)
+		case utils.NEO4JMalwareScan:
+			response.MalwareScan, err = GetScansFor(ctx, scanType, statues)
 			if err != nil {
 				return response, err
 			}
-		case "posture":
-			response.PostureScan, err = GetScansFor(ctx, "PostureScan", statues)
+		case utils.NEO4JComplianceScan:
+			response.ComplianceScan, err = GetScansFor(ctx, scanType, statues)
 			if err != nil {
 				return response, err
 			}
-		case "all":
-			response.VulnerabilityScan, err = GetScansFor(ctx, "VulnerabilityScan", statues)
+		case utils.NEO4JCloudComplianceScan:
+			response.CloudComplianceScan, err = GetScansFor(ctx, scanType, statues)
 			if err != nil {
 				return response, err
 			}
-			response.SecretScan, err = GetScansFor(ctx, "SecretScan", statues)
+		case "":
+			response.VulnerabilityScan, err = GetScansFor(ctx, utils.NEO4JVulnerabilityScan, statues)
 			if err != nil {
 				return response, err
 			}
-			response.MalwareScan, err = GetScansFor(ctx, "MalwareScan", statues)
+			response.SecretScan, err = GetScansFor(ctx, utils.NEO4JSecretScan, statues)
 			if err != nil {
 				return response, err
 			}
-			response.PostureScan, err = GetScansFor(ctx, "PostureScan", statues)
+			response.MalwareScan, err = GetScansFor(ctx, utils.NEO4JMalwareScan, statues)
+			if err != nil {
+				return response, err
+			}
+			response.ComplianceScan, err = GetScansFor(ctx, utils.NEO4JComplianceScan, statues)
+			if err != nil {
+				return response, err
+			}
+			response.CloudComplianceScan, err = GetScansFor(ctx, utils.NEO4JCloudComplianceScan, statues)
 			if err != nil {
 				return response, err
 			}
 		default:
-			return response, fmt.Errorf("Invalid scan type")
+			return response, fmt.Errorf("invalid scan type")
 		}
 	}
 	return response, nil
 }
 
-func GetScansFor(ctx context.Context, scanType string, statues []string) ([]model.Scan, error) {
+func GetScansFor(ctx context.Context, scanType utils.Neo4jScanType, statues []string) ([]model.Scan, error) {
 	scans := []model.Scan{}
 	driver, err := directory.Neo4jClient(ctx)
 	if err != nil {
@@ -79,13 +90,13 @@ func GetScansFor(ctx context.Context, scanType string, statues []string) ([]mode
 	}
 	defer tx.Close(ctx)
 	query := `
-	MATCH (n:` + scanType + `)
+	MATCH (n:` + string(scanType) + `)
 	WHERE n.status IN $statues
 	AND n.acknowledged_at IS NULL
 	RETURN n.created_at, n.updated_at, n.node_id, n.is_priority, n.status, n.status_message, n.trigger_action, n.retries`
 	if len(statues) == 0 {
 		query = `
-		MATCH (n:` + scanType + `)
+		MATCH (n:` + string(scanType) + `)
 		WHERE n.acknowledged_at IS NULL
 		RETURN n.created_at, n.updated_at, n.node_id, n.is_priority, n.status, n.status_message, n.trigger_action, n.retries`
 	}