Skip to content

Commit

Permalink
Clean up start scan actions #728
Browse files Browse the repository at this point in the history 
- Remove actions from host node and move them to scan only
- Clean up startScan logic
  • Loading branch information
@noboruma
noboruma committed Dec 21, 2022
1 parent db3c08d commit 9d8b6fb
Show file tree
Hide file tree
Showing 5 changed files with 56 additions and 76 deletions.
86 changes: 29 additions & 57 deletions deepfence_server/controls/agent.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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
Expand Down
26 changes: 11 additions & 15 deletions deepfence_server/handler/scan_reports.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand All @@ -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) {
Expand Down Expand Up @@ -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) {
Expand All @@ -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) {
Expand All @@ -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) {
Expand All @@ -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)
Expand Down
7 changes: 5 additions & 2 deletions deepfence_server/ingesters/agent_ingester.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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 {
Expand Down
4 changes: 2 additions & 2 deletions deepfence_server/ingesters/scan_status.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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),
}
}

Expand Down
9 changes: 9 additions & 0 deletions deepfence_utils/utils/constants.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit 9d8b6fb

Please sign in to comment.