diff --git a/deepfence_server/controls/agent.go b/deepfence_server/controls/agent.go index d60a1db2ec..634d32d935 100644 --- a/deepfence_server/controls/agent.go +++ b/deepfence_server/controls/agent.go @@ -13,92 +13,64 @@ import ( ) func GetAgentActions(ctx context.Context, nodeId string) ([]controls.Action, error) { + // Aappend more actions here + return ExtractStartingAgentScans(ctx, nodeId) +} + +func GetPendingAgentScans(ctx context.Context, nodeId string) ([]controls.Action, error) { + res := []controls.Action{} + if len(nodeId) == 0 { + return res, errors.New("Missing node_id") + } + client, err := directory.Neo4jClient(ctx) if err != nil { - return nil, err + return res, err } session, err := client.Session(neo4j.AccessModeWrite) if err != nil { - return nil, err + return res, err } defer session.Close() tx, err := session.BeginTransaction() if err != nil { - return nil, err + return res, err } defer tx.Close() - r, err := tx.Run("match (n:Node{node_id:$id}) RETURN n.actions", map[string]interface{}{"id": nodeId}) + r, err := tx.Run(`MATCH (s) -[:SCANNED]-> (n:Node{node_id:$id}) WHERE NOT (s.status = '`+utils.SCAN_STATUS_SUCCESS+`') AND s.retries < 3 SET s.retries = s.retries + 1 WITH s RETURN s.trigger_action`, map[string]interface{}{"id": nodeId}) if err != nil { - log.Error().Msgf("neo4j req err: %v", err) - return nil, err + return res, err } - actions, err := r.Single() + records, err := r.Collect() if err != nil { - return nil, err + return res, err } - res := []controls.Action{} - if actions.Values[0] != nil { - for _, action := range actions.Values[0].([]interface{}) { - entry := controls.Action{} - err = json.Unmarshal([]byte(action.(string)), &entry) - if err != nil { - return res, err - } - res = append(res, entry) + for _, record := range records { + var action controls.Action + if record.Values[0] == nil { + log.Error().Msgf("Invalid neo4j trigger_action result, skipping") + continue } - } - - _, err = tx.Run("match (n:Node{node_id:$id}) SET n.actions = []", map[string]interface{}{"id": nodeId}) - - return res, tx.Commit() -} - -func SetAgentActions(ctx context.Context, nodeId string, actions []controls.Action) error { - client, err := directory.Neo4jClient(ctx) - if err != nil { - return err - } - - session, err := client.Session(neo4j.AccessModeWrite) - if err != nil { - return err - } - defer session.Close() - - tx, err := session.BeginTransaction() - if err != nil { - return err - } - defer tx.Close() - - act_bytes := []string{} - for _, act := range actions { - bytes, err := json.Marshal(act) + err := json.Unmarshal([]byte(record.Values[0].(string)), &action) if err != nil { - log.Error().Msgf("neo4j marshal err: %v", err) - return err + log.Error().Msgf("Unmarshal of action failed: %v", err) + continue } - act_bytes = append(act_bytes, string(bytes)) - } - - _, err = tx.Run("match (n:Node{node_id:$id}) UNWIND $data as act SET n.actions = coalesce(n.actions,[]) + act", map[string]interface{}{"id": nodeId, "data": act_bytes}) - - if err != nil { - return err + res = append(res, action) } - return tx.Commit() + return res, tx.Commit() } -func GetPendingAgentScans(ctx context.Context, nodeId string) ([]controls.Action, error) { +func ExtractStartingAgentScans(ctx context.Context, nodeId string) ([]controls.Action, error) { res := []controls.Action{} if len(nodeId) == 0 { return res, errors.New("Missing node_id") @@ -121,7 +93,7 @@ func GetPendingAgentScans(ctx context.Context, nodeId string) ([]controls.Action } defer tx.Close() - r, err := tx.Run(`MATCH (s) -[:SCANNED]-> (n:Node{node_id:$id}) WHERE NOT (s.status = '`+utils.SCAN_STATUS_SUCCESS+`') AND s.retries < 3 SET s.retries = s.retries + 1 WITH s RETURN s.trigger_action`, map[string]interface{}{"id": nodeId}) + r, err := tx.Run(`MATCH (s) -[:SCANNED]-> (n:Node{node_id:$id}) WHERE s.status = '`+utils.SCAN_STATUS_STARTING+`' AND s.retries < 3 SET s.status = '`+utils.SCAN_STATUS_INPROGRESS+`' WITH s RETURN s.trigger_action`, map[string]interface{}{"id": nodeId}) if err != nil { return res, err diff --git a/deepfence_server/handler/scan_reports.go b/deepfence_server/handler/scan_reports.go index 4ba6f04ded..409ae6e344 100644 --- a/deepfence_server/handler/scan_reports.go +++ b/deepfence_server/handler/scan_reports.go @@ -7,12 +7,12 @@ import ( "net/http" "time" - "github.com/deepfence/ThreatMapper/deepfence_server/controls" "github.com/deepfence/ThreatMapper/deepfence_server/ingesters" "github.com/deepfence/ThreatMapper/deepfence_server/model" "github.com/deepfence/ThreatMapper/deepfence_server/reporters" ctl "github.com/deepfence/ThreatMapper/deepfence_utils/controls" "github.com/deepfence/ThreatMapper/deepfence_utils/log" + "github.com/deepfence/ThreatMapper/deepfence_utils/utils" httpext "github.com/go-playground/pkg/v5/net/http" "github.com/twmb/franz-go/pkg/kgo" ) @@ -34,7 +34,7 @@ func (h *Handler) StartVulnerabilityScanHandler(w http.ResponseWriter, r *http.R RequestPayload: "", } - startScan(w, r, scanId, req.NodeId, action) + startScan(w, r, utils.NEO4J_VULNERABILTY_SCAN, scanId, req.NodeId, action) } func (h *Handler) StartSecretScanHandler(w http.ResponseWriter, r *http.Request) { @@ -73,13 +73,7 @@ func (h *Handler) StartSecretScanHandler(w http.ResponseWriter, r *http.Request) return } - err = ingesters.AddNewScan(r.Context(), "SecretScan", scanId, req.NodeId, action) - if err != nil { - httpext.JSON(w, http.StatusInternalServerError, model.Response{Success: false, Data: err}) - return - } - - startScan(w, r, scanId, req.NodeId, action) + startScan(w, r, utils.NEO4J_SECRET_SCAN, scanId, req.NodeId, action) } func (h *Handler) StartComplianceScanHandler(w http.ResponseWriter, r *http.Request) { @@ -95,7 +89,7 @@ func (h *Handler) StartComplianceScanHandler(w http.ResponseWriter, r *http.Requ RequestPayload: "", } - startScan(w, r, scanId, req.NodeId, action) + startScan(w, r, utils.NEO4J_COMPLIANCE_SCAN, scanId, req.NodeId, action) } func (h *Handler) StartMalwareScanHandler(w http.ResponseWriter, r *http.Request) { @@ -111,7 +105,7 @@ func (h *Handler) StartMalwareScanHandler(w http.ResponseWriter, r *http.Request RequestPayload: "", } - startScan(w, r, scanId, req.NodeId, action) + startScan(w, r, utils.NEO4J_MALWARE_SCAN, scanId, req.NodeId, action) } func (h *Handler) StopVulnerabilityScanHandler(w http.ResponseWriter, r *http.Request) { @@ -132,13 +126,15 @@ func (h *Handler) StopMalwareScanHandler(w http.ResponseWriter, r *http.Request) func startScan( w http.ResponseWriter, r *http.Request, + scanType utils.Neo4jScanType, scanId string, nodeId string, action ctl.Action) { - ctx := r.Context() - err := controls.SetAgentActions(ctx, nodeId, []ctl.Action{ - action, - }) + err := ingesters.AddNewScan(r.Context(), scanType, scanId, nodeId, action) + if err != nil { + httpext.JSON(w, http.StatusInternalServerError, model.Response{Success: false, Data: err}) + return + } if err != nil { log.Error().Msgf("%v", err) diff --git a/deepfence_server/ingesters/agent_ingester.go b/deepfence_server/ingesters/agent_ingester.go index f44620ed14..88a0ce7827 100644 --- a/deepfence_server/ingesters/agent_ingester.go +++ b/deepfence_server/ingesters/agent_ingester.go @@ -11,6 +11,7 @@ import ( "github.com/deepfence/ThreatMapper/deepfence_utils/directory" "github.com/deepfence/ThreatMapper/deepfence_utils/log" + "github.com/deepfence/ThreatMapper/deepfence_utils/utils" redis2 "github.com/go-redis/redis/v8" "github.com/neo4j/neo4j-go-driver/v4/neo4j" "github.com/weaveworks/scope/report" @@ -624,12 +625,14 @@ func (nc *neo4jIngester) applyDBConstraints() error { tx.Run("CREATE CONSTRAINT ON (n:Pod) ASSERT n.node_id IS UNIQUE", map[string]interface{}{}) tx.Run("CREATE CONSTRAINT ON (n:Process) ASSERT n.node_id IS UNIQUE", map[string]interface{}{}) tx.Run("CREATE CONSTRAINT ON (n:KCluster) ASSERT n.node_id IS UNIQUE", map[string]interface{}{}) - tx.Run("CREATE CONSTRAINT ON (n:SecretScan) ASSERT n.node_id IS UNIQUE", map[string]interface{}{}) tx.Run("CREATE CONSTRAINT ON (n:Secret) ASSERT n.rule_id IS UNIQUE", map[string]interface{}{}) tx.Run("CREATE CONSTRAINT ON (n:Cve) ASSERT n.node_id IS UNIQUE", map[string]interface{}{}) - tx.Run("CREATE CONSTRAINT ON (n:CveScan) ASSERT n.node_id IS UNIQUE", map[string]interface{}{}) tx.Run("CREATE CONSTRAINT ON (n:SecurityGroup) ASSERT n.node_id IS UNIQUE", map[string]interface{}{}) tx.Run("CREATE CONSTRAINT ON (n:CloudResource) ASSERT n.node_id IS UNIQUE", map[string]interface{}{}) + tx.Run(fmt.Sprintf("CREATE CONSTRAINT ON (n:%s) ASSERT n.node_id IS UNIQUE", utils.NEO4J_SECRET_SCAN), map[string]interface{}{}) + tx.Run(fmt.Sprintf("CREATE CONSTRAINT ON (n:%s) ASSERT n.node_id IS UNIQUE", utils.NEO4J_VULNERABILTY_SCAN), map[string]interface{}{}) + tx.Run(fmt.Sprintf("CREATE CONSTRAINT ON (n:%s) ASSERT n.node_id IS UNIQUE", utils.NEO4J_COMPLIANCE_SCAN), map[string]interface{}{}) + tx.Run(fmt.Sprintf("CREATE CONSTRAINT ON (n:%s) ASSERT n.node_id IS UNIQUE", utils.NEO4J_MALWARE_SCAN), map[string]interface{}{}) err = tx.Commit() if err != nil { diff --git a/deepfence_server/ingesters/scan_status.go b/deepfence_server/ingesters/scan_status.go index 856c0cc712..d7044f7e32 100644 --- a/deepfence_server/ingesters/scan_status.go +++ b/deepfence_server/ingesters/scan_status.go @@ -21,7 +21,7 @@ func (ve *AlreadyRunningScanError) Error() string { return fmt.Sprintf("Scan of type %s already running for %s, id: %s", ve.scan_type, ve.node_id, ve.scan_id) } -func AddNewScan(ctx context.Context, scan_type string, scan_id string, node_id string, action controls.Action) error { +func AddNewScan(ctx context.Context, scan_type utils.Neo4jScanType, scan_id string, node_id string, action controls.Action) error { driver, err := directory.Neo4jClient(ctx) @@ -55,7 +55,7 @@ func AddNewScan(ctx context.Context, scan_type string, scan_id string, node_id s return &AlreadyRunningScanError{ scan_id: rec.Values[0].(string), node_id: node_id, - scan_type: scan_type, + scan_type: string(scan_type), } } diff --git a/deepfence_utils/utils/constants.go b/deepfence_utils/utils/constants.go index db729776da..95dbf15c1c 100644 --- a/deepfence_utils/utils/constants.go +++ b/deepfence_utils/utils/constants.go @@ -23,6 +23,15 @@ const ( SCAN_STATUS_INPROGRESS = "IN_PROGRESS" ) +type Neo4jScanType string + +const ( + NEO4J_SECRET_SCAN Neo4jScanType = "SecretScan" + NEO4J_VULNERABILTY_SCAN Neo4jScanType = "VulnerabilityScan" + NEO4J_MALWARE_SCAN Neo4jScanType = "MalwareScan" + NEO4J_COMPLIANCE_SCAN Neo4jScanType = "ComplianceScan" +) + var Topics = []string{ VULNERABILITY_SCAN, VULNERABILITY_SCAN_STATUS, SECRET_SCAN, SECRET_SCAN_STATUS,