ISSUE-1124-3: Simplifying the api (#1684)

ISSUE-1124-3: Simplifying the api using enums
deepfence · Oct 27, 2023 · 0382198 · 0382198
1 parent f6d0400
commit 0382198
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 28 deletions.
diff --git a/deepfence_server/model/scans.go b/deepfence_server/model/scans.go
@@ -128,12 +128,10 @@ type CloudComplianceScanListResp struct {
 }
 
 type ScanResultsMaskRequest struct {
-	ScanID                   string   `json:"scan_id" validate:"required" required:"true"`
-	ResultIDs                []string `json:"result_ids" validate:"required,gt=0,dive,min=1" required:"true"`
-	ScanType                 string   `json:"scan_type" validate:"required,oneof=SecretScan VulnerabilityScan MalwareScan ComplianceScan CloudComplianceScan" required:"true" enum:"SecretScan,VulnerabilityScan,MalwareScan,ComplianceScan,CloudComplianceScan"`
-	MaskAcrossHostsAndImages bool     `json:"mask_across_hosts_and_images"`
-	MaskAcrossImageTags      bool     `json:"mask_in_this_host_or_image_tags"`
-	MaskForImageTag          bool     `json:"mask_in_this_image_tag"`
+	ScanID     string   `json:"scan_id" validate:"required" required:"true"`
+	ResultIDs  []string `json:"result_ids" validate:"required,gt=0,dive,min=1" required:"true"`
+	ScanType   string   `json:"scan_type" validate:"required,oneof=SecretScan VulnerabilityScan MalwareScan ComplianceScan CloudComplianceScan" required:"true" enum:"SecretScan,VulnerabilityScan,MalwareScan,ComplianceScan,CloudComplianceScan"`
+	MaskAction string   `json:"mask_action" validate:"required,oneof=mask_global mask_all_image_tag mask_entity mask_image_tag" required:"true" enum:"mask_global,mask_all_image_tag,mask_entity,mask_image_tag"`
 }
 
 type ScanResultsActionRequest struct {

diff --git a/deepfence_server/reporters/scan/scan_result_actions.go b/deepfence_server/reporters/scan/scan_result_actions.go
@@ -62,45 +62,50 @@ func UpdateScanResultMasked(ctx context.Context, req *model.ScanResultsMaskReque
 	}
 	defer tx.Close()
 
-	if req.MaskAcrossHostsAndImages {
+	switch req.MaskAction {
+	case utils.MASK_GLOBAL:
 		nodeTag := utils.ScanTypeDetectedNode[utils.Neo4jScanType(req.ScanType)]
 		_, err = tx.Run(`
-		MATCH (o:`+nodeTag+`) -[:IS]-> (r)
-		WHERE o.node_id IN $node_ids
-		MATCH (n:`+nodeTag+`) -[:IS]-> (r)
-		MATCH (s) - [d:DETECTED] -> (n)
-		SET r.masked = $value, n.masked = $value, d.masked = $value
-		WITH s, n
-		MATCH (s) -[:SCANNED] ->(e)
-		MATCH (c:ContainerImage{node_id: e.docker_image_id}) -[:ALIAS] ->(t)
-		MERGE (t) -[m:MASKED]->(n)
-		SET m.masked = $value`,
+        MATCH (o:`+nodeTag+`) -[:IS]-> (r)
+        WHERE o.node_id IN $node_ids
+        MATCH (n:`+nodeTag+`) -[:IS]-> (r)
+        MATCH (s) - [d:DETECTED] -> (n)
+        SET r.masked = $value, n.masked = $value, d.masked = $value
+        WITH s, n
+        MATCH (s) -[:SCANNED] ->(e)
+        MATCH (c:ContainerImage{node_id: e.docker_image_id}) -[:ALIAS] ->(t)
+        MERGE (t) -[m:MASKED]->(n)
+        SET m.masked = $value`,
 			map[string]interface{}{"node_ids": req.ResultIDs, "value": value})
-	} else if req.MaskAcrossImageTags {
+
+	case utils.MASK_ALL_IMAGE_TAG, utils.MASK_ENTITY:
 		_, err = tx.Run(`
-		MATCH (s:`+string(req.ScanType)+`) - [d:DETECTED] -> (n)
+        MATCH (s:`+string(req.ScanType)+`) - [d:DETECTED] -> (n)
         WHERE n.node_id IN $node_ids
         SET n.masked = $value, d.masked = $value`,
 			map[string]interface{}{"node_ids": req.ResultIDs, "value": value})
-	} else if req.MaskForImageTag {
+
+	case utils.MASK_IMAGE_TAG:
 		_, err = tx.Run(`
         MATCH (s:`+string(req.ScanType)+`) -[d:DETECTED] -> (n)
         WHERE n.node_id IN $node_ids AND s.node_id=$scan_id
-		MATCH (s) -[:SCANNED] ->(e)
-		MATCH (c:ContainerImage{node_id: e.docker_image_id}) -[:ALIAS] ->(t)
-		MERGE (t) -[m:MASKED]->(n)
-		SET m.masked = $value, d.masked = $value`,
+        MATCH (s) -[:SCANNED] ->(e)
+        MATCH (c:ContainerImage{node_id: e.docker_image_id}) -[:ALIAS] ->(t)
+        MERGE (t) -[m:MASKED]->(n)
+        SET m.masked = $value, d.masked = $value`,
 			map[string]interface{}{"node_ids": req.ResultIDs, "value": value,
 				"scan_id": req.ScanID})
 
-	} else {
+	default:
 		_, err = tx.Run(`
-		MATCH (m:`+string(req.ScanType)+`) -[d:DETECTED] -> (n)
-		WHERE n.node_id IN $node_ids AND m.node_id=$scan_id
-		SET d.masked = $value`,
+        MATCH (m:`+string(req.ScanType)+`) -[d:DETECTED] -> (n)
+        WHERE n.node_id IN $node_ids AND m.node_id=$scan_id
+        SET d.masked = $value`,
 			map[string]interface{}{"node_ids": req.ResultIDs, "value": value,
 				"scan_id": req.ScanID})
+
 	}
+
 	if err != nil {
 		return err
 	}

diff --git a/deepfence_utils/utils/constants.go b/deepfence_utils/utils/constants.go
@@ -217,3 +217,10 @@ const (
 	ReportXLSX ReportType = "xlsx"
 	ReportPDF  ReportType = "pdf"
 )
+
+const (
+	MASK_GLOBAL        = "mask_global"
+	MASK_ALL_IMAGE_TAG = "mask_all_image_tag"
+	MASK_ENTITY        = "mask_entity"
+	MASK_IMAGE_TAG     = "mask_image_tag"
+)