From fffb5f1718ce4fa21847d1a28cf2d6a49a6f0585 Mon Sep 17 00:00:00 2001
From: Thomas Legris <thomas@deepfence.io>
Date: Tue, 28 Nov 2023 22:09:25 +0900
Subject: [PATCH] Fix scan status

---
 deepfence_worker/cronjobs/neo4j.go   | 33 ++++++++++++++++------------
 deepfence_worker/ingesters/common.go | 10 ++++-----
 2 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/deepfence_worker/cronjobs/neo4j.go b/deepfence_worker/cronjobs/neo4j.go
index 6f6f7aa421..33544fbfed 100644
--- a/deepfence_worker/cronjobs/neo4j.go
+++ b/deepfence_worker/cronjobs/neo4j.go
@@ -12,6 +12,7 @@ import (
 	"github.com/deepfence/ThreatMapper/deepfence_utils/log"
 	"github.com/deepfence/ThreatMapper/deepfence_utils/utils"
 	"github.com/neo4j/neo4j-go-driver/v4/neo4j"
+	ingestersUtil "github.com/deepfence/ThreatMapper/deepfence_utils/utils/ingesters"
 )
 
 const (
@@ -219,7 +220,7 @@ func CleanUpDB(ctx context.Context, task *asynq.Task) error {
 	if _, err = session.Run(`
 		MATCH (n:ContainerImage)
 		WHERE n.active = false
-		AND ((NOT exists((n) <-[:SCANNED]-()) 
+		AND ((NOT exists((n) <-[:SCANNED]-())
 		AND n.updated_at < TIMESTAMP() - $delete_threshold_ms)
 		OR n.updated_at < TIMESTAMP()-$old_time_ms)
 		WITH n LIMIT 10000
@@ -235,8 +236,8 @@ func CleanUpDB(ctx context.Context, task *asynq.Task) error {
 	if _, err = session.Run(`
 		MATCH (n:Container)
 		WHERE n.active = false
-		AND ((NOT exists((n) <-[:SCANNED]-()) 
-		AND n.updated_at < TIMESTAMP() - $delete_threshold_ms) 
+		AND ((NOT exists((n) <-[:SCANNED]-())
+		AND n.updated_at < TIMESTAMP() - $delete_threshold_ms)
 		OR n.updated_at < TIMESTAMP()-$old_time_ms)
 		WITH n LIMIT 10000
 		DETACH DELETE n`,
@@ -282,17 +283,21 @@ func CleanUpDB(ctx context.Context, task *asynq.Task) error {
 		return err
 	}
 
-	if _, err = session.Run(`
-		MATCH (n) -[:SCANNED]-> ()
-		WHERE n.retries >= 3
-		WITH n LIMIT 10000
-		SET n.status = $new_status`,
-		map[string]interface{}{
-			"time_ms":    dbScanTimeout.Milliseconds(),
-			"new_status": utils.ScanStatusFailed,
-		}, txConfig); err != nil {
-		log.Error().Msgf("Error in Clean up DB task: %v", err)
-		return err
+	for ts := range ingestersUtil.ScanStatusField {
+		if _, err = session.Run(`
+			MATCH (n:`+string(ts)+`) -[:SCANNED]-> (r)
+			WHERE n.retries >= 3
+			WITH n, r LIMIT 10000
+			SET n.status = $new_status,
+				r.`+ingestersUtil.ScanStatusField[ts]+`=n.status,
+				r.`+ingestersUtil.LatestScanIDField[ts]+`=n.node_id`,
+			map[string]interface{}{
+				"time_ms":    dbScanTimeout.Milliseconds(),
+				"new_status": utils.ScanStatusFailed,
+			}, txConfig); err != nil {
+			log.Error().Msgf("Error in Clean up DB task: %v", err)
+			return err
+		}
 	}
 
 	if _, err = session.Run(`
diff --git a/deepfence_worker/ingesters/common.go b/deepfence_worker/ingesters/common.go
index 46aae4754f..2b399f10ac 100644
--- a/deepfence_worker/ingesters/common.go
+++ b/deepfence_worker/ingesters/common.go
@@ -46,7 +46,7 @@ func CommitFuncStatus[Status any](ts utils.Neo4jScanType) func(ns string, data [
 			n.status_message = row.scan_message,
 			n.updated_at = TIMESTAMP()
 		WITH n
-		OPTIONAL MATCH (m) -[:DETECTED]- (n)
+		OPTIONAL MATCH (n) -[:DETECTED]- (m)
 		WITH n, count(m) as m_count
 		MATCH (n) -[:SCANNED]- (r)
 		SET r.` + ingestersUtil.ScanStatusField[ts] + `=n.status,
@@ -214,10 +214,10 @@ func getEntityIdFromScanID(scanId, scanType string,
 	entityId := ""
 	query := `MATCH (s:` + scanType + `{node_id:'` + scanId + `'}) - [:SCANNED] -> (n)
 		WITH labels(n) as label, n
-		RETURN 
-		CASE 
-    		WHEN 'ContainerImage' IN label or 'Container' in label 
-			THEN [(ci:ContainerImage{node_id:n.docker_image_id}) - [:IS] -> (cis) | cis.node_id] 
+		RETURN
+		CASE
+    		WHEN 'ContainerImage' IN label or 'Container' in label
+			THEN [(ci:ContainerImage{node_id:n.docker_image_id}) - [:IS] -> (cis) | cis.node_id]
     		ELSE [n.node_id]
 		END`
 	res, err := tx.Run(query, map[string]interface{}{})