agentshield/rules/
mod.rs

1pub mod builtin;
2pub mod custom;
3pub mod finding;
4pub mod policy;
5
6use std::collections::HashSet;
7use std::path::{Path, PathBuf};
8
9use crate::analysis::DetectionInput;
10use crate::ir::ScanTarget;
11use crate::ir::SourceLocation;
12
13pub use custom::{CustomRuleDef, CustomRuleDetector, load_custom_rules_from_dir};
14pub use finding::{
15    AttackCategory, Confidence, Evidence, Finding, OwaspMcp, RuleMetadata, Severity,
16};
17
18/// A detector checks a `ScanTarget` and produces findings.
19pub trait Detector: Send + Sync {
20    /// Metadata about this rule (id, name, severity, CWE).
21    fn metadata(&self) -> RuleMetadata;
22
23    /// Run the detector against a scan target.
24    fn run(&self, target: &ScanTarget) -> Vec<Finding>;
25}
26
27pub(crate) trait ContextDetector: Send + Sync {
28    fn metadata(&self) -> RuleMetadata;
29
30    fn run(&self, input: &DetectionInput<'_>) -> Vec<Finding>;
31}
32
33/// The rule engine runs all registered detectors against a target.
34pub struct RuleEngine {
35    detectors: Vec<Box<dyn Detector>>,
36    context_detectors: Vec<Box<dyn ContextDetector>>,
37    custom_detectors: Vec<CustomRuleDetector>,
38}
39
40impl RuleEngine {
41    /// Create a new engine with all built-in detectors registered.
42    pub fn new() -> Self {
43        Self {
44            detectors: builtin::all_detectors(),
45            context_detectors: builtin::all_context_detectors(),
46            custom_detectors: Vec::new(),
47        }
48    }
49
50    /// Add custom rule detectors to the engine.
51    pub fn with_custom_rules(mut self, custom: Vec<CustomRuleDetector>) -> Self {
52        self.custom_detectors.extend(custom);
53        self
54    }
55
56    /// Load custom rules from a directory into this engine.
57    pub fn load_custom_rules_from(&mut self, dir: &Path) -> crate::error::Result<()> {
58        let loaded = custom::load_custom_rules_from_dir(dir)?;
59        self.custom_detectors.extend(loaded);
60        Ok(())
61    }
62
63    /// Run all detectors against a scan target.
64    pub fn run(&self, target: &ScanTarget) -> Vec<Finding> {
65        let mut findings: Vec<Finding> =
66            self.detectors.iter().flat_map(|d| d.run(target)).collect();
67        findings.extend(self.custom_detectors.iter().flat_map(|d| d.run(target)));
68        apply_overlapping_rule_suppression(findings)
69    }
70
71    /// Run all built-in detectors, including contextual and custom detectors.
72    pub(crate) fn run_with_context(&self, input: &DetectionInput<'_>) -> Vec<Finding> {
73        let mut findings = self
74            .detectors
75            .iter()
76            .flat_map(|detector| detector.run(input.target))
77            .collect::<Vec<_>>();
78        findings.extend(
79            self.context_detectors
80                .iter()
81                .flat_map(|detector| detector.run(input)),
82        );
83        findings.extend(
84            self.custom_detectors
85                .iter()
86                .flat_map(|detector| detector.run(input.target)),
87        );
88        apply_overlapping_rule_suppression(findings)
89    }
90
91    /// List metadata for all registered rules, including custom rules.
92    pub fn list_rules(&self) -> Vec<RuleMetadata> {
93        let mut rules: Vec<RuleMetadata> = self.detectors.iter().map(|d| d.metadata()).collect();
94        rules.extend(self.custom_detectors.iter().map(|d| d.metadata()));
95        rules
96    }
97
98    /// List metadata for all scanner rules, including future contextual detectors.
99    pub fn list_scanner_rules(&self) -> Vec<RuleMetadata> {
100        let mut rules = self
101            .detectors
102            .iter()
103            .map(|d| d.metadata())
104            .collect::<Vec<_>>();
105        rules.extend(self.context_detectors.iter().map(|d| d.metadata()));
106        rules.extend(self.custom_detectors.iter().map(|d| d.metadata()));
107        rules
108    }
109}
110
111impl Default for RuleEngine {
112    fn default() -> Self {
113        Self::new()
114    }
115}
116
117const OVERLAPPING_RULE_PAIRS: &[(&str, &str)] = &[
118    // Keep precise/specific signal and suppress broader overlap.
119    ("SHIELD-013", "SHIELD-003"), // Metadata/private SSRF suppresses generic SSRF
120    ("SHIELD-002", "SHIELD-018"), // Credential exfil suppresses generic secret leakage
121    ("SHIELD-004", "SHIELD-015"), // Arbitrary file access suppresses overbroad filesystem scope
122    ("SHIELD-011", "SHIELD-016"), // Dynamic eval/import suppression suppresses unsafe deserialization overlap
123];
124
125#[derive(Debug, Clone, PartialEq, Eq, Hash)]
126struct SourceLocationKey {
127    file: PathBuf,
128    line: usize,
129    column: usize,
130    end_line: Option<usize>,
131    end_column: Option<usize>,
132}
133
134impl From<&SourceLocation> for SourceLocationKey {
135    fn from(location: &SourceLocation) -> Self {
136        Self {
137            file: location.file.clone(),
138            line: location.line,
139            column: location.column,
140            end_line: location.end_line,
141            end_column: location.end_column,
142        }
143    }
144}
145
146#[derive(Debug, Clone, PartialEq, Eq, Hash)]
147struct SuppressionKey {
148    dominant_rule: &'static str,
149    location: SourceLocationKey,
150}
151
152fn apply_overlapping_rule_suppression(findings: Vec<Finding>) -> Vec<Finding> {
153    // Only dominant rules can suppress a candidate. Indexing those keys avoids
154    // comparing every finding with every other finding on the scan path.
155    let suppressors: HashSet<SuppressionKey> = findings
156        .iter()
157        .filter_map(|dominant| {
158            let dominant_rule = dominant_rule_for_dominant(dominant.rule_id.as_str())?;
159            let location = dominant.location.as_ref()?;
160            dominant_finding_can_suppress(dominant).then_some(SuppressionKey {
161                dominant_rule,
162                location: location.into(),
163            })
164        })
165        .collect();
166
167    findings
168        .into_iter()
169        .filter(|candidate| {
170            let Some(dominant_rule) = dominant_rule_for_candidate(candidate.rule_id.as_str())
171            else {
172                return true;
173            };
174            let Some(location) = candidate.location.as_ref() else {
175                return true;
176            };
177
178            !suppressors.contains(&SuppressionKey {
179                dominant_rule,
180                location: location.into(),
181            })
182        })
183        .collect()
184}
185
186fn dominant_rule_for_candidate(candidate_rule: &str) -> Option<&'static str> {
187    OVERLAPPING_RULE_PAIRS
188        .iter()
189        .find_map(|(dominant, dominated)| (*dominated == candidate_rule).then_some(*dominant))
190}
191
192fn dominant_rule_for_dominant(dominant_rule: &str) -> Option<&'static str> {
193    OVERLAPPING_RULE_PAIRS
194        .iter()
195        .find_map(|(dominant, _)| (*dominant == dominant_rule).then_some(*dominant))
196}
197
198fn dominant_finding_can_suppress(dominant: &Finding) -> bool {
199    // A tainted URL is enough to retain the critical SHIELD-013 signal, but it
200    // is not proof that the destination is metadata/private. Keep the generic
201    // SHIELD-003 finding in that uncertain case; suppress it only when the
202    // dominant finding carries a concrete metadata/private indication (or is a
203    // synthetic dominant finding with no taint path).
204    if dominant.rule_id == "SHIELD-013"
205        && dominant.taint_path.is_some()
206        && !dominant
207            .evidence
208            .iter()
209            .any(|evidence| evidence.description.starts_with("Sink: HTTP request to "))
210    {
211        return false;
212    }
213
214    true
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220    use crate::ir::ArgumentSource;
221    use crate::ir::data_surface::*;
222    use crate::ir::execution_surface::*;
223    use std::path::PathBuf;
224
225    fn loc() -> SourceLocation {
226        SourceLocation {
227            file: PathBuf::from("server.py"),
228            line: 10,
229            column: 0,
230            end_line: None,
231            end_column: None,
232        }
233    }
234
235    fn empty_target() -> ScanTarget {
236        ScanTarget {
237            name: "test".into(),
238            framework: crate::ir::Framework::Mcp,
239            root_path: PathBuf::from("."),
240            tools: vec![],
241            execution: ExecutionSurface::default(),
242            data: DataSurface::default(),
243            dependencies: Default::default(),
244            provenance: Default::default(),
245            source_files: vec![],
246        }
247    }
248
249    fn simple_finding(rule_id: &str, location: Option<SourceLocation>) -> Finding {
250        Finding {
251            rule_id: rule_id.to_string(),
252            rule_name: rule_id.to_string(),
253            severity: Severity::Critical,
254            confidence: Confidence::High,
255            attack_category: AttackCategory::ArbitraryFileAccess,
256            message: "test".into(),
257            location,
258            evidence: vec![],
259            taint_path: None,
260            remediation: None,
261            cwe_id: None,
262        }
263    }
264
265    #[test]
266    fn all_builtin_rules_have_owasp_mcp_mapping() {
267        let engine = RuleEngine::new();
268        let rules = engine.list_rules();
269        assert!(!rules.is_empty());
270        for rule in &rules {
271            assert!(
272                rule.owasp_mcp.is_some(),
273                "rule {} is missing an OWASP MCP Top 10 mapping",
274                rule.id
275            );
276        }
277    }
278
279    #[test]
280    fn suppresses_overlapping_014_pairs_in_engine_output() {
281        let target = {
282            let mut target = empty_target();
283            let overlap_loc = loc();
284
285            target.data.taint_paths.push(TaintPath {
286                source: TaintSource {
287                    source_type: TaintSourceType::ToolArgument,
288                    description: "url".into(),
289                    location: overlap_loc.clone(),
290                },
291                sink: TaintSink {
292                    sink_type: TaintSinkType::HttpRequest,
293                    description: "requests.get".into(),
294                    location: overlap_loc.clone(),
295                },
296                through: vec![],
297                confidence: 0.9,
298            });
299            target.execution.network_operations.push(NetworkOperation {
300                function: "requests.get".into(),
301                url_arg: ArgumentSource::Literal("http://169.254.169.254/latest/meta-data/".into()),
302                method: Some("GET".into()),
303                sends_data: false,
304                location: overlap_loc.clone(),
305            });
306
307            target.execution.file_operations.push(FileOperation {
308                operation: FileOpType::Read,
309                path_arg: ArgumentSource::Parameter {
310                    name: "path".into(),
311                },
312                location: overlap_loc.clone(),
313            });
314
315            target
316        };
317
318        let findings = RuleEngine::new().run(&target);
319        let has_metadata_ssrf = findings.iter().any(|f| f.rule_id == "SHIELD-013");
320        let has_ssrf = findings.iter().any(|f| f.rule_id == "SHIELD-003");
321        assert!(has_metadata_ssrf, "should keep SHIELD-013");
322        assert!(
323            !has_ssrf,
324            "SHIELD-003 should be suppressed when SHIELD-013 is present at same location"
325        );
326        assert!(has_metadata_ssrf || has_ssrf);
327    }
328
329    #[test]
330    fn suppresses_overlapping_findings_for_arbitrary_file_and_overbroad_fs() {
331        let target = {
332            let mut target = empty_target();
333            target.execution.file_operations.push(FileOperation {
334                operation: FileOpType::Write,
335                path_arg: ArgumentSource::Parameter {
336                    name: "file_path".into(),
337                },
338                location: loc(),
339            });
340            target
341        };
342
343        let findings = RuleEngine::new().run(&target);
344        let has_arf = findings.iter().any(|f| f.rule_id == "SHIELD-004");
345        let has_overbroad = findings.iter().any(|f| f.rule_id == "SHIELD-015");
346        assert!(has_arf);
347        assert!(!has_overbroad);
348    }
349
350    #[test]
351    fn suppresses_overlapping_pairs_by_rule_with_same_location() {
352        let overlap_loc = loc();
353        let findings = vec![
354            simple_finding("SHIELD-016", Some(overlap_loc.clone())),
355            simple_finding("SHIELD-011", Some(overlap_loc.clone())),
356            simple_finding("SHIELD-018", Some(overlap_loc.clone())),
357            simple_finding("SHIELD-002", Some(overlap_loc.clone())),
358            simple_finding("SHIELD-003", Some(overlap_loc.clone())),
359            simple_finding("SHIELD-013", Some(overlap_loc)),
360        ];
361        let filtered = apply_overlapping_rule_suppression(findings);
362        let ids: Vec<_> = filtered.iter().map(|f| f.rule_id.as_str()).collect();
363        assert!(ids.contains(&"SHIELD-013"));
364        assert!(!ids.contains(&"SHIELD-003"));
365        assert!(ids.contains(&"SHIELD-002"));
366        assert!(!ids.contains(&"SHIELD-018"));
367        assert!(ids.contains(&"SHIELD-011"));
368        assert!(!ids.contains(&"SHIELD-016"));
369    }
370
371    #[test]
372    fn keeps_overlapping_rules_for_distinct_spans_on_the_same_line() {
373        let mut first_span = loc();
374        first_span.column = 4;
375        first_span.end_line = Some(first_span.line);
376        first_span.end_column = Some(14);
377
378        let mut second_span = first_span.clone();
379        second_span.column = 16;
380        second_span.end_column = Some(33);
381
382        let findings = vec![
383            simple_finding("SHIELD-004", Some(first_span)),
384            simple_finding("SHIELD-015", Some(second_span)),
385        ];
386
387        let filtered = apply_overlapping_rule_suppression(findings);
388        let ids: Vec<_> = filtered
389            .iter()
390            .map(|finding| finding.rule_id.as_str())
391            .collect();
392
393        assert_eq!(filtered.len(), 2);
394        assert!(ids.contains(&"SHIELD-004"));
395        assert!(ids.contains(&"SHIELD-015"));
396    }
397}
agentshield/rules/mod.rs

agentshield/rules/
mod.rs
