Skip to main content

agentshield/
lib.rs

1//! AgentShield — Security scanner for AI agent extensions.
2//!
3//! Offline-first, multi-framework, SARIF output. Scans MCP servers,
4//! OpenClaw skills, and other agent extension formats for security issues.
5//!
6//! # Quick Start
7//!
8//! ```no_run
9//! use std::path::Path;
10//! use agentshield::{scan, ScanOptions};
11//!
12//! let options = ScanOptions::default();
13//! let report = scan(Path::new("./my-mcp-server"), &options).unwrap();
14//! println!("Pass: {}, Findings: {}", report.verdict.pass, report.findings.len());
15//! ```
16
17pub mod adapter;
18pub mod analysis;
19pub mod baseline;
20pub mod certify;
21pub mod config;
22pub mod doctor;
23pub mod egress;
24pub mod error;
25pub mod ir;
26pub mod output;
27pub mod parser;
28pub mod rules;
29
30use std::path::Path;
31
32use config::Config;
33use error::Result;
34use ir::ScanTarget;
35use output::OutputFormat;
36use rules::policy::PolicyVerdict;
37use rules::{Finding, RuleEngine};
38
39/// Options for a scan invocation.
40#[derive(Debug, Clone)]
41pub struct ScanOptions {
42    /// Path to config file (defaults to `.agentshield.toml` in scan dir).
43    pub config_path: Option<std::path::PathBuf>,
44    /// Output format.
45    pub format: OutputFormat,
46    /// CLI override for fail_on threshold.
47    pub fail_on_override: Option<rules::Severity>,
48    /// Skip test files (test/, tests/, *.test.ts, *.spec.ts, etc.).
49    pub ignore_tests: bool,
50}
51
52impl Default for ScanOptions {
53    fn default() -> Self {
54        Self {
55            config_path: None,
56            format: OutputFormat::Console,
57            fail_on_override: None,
58            ignore_tests: false,
59        }
60    }
61}
62
63/// Complete scan report.
64#[derive(Debug)]
65pub struct ScanReport {
66    pub target_name: String,
67    pub findings: Vec<Finding>,
68    pub verdict: PolicyVerdict,
69    /// Absolute (or canonicalized) path to the scanned directory.
70    /// Passed to output renderers for stable fingerprint computation.
71    pub scan_root: std::path::PathBuf,
72    /// Raw scan targets produced by the adapter pipeline.
73    /// Used by callers that need to inspect the IR (e.g., `--emit-egress-policy`).
74    pub targets: Vec<ScanTarget>,
75}
76
77/// Run a complete scan: detect framework, parse, analyze, evaluate policy.
78pub fn scan(path: &Path, options: &ScanOptions) -> Result<ScanReport> {
79    // Load config
80    let config_path = options
81        .config_path
82        .clone()
83        .unwrap_or_else(|| path.join(".agentshield.toml"));
84    let mut config = Config::load(&config_path)?;
85
86    // Apply CLI override
87    if let Some(fail_on) = options.fail_on_override {
88        config.policy.fail_on = fail_on;
89    }
90
91    // Auto-detect framework and load IR
92    let ignore_tests = options.ignore_tests || config.scan.ignore_tests;
93    let targets = adapter::auto_detect_and_load(path, ignore_tests)?;
94
95    // Run detectors on all targets
96    let engine = RuleEngine::new();
97    let mut all_findings: Vec<Finding> = Vec::new();
98
99    let target_name = if let Some(first) = targets.first() {
100        first.name.clone()
101    } else {
102        path.file_name()
103            .map(|n| n.to_string_lossy().into_owned())
104            .unwrap_or_else(|| "unknown".into())
105    };
106
107    for target in &targets {
108        let findings = engine.run(target);
109        all_findings.extend(findings);
110    }
111
112    // Canonicalize for stable fingerprints; fall back to the raw path on error.
113    let scan_root = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
114
115    // Apply policy (ignore rules, overrides, suppressions)
116    let effective_findings = config.policy.apply(&all_findings, &scan_root);
117    let verdict = config.policy.evaluate(&effective_findings);
118
119    Ok(ScanReport {
120        target_name,
121        findings: effective_findings,
122        verdict,
123        scan_root,
124        targets,
125    })
126}
127
128/// Render a scan report in the specified format.
129pub fn render_report(report: &ScanReport, format: OutputFormat) -> Result<String> {
130    output::render(
131        &report.findings,
132        &report.verdict,
133        format,
134        &report.target_name,
135        &report.scan_root,
136    )
137}
138
139#[cfg(test)]
140mod integration_tests {
141    use super::*;
142    use std::path::Path;
143
144    #[test]
145    fn safe_calculator_zero_findings() {
146        let opts = ScanOptions::default();
147        let report = scan(
148            Path::new("tests/fixtures/mcp_servers/safe_calculator"),
149            &opts,
150        )
151        .unwrap();
152        // No code-level security findings (SHIELD-001 through SHIELD-006, SHIELD-011)
153        assert!(
154            !report
155                .findings
156                .iter()
157                .any(|f| f.severity >= rules::Severity::High),
158            "safe calculator should have no High+ findings"
159        );
160        assert!(report.verdict.pass);
161    }
162
163    #[test]
164    fn vuln_cmd_inject_detected() {
165        let opts = ScanOptions::default();
166        let report = scan(
167            Path::new("tests/fixtures/mcp_servers/vuln_cmd_inject"),
168            &opts,
169        )
170        .unwrap();
171        assert!(report.findings.iter().any(|f| f.rule_id == "SHIELD-001"));
172        assert!(!report.verdict.pass);
173    }
174
175    #[test]
176    fn vuln_ssrf_detected() {
177        let opts = ScanOptions::default();
178        let report = scan(Path::new("tests/fixtures/mcp_servers/vuln_ssrf"), &opts).unwrap();
179        assert!(report.findings.iter().any(|f| f.rule_id == "SHIELD-003"));
180        assert!(!report.verdict.pass);
181    }
182
183    #[test]
184    fn vuln_cred_exfil_detected() {
185        let opts = ScanOptions::default();
186        let report = scan(
187            Path::new("tests/fixtures/mcp_servers/vuln_cred_exfil"),
188            &opts,
189        )
190        .unwrap();
191        assert!(report.findings.iter().any(|f| f.rule_id == "SHIELD-002"));
192        assert!(!report.verdict.pass);
193    }
194
195    #[test]
196    fn baseline_write_and_filter_round_trip() {
197        use crate::baseline::{BaselineEntry, BaselineFile};
198        use tempfile::NamedTempFile;
199
200        let fixture = Path::new("tests/fixtures/mcp_servers/vuln_cmd_inject");
201        let opts = ScanOptions::default();
202
203        // Step 1: scan and get findings
204        let report = scan(fixture, &opts).unwrap();
205        assert!(
206            !report.findings.is_empty(),
207            "vuln_cmd_inject should produce findings"
208        );
209
210        // Step 2: write baseline from findings
211        let baseline_file = NamedTempFile::new().unwrap();
212        let now = chrono::Utc::now().to_rfc3339();
213        let entries: Vec<BaselineEntry> = report
214            .findings
215            .iter()
216            .map(|f| BaselineEntry {
217                fingerprint: f.fingerprint(&report.scan_root),
218                rule_id: f.rule_id.clone(),
219                first_seen: now.clone(),
220            })
221            .collect();
222        let baseline = BaselineFile::new(entries);
223        baseline.save(baseline_file.path()).unwrap();
224
225        // Step 3: re-scan and filter with baseline
226        let report2 = scan(fixture, &opts).unwrap();
227        let loaded_baseline = BaselineFile::load(baseline_file.path()).unwrap();
228        let filtered: Vec<_> = report2
229            .findings
230            .into_iter()
231            .filter(|f| {
232                let fp = f.fingerprint(&report2.scan_root);
233                !loaded_baseline.contains(&fp)
234            })
235            .collect();
236
237        // Step 4: all findings should be filtered out
238        assert!(
239            filtered.is_empty(),
240            "All findings should be filtered by baseline, but {} remain: {:?}",
241            filtered.len(),
242            filtered.iter().map(|f| &f.rule_id).collect::<Vec<_>>()
243        );
244    }
245
246    #[test]
247    fn suppress_command_roundtrip() {
248        use crate::config::Config;
249        use crate::rules::policy::Suppression;
250        use tempfile::TempDir;
251
252        // Use a temp dir so we don't pollute fixture directories
253        let tmp = TempDir::new().unwrap();
254        let fixture = Path::new("tests/fixtures/mcp_servers/vuln_cmd_inject");
255
256        // Step 1: Scan the fixture to get a real fingerprint
257        let opts = ScanOptions::default();
258        let report = scan(fixture, &opts).unwrap();
259        assert!(
260            !report.findings.is_empty(),
261            "vuln_cmd_inject should produce findings"
262        );
263
264        let first_finding = &report.findings[0];
265        let fp = first_finding.fingerprint(&report.scan_root);
266        let rule_id = first_finding.rule_id.clone();
267
268        // Step 2: Write a config with the suppression into a temp dir
269        let config_path = tmp.path().join(".agentshield.toml");
270        let mut cfg = Config::default();
271        cfg.policy.suppressions.push(Suppression {
272            fingerprint: fp.clone(),
273            reason: "Integration test suppression".into(),
274            expires: None,
275            created_at: Some("2026-03-21".into()),
276        });
277        let toml_str = toml::to_string_pretty(&cfg).unwrap();
278        std::fs::write(&config_path, &toml_str).unwrap();
279
280        // Step 3: Verify the config round-trips correctly
281        let loaded = Config::load(&config_path).unwrap();
282        assert_eq!(loaded.policy.suppressions.len(), 1);
283        assert_eq!(loaded.policy.suppressions[0].fingerprint, fp);
284        assert_eq!(
285            loaded.policy.suppressions[0].reason,
286            "Integration test suppression"
287        );
288
289        // Step 4: Re-scan using the config with the suppression
290        let opts_with_config = ScanOptions {
291            config_path: Some(config_path.clone()),
292            ..ScanOptions::default()
293        };
294        let report2 = scan(fixture, &opts_with_config).unwrap();
295
296        // The suppressed finding should no longer appear
297        let still_present = report2
298            .findings
299            .iter()
300            .any(|f| f.rule_id == rule_id && f.fingerprint(&report2.scan_root) == fp);
301
302        assert!(
303            !still_present,
304            "Suppressed finding {} should not appear in re-scan",
305            fp
306        );
307    }
308
309    /// Verify that dependency findings with manifest file locations appear
310    /// consistently across all 4 output formats (console, JSON, SARIF, HTML).
311    ///
312    /// Prior to T7, SHIELD-009 and SHIELD-012 had `location: None` and were
313    /// silently dropped from SARIF output while appearing as "-" in console
314    /// and HTML. Now that the adapter populates manifest file locations, all
315    /// formats must agree on the same location.
316    #[test]
317    fn dep_findings_location_parity_across_output_formats() {
318        use crate::output::OutputFormat;
319
320        let fixture = Path::new("tests/fixtures/mcp_servers/vuln_unpinned_deps");
321        let opts = ScanOptions::default();
322        let report = scan(fixture, &opts).unwrap();
323
324        // The fixture has requirements.txt with >=versions and no lockfile.
325        // Expect at least one SHIELD-009 (Unpinned Dependencies) finding.
326        let dep_finding = report
327            .findings
328            .iter()
329            .find(|f| f.rule_id == "SHIELD-009")
330            .expect("Expected at least one SHIELD-009 finding from vuln_unpinned_deps fixture");
331
332        // The finding must have a location pointing to requirements.txt.
333        let loc = dep_finding
334            .location
335            .as_ref()
336            .expect("SHIELD-009 finding must carry a manifest file location");
337        assert!(
338            loc.file.to_string_lossy().contains("requirements.txt"),
339            "SHIELD-009 location file should be requirements.txt, got: {}",
340            loc.file.display()
341        );
342        assert!(loc.line >= 1, "SHIELD-009 location line must be >= 1");
343
344        let expected_file = loc.file.to_string_lossy().to_string();
345
346        // Render all 4 formats and verify the location is present in each.
347        let console_out =
348            render_report(&report, OutputFormat::Console).expect("console render failed");
349        assert!(
350            console_out.contains("requirements.txt"),
351            "Console output should contain requirements.txt for dep findings"
352        );
353
354        let json_out = render_report(&report, OutputFormat::Json).expect("json render failed");
355        let json_val: serde_json::Value =
356            serde_json::from_str(&json_out).expect("JSON output must be valid JSON");
357        let json_findings = json_val["findings"]
358            .as_array()
359            .expect("JSON must have findings array");
360        let json_dep = json_findings
361            .iter()
362            .find(|f| f["rule_id"].as_str() == Some("SHIELD-009"))
363            .expect("JSON output must contain SHIELD-009 finding");
364        let json_file = json_dep["location"]["file"]
365            .as_str()
366            .expect("JSON SHIELD-009 finding must have location.file");
367        assert!(
368            json_file.contains("requirements.txt"),
369            "JSON location.file should contain requirements.txt, got: {json_file}"
370        );
371
372        let sarif_out = render_report(&report, OutputFormat::Sarif).expect("SARIF render failed");
373        let sarif_val: serde_json::Value =
374            serde_json::from_str(&sarif_out).expect("SARIF output must be valid JSON");
375        let sarif_results = sarif_val["runs"][0]["results"]
376            .as_array()
377            .expect("SARIF must have runs[0].results array");
378        let sarif_dep = sarif_results
379            .iter()
380            .find(|r| r["ruleId"].as_str() == Some("SHIELD-009"))
381            .expect(
382                "SARIF output must contain SHIELD-009 result (dep findings now have locations)",
383            );
384        let sarif_uri = sarif_dep["locations"][0]["physicalLocation"]["artifactLocation"]["uri"]
385            .as_str()
386            .expect("SARIF SHIELD-009 result must have a physicalLocation URI");
387        assert!(
388            sarif_uri.contains("requirements.txt"),
389            "SARIF artifactLocation URI should contain requirements.txt, got: {sarif_uri}"
390        );
391        // Verify the URI matches the JSON location (both point to the same file)
392        assert!(
393            expected_file.contains("requirements.txt"),
394            "Location file {expected_file} must reference requirements.txt"
395        );
396
397        let html_out = render_report(&report, OutputFormat::Html).expect("HTML render failed");
398        assert!(
399            html_out.contains("requirements.txt"),
400            "HTML output should contain requirements.txt for dep findings"
401        );
402        // HTML must NOT show "-" for the dep finding location
403        // (the table shows the location as "file:line")
404        assert!(
405            !html_out.contains("<code>-</code>"),
406            "HTML output must not show '-' for dep finding locations that have a manifest file"
407        );
408    }
409
410    #[test]
411    fn vuln_metadata_ssrf_detected() {
412        let opts = ScanOptions::default();
413        let report = scan(
414            Path::new("tests/fixtures/mcp_servers/vuln_metadata_ssrf"),
415            &opts,
416        )
417        .unwrap();
418        // The fixture has a tool that passes user-controlled `url` to requests.get,
419        // which should trigger SHIELD-003 (general SSRF) and potentially SHIELD-013
420        // (metadata SSRF) via taint paths if populated.
421        // At minimum, SHIELD-003 must fire (parameter -> network call).
422        assert!(
423            report.findings.iter().any(|f| f.rule_id == "SHIELD-003"),
424            "Expected SHIELD-003 (general SSRF) from vuln_metadata_ssrf fixture"
425        );
426        assert!(!report.verdict.pass);
427    }
428
429    #[test]
430    fn safe_filesystem_no_file_access_findings() {
431        // This fixture has a handler that validates paths via validatePath()
432        // before passing them to helper functions. Cross-file analysis should
433        // downgrade the helpers' operations from tainted to sanitized.
434        let opts = ScanOptions::default();
435        let report = scan(
436            Path::new("tests/fixtures/mcp_servers/safe_filesystem"),
437            &opts,
438        )
439        .unwrap();
440
441        let file_access_findings: Vec<_> = report
442            .findings
443            .iter()
444            .filter(|f| f.rule_id == "SHIELD-004")
445            .collect();
446
447        assert!(
448            file_access_findings.is_empty(),
449            "Expected 0 SHIELD-004 findings (cross-file sanitization should eliminate FPs), \
450             but got {}: {:?}",
451            file_access_findings.len(),
452            file_access_findings
453                .iter()
454                .map(|f| &f.message)
455                .collect::<Vec<_>>()
456        );
457    }
458}