Skip to main content

codemem_engine/enrichment/
security.rs

1//! Security analysis: sensitive files, endpoints, security functions.
2
3use super::EnrichResult;
4use crate::CodememEngine;
5use codemem_core::{CodememError, GraphBackend, NodeKind};
6use serde_json::json;
7
8impl CodememEngine {
9    /// Enrich the graph with security analysis: sensitive files, endpoints, security functions.
10    pub fn enrich_security(&self, namespace: Option<&str>) -> Result<EnrichResult, CodememError> {
11        use std::sync::LazyLock;
12
13        static SECURITY_PATTERN: LazyLock<regex::Regex> = LazyLock::new(|| {
14            regex::Regex::new(
15                r"(?i)(auth|secret|key|password|token|credential|\.env|private|encrypt|decrypt|cert|permission|rbac|oauth|jwt|session|cookie|csrf|xss|injection|sanitize)"
16            ).expect("valid regex")
17        });
18
19        static SECURITY_FN_PATTERN: LazyLock<regex::Regex> = LazyLock::new(|| {
20            regex::Regex::new(
21                r"(?i)(hash|verify|sign|encrypt|authenticate|authorize|validate_token|check_permission)"
22            ).expect("valid regex")
23        });
24
25        let security_pattern = &*SECURITY_PATTERN;
26        let security_fn_pattern = &*SECURITY_FN_PATTERN;
27
28        let graph = self.lock_graph()?;
29        let all_nodes = graph.get_all_nodes();
30        drop(graph);
31
32        let mut sensitive_files: Vec<String> = Vec::new();
33        let mut endpoints: Vec<String> = Vec::new();
34        let mut security_functions: Vec<(String, String, String)> = Vec::new();
35        let mut nodes_to_annotate: Vec<(String, Vec<String>)> = Vec::new();
36
37        for node in &all_nodes {
38            match node.kind {
39                NodeKind::File => {
40                    if security_pattern.is_match(&node.label) {
41                        sensitive_files.push(node.label.clone());
42                        nodes_to_annotate.push((
43                            node.id.clone(),
44                            vec!["sensitive".into(), "auth_related".into()],
45                        ));
46                    }
47                }
48                NodeKind::Endpoint => {
49                    endpoints.push(node.label.clone());
50                    nodes_to_annotate.push((node.id.clone(), vec!["exposed_endpoint".into()]));
51                }
52                NodeKind::Function | NodeKind::Method => {
53                    if security_fn_pattern.is_match(&node.label) {
54                        let file = node
55                            .payload
56                            .get("file_path")
57                            .and_then(|v| v.as_str())
58                            .unwrap_or("unknown")
59                            .to_string();
60                        security_functions.push((node.label.clone(), file, node.id.clone()));
61                        nodes_to_annotate.push((node.id.clone(), vec!["security_function".into()]));
62                    }
63                }
64                _ => {}
65            }
66        }
67
68        // Annotate nodes with security flags
69        {
70            let mut graph = self.lock_graph()?;
71            for (node_id, flags) in &nodes_to_annotate {
72                if let Ok(Some(mut node)) = graph.get_node(node_id) {
73                    node.payload.insert("security_flags".into(), json!(flags));
74                    let _ = graph.add_node(node);
75                }
76            }
77        }
78
79        // Store insights
80        let mut insights_stored = 0;
81
82        for file_path in &sensitive_files {
83            let content = format!(
84                "Sensitive file: {} — contains security-critical code (auth/credentials)",
85                file_path
86            );
87            if self
88                .store_insight(
89                    &content,
90                    "security",
91                    &["severity:high"],
92                    0.8,
93                    namespace,
94                    &[format!("file:{file_path}")],
95                )
96                .is_some()
97            {
98                insights_stored += 1;
99            }
100        }
101
102        if !endpoints.is_empty() {
103            let content = format!(
104                "{} exposed API endpoints detected — review access controls",
105                endpoints.len()
106            );
107            if self
108                .store_insight(
109                    &content,
110                    "security",
111                    &["severity:medium", "endpoints"],
112                    0.7,
113                    namespace,
114                    &[],
115                )
116                .is_some()
117            {
118                insights_stored += 1;
119            }
120        }
121
122        for (name, file, node_id) in &security_functions {
123            let content = format!(
124                "Security-critical function: {} in {} — ensure proper testing",
125                name, file
126            );
127            if self
128                .store_insight(
129                    &content,
130                    "security",
131                    &["severity:medium"],
132                    0.6,
133                    namespace,
134                    std::slice::from_ref(node_id),
135                )
136                .is_some()
137            {
138                insights_stored += 1;
139            }
140        }
141
142        self.save_index();
143
144        Ok(EnrichResult {
145            insights_stored,
146            details: json!({
147                "sensitive_file_count": sensitive_files.len(),
148                "endpoint_count": endpoints.len(),
149                "security_function_count": security_functions.len(),
150                "insights_stored": insights_stored,
151            }),
152        })
153    }
154}