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

agentshield/rules/
mod.rs
