refactor: modularize codebase and add 404 unit tests (#13)

compliance-agent/src/pentest/prompt_builder.rs

@@ -0,0 +1,504 @@
+use compliance_core::models::dast::DastTarget;
+use compliance_core::models::finding::{Finding, FindingStatus, Severity};
+use compliance_core::models::pentest::*;
+use compliance_core::models::sbom::SbomEntry;
+
+use super::orchestrator::PentestOrchestrator;
+
+/// Return strategy guidance text for the given strategy.
+fn strategy_guidance(strategy: &PentestStrategy) -> &'static str {
+    match strategy {
+        PentestStrategy::Quick => {
+            "Focus on the most common and impactful vulnerabilities. Run a quick recon, then target the highest-risk areas."
+        }
+        PentestStrategy::Comprehensive => {
+            "Perform a thorough assessment covering all vulnerability types. Start with recon, then systematically test each attack surface."
+        }
+        PentestStrategy::Targeted => {
+            "Focus specifically on areas highlighted by SAST findings and known CVEs. Prioritize exploiting known weaknesses."
+        }
+        PentestStrategy::Aggressive => {
+            "Use all available tools aggressively. Test with maximum payloads and attempt full exploitation."
+        }
+        PentestStrategy::Stealth => {
+            "Minimize noise. Use fewer requests, avoid aggressive payloads. Focus on passive analysis and targeted probes."
+        }
+    }
+}
+
+/// Build the SAST findings section for the system prompt.
+fn build_sast_section(sast_findings: &[Finding]) -> String {
+    if sast_findings.is_empty() {
+        return String::from("No SAST findings available for this target.");
+    }
+
+    let critical = sast_findings
+        .iter()
+        .filter(|f| f.severity == Severity::Critical)
+        .count();
+    let high = sast_findings
+        .iter()
+        .filter(|f| f.severity == Severity::High)
+        .count();
+
+    let mut section = format!(
+        "{} open findings ({} critical, {} high):\n",
+        sast_findings.len(),
+        critical,
+        high
+    );
+
+    // List the most important findings (critical/high first, up to 20)
+    for f in sast_findings.iter().take(20) {
+        let file_info = f
+            .file_path
+            .as_ref()
+            .map(|p| format!(" in {}:{}", p, f.line_number.unwrap_or(0)))
+            .unwrap_or_default();
+        let status_note = match f.status {
+            FindingStatus::Triaged => " [TRIAGED]",
+            _ => "",
+        };
+        section.push_str(&format!(
+            "- [{sev}] {title}{file}{status}\n",
+            sev = f.severity,
+            title = f.title,
+            file = file_info,
+            status = status_note,
+        ));
+        if let Some(cwe) = &f.cwe {
+            section.push_str(&format!("  CWE: {cwe}\n"));
+        }
+    }
+    if sast_findings.len() > 20 {
+        section.push_str(&format!(
+            "... and {} more findings\n",
+            sast_findings.len() - 20
+        ));
+    }
+    section
+}
+
+/// Build the SBOM/CVE section for the system prompt.
+fn build_sbom_section(sbom_entries: &[SbomEntry]) -> String {
+    if sbom_entries.is_empty() {
+        return String::from("No vulnerable dependencies identified.");
+    }
+
+    let mut section = format!(
+        "{} dependencies with known vulnerabilities:\n",
+        sbom_entries.len()
+    );
+    for entry in sbom_entries.iter().take(15) {
+        let cve_ids: Vec<&str> = entry
+            .known_vulnerabilities
+            .iter()
+            .map(|v| v.id.as_str())
+            .collect();
+        section.push_str(&format!(
+            "- {} {} ({}): {}\n",
+            entry.name,
+            entry.version,
+            entry.package_manager,
+            cve_ids.join(", ")
+        ));
+    }
+    if sbom_entries.len() > 15 {
+        section.push_str(&format!(
+            "... and {} more vulnerable dependencies\n",
+            sbom_entries.len() - 15
+        ));
+    }
+    section
+}
+
+/// Build the code context section for the system prompt.
+fn build_code_section(code_context: &[CodeContextHint]) -> String {
+    if code_context.is_empty() {
+        return String::from("No code knowledge graph available for this target.");
+    }
+
+    let with_vulns = code_context
+        .iter()
+        .filter(|c| !c.known_vulnerabilities.is_empty())
+        .count();
+
+    let mut section = format!(
+        "{} entry points identified ({} with linked SAST findings):\n",
+        code_context.len(),
+        with_vulns
+    );
+
+    for hint in code_context.iter().take(20) {
+        section.push_str(&format!(
+            "- {} ({})\n",
+            hint.endpoint_pattern, hint.file_path
+        ));
+        for vuln in &hint.known_vulnerabilities {
+            section.push_str(&format!("  SAST: {vuln}\n"));
+        }
+    }
+    section
+}
+
+impl PentestOrchestrator {
+    pub(crate) async fn build_system_prompt(
+        &self,
+        session: &PentestSession,
+        target: &DastTarget,
+        sast_findings: &[Finding],
+        sbom_entries: &[SbomEntry],
+        code_context: &[CodeContextHint],
+    ) -> String {
+        let tool_names = self.tool_registry.list_names().join(", ");
+        let guidance = strategy_guidance(&session.strategy);
+        let sast_section = build_sast_section(sast_findings);
+        let sbom_section = build_sbom_section(sbom_entries);
+        let code_section = build_code_section(code_context);
+
+        format!(
+            r#"You are an expert penetration tester conducting an authorized security assessment.
+
+## Target
+- **Name**: {target_name}
+- **URL**: {base_url}
+- **Type**: {target_type}
+- **Rate Limit**: {rate_limit} req/s
+- **Destructive Tests Allowed**: {allow_destructive}
+- **Linked Repository**: {repo_linked}
+
+## Strategy
+{strategy_guidance}
+
+## SAST Findings (Static Analysis)
+{sast_section}
+
+## Vulnerable Dependencies (SBOM)
+{sbom_section}
+
+## Code Entry Points (Knowledge Graph)
+{code_section}
+
+## Available Tools
+{tool_names}
+
+## Instructions
+1. Start by running reconnaissance (recon tool) to fingerprint the target and discover technologies.
+2. Run the OpenAPI parser to discover API endpoints from specs.
+3. Check infrastructure: DNS, DMARC, TLS, security headers, cookies, CSP, CORS.
+4. Based on SAST findings, prioritize testing endpoints where vulnerabilities were found in code.
+5. For each vulnerability type found in SAST, use the corresponding DAST tool to verify exploitability.
+6. If vulnerable dependencies are listed, try to trigger known CVE conditions against the running application.
+7. Test rate limiting on critical endpoints (login, API).
+8. Check for console.log leakage in frontend JavaScript.
+9. Analyze tool results and chain findings — if one vulnerability enables others, explore the chain.
+10. When testing is complete, provide a structured summary with severity and remediation.
+11. Always explain your reasoning before invoking each tool.
+12. When done, say "Testing complete" followed by a final summary.
+
+## Important
+- This is an authorized penetration test. All testing is permitted within the target scope.
+- Respect the rate limit of {rate_limit} requests per second.
+- Only use destructive tests if explicitly allowed ({allow_destructive}).
+- Use SAST findings to guide your testing — they tell you WHERE in the code vulnerabilities exist.
+- Use SBOM data to understand what technologies and versions the target runs.
+"#,
+            target_name = target.name,
+            base_url = target.base_url,
+            target_type = target.target_type,
+            rate_limit = target.rate_limit,
+            allow_destructive = target.allow_destructive,
+            repo_linked = target.repo_id.as_deref().unwrap_or("None"),
+            strategy_guidance = guidance,
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use compliance_core::models::finding::Severity;
+    use compliance_core::models::sbom::VulnRef;
+    use compliance_core::models::scan::ScanType;
+
+    fn make_finding(
+        severity: Severity,
+        title: &str,
+        file_path: Option<&str>,
+        line: Option<u32>,
+        status: FindingStatus,
+        cwe: Option<&str>,
+    ) -> Finding {
+        let mut f = Finding::new(
+            "repo-1".into(),
+            format!("fp-{title}"),
+            "semgrep".into(),
+            ScanType::Sast,
+            title.into(),
+            "desc".into(),
+            severity,
+        );
+        f.file_path = file_path.map(|s| s.to_string());
+        f.line_number = line;
+        f.status = status;
+        f.cwe = cwe.map(|s| s.to_string());
+        f
+    }
+
+    fn make_sbom_entry(name: &str, version: &str, cves: &[&str]) -> SbomEntry {
+        let mut entry = SbomEntry::new("repo-1".into(), name.into(), version.into(), "npm".into());
+        entry.known_vulnerabilities = cves
+            .iter()
+            .map(|id| VulnRef {
+                id: id.to_string(),
+                source: "nvd".into(),
+                severity: None,
+                url: None,
+            })
+            .collect();
+        entry
+    }
+
+    fn make_code_hint(endpoint: &str, file: &str, vulns: Vec<String>) -> CodeContextHint {
+        CodeContextHint {
+            endpoint_pattern: endpoint.into(),
+            handler_function: "handler".into(),
+            file_path: file.into(),
+            code_snippet: String::new(),
+            known_vulnerabilities: vulns,
+        }
+    }
+
+    // ── strategy_guidance ────────────────────────────────────────────
+
+    #[test]
+    fn strategy_guidance_quick() {
+        let g = strategy_guidance(&PentestStrategy::Quick);
+        assert!(g.contains("most common"));
+        assert!(g.contains("quick recon"));
+    }
+
+    #[test]
+    fn strategy_guidance_comprehensive() {
+        let g = strategy_guidance(&PentestStrategy::Comprehensive);
+        assert!(g.contains("thorough assessment"));
+    }
+
+    #[test]
+    fn strategy_guidance_targeted() {
+        let g = strategy_guidance(&PentestStrategy::Targeted);
+        assert!(g.contains("SAST findings"));
+        assert!(g.contains("known CVEs"));
+    }
+
+    #[test]
+    fn strategy_guidance_aggressive() {
+        let g = strategy_guidance(&PentestStrategy::Aggressive);
+        assert!(g.contains("aggressively"));
+        assert!(g.contains("full exploitation"));
+    }
+
+    #[test]
+    fn strategy_guidance_stealth() {
+        let g = strategy_guidance(&PentestStrategy::Stealth);
+        assert!(g.contains("Minimize noise"));
+        assert!(g.contains("passive analysis"));
+    }
+
+    // ── build_sast_section ───────────────────────────────────────────
+
+    #[test]
+    fn sast_section_empty() {
+        let section = build_sast_section(&[]);
+        assert_eq!(section, "No SAST findings available for this target.");
+    }
+
+    #[test]
+    fn sast_section_single_critical() {
+        let findings = vec![make_finding(
+            Severity::Critical,
+            "SQL Injection",
+            Some("src/db.rs"),
+            Some(42),
+            FindingStatus::Open,
+            Some("CWE-89"),
+        )];
+        let section = build_sast_section(&findings);
+        assert!(section.contains("1 open findings (1 critical, 0 high)"));
+        assert!(section.contains("[critical] SQL Injection in src/db.rs:42"));
+        assert!(section.contains("CWE: CWE-89"));
+    }
+
+    #[test]
+    fn sast_section_triaged_finding_shows_marker() {
+        let findings = vec![make_finding(
+            Severity::High,
+            "XSS",
+            None,
+            None,
+            FindingStatus::Triaged,
+            None,
+        )];
+        let section = build_sast_section(&findings);
+        assert!(section.contains("[TRIAGED]"));
+    }
+
+    #[test]
+    fn sast_section_no_file_path_omits_location() {
+        let findings = vec![make_finding(
+            Severity::Medium,
+            "Open Redirect",
+            None,
+            None,
+            FindingStatus::Open,
+            None,
+        )];
+        let section = build_sast_section(&findings);
+        assert!(section.contains("- [medium] Open Redirect\n"));
+        assert!(!section.contains(" in "));
+    }
+
+    #[test]
+    fn sast_section_counts_critical_and_high() {
+        let findings = vec![
+            make_finding(
+                Severity::Critical,
+                "F1",
+                None,
+                None,
+                FindingStatus::Open,
+                None,
+            ),
+            make_finding(
+                Severity::Critical,
+                "F2",
+                None,
+                None,
+                FindingStatus::Open,
+                None,
+            ),
+            make_finding(Severity::High, "F3", None, None, FindingStatus::Open, None),
+            make_finding(
+                Severity::Medium,
+                "F4",
+                None,
+                None,
+                FindingStatus::Open,
+                None,
+            ),
+        ];
+        let section = build_sast_section(&findings);
+        assert!(section.contains("4 open findings (2 critical, 1 high)"));
+    }
+
+    #[test]
+    fn sast_section_truncates_at_20() {
+        let findings: Vec<Finding> = (0..25)
+            .map(|i| {
+                make_finding(
+                    Severity::Low,
+                    &format!("Finding {i}"),
+                    None,
+                    None,
+                    FindingStatus::Open,
+                    None,
+                )
+            })
+            .collect();
+        let section = build_sast_section(&findings);
+        assert!(section.contains("... and 5 more findings"));
+        // Should contain Finding 19 (the 20th) but not Finding 20 (the 21st)
+        assert!(section.contains("Finding 19"));
+        assert!(!section.contains("Finding 20"));
+    }
+
+    // ── build_sbom_section ───────────────────────────────────────────
+
+    #[test]
+    fn sbom_section_empty() {
+        let section = build_sbom_section(&[]);
+        assert_eq!(section, "No vulnerable dependencies identified.");
+    }
+
+    #[test]
+    fn sbom_section_single_entry() {
+        let entries = vec![make_sbom_entry("lodash", "4.17.20", &["CVE-2021-23337"])];
+        let section = build_sbom_section(&entries);
+        assert!(section.contains("1 dependencies with known vulnerabilities"));
+        assert!(section.contains("- lodash 4.17.20 (npm): CVE-2021-23337"));
+    }
+
+    #[test]
+    fn sbom_section_multiple_cves() {
+        let entries = vec![make_sbom_entry(
+            "openssl",
+            "1.1.1",
+            &["CVE-2022-0001", "CVE-2022-0002"],
+        )];
+        let section = build_sbom_section(&entries);
+        assert!(section.contains("CVE-2022-0001, CVE-2022-0002"));
+    }
+
+    #[test]
+    fn sbom_section_truncates_at_15() {
+        let entries: Vec<SbomEntry> = (0..18)
+            .map(|i| make_sbom_entry(&format!("pkg-{i}"), "1.0.0", &["CVE-2024-0001"]))
+            .collect();
+        let section = build_sbom_section(&entries);
+        assert!(section.contains("... and 3 more vulnerable dependencies"));
+        assert!(section.contains("pkg-14"));
+        assert!(!section.contains("pkg-15"));
+    }
+
+    // ── build_code_section ───────────────────────────────────────────
+
+    #[test]
+    fn code_section_empty() {
+        let section = build_code_section(&[]);
+        assert_eq!(
+            section,
+            "No code knowledge graph available for this target."
+        );
+    }
+
+    #[test]
+    fn code_section_single_entry_no_vulns() {
+        let hints = vec![make_code_hint("GET /api/users", "src/routes.rs", vec![])];
+        let section = build_code_section(&hints);
+        assert!(section.contains("1 entry points identified (0 with linked SAST findings)"));
+        assert!(section.contains("- GET /api/users (src/routes.rs)"));
+    }
+
+    #[test]
+    fn code_section_with_linked_vulns() {
+        let hints = vec![make_code_hint(
+            "POST /login",
+            "src/auth.rs",
+            vec!["[critical] semgrep: SQL Injection (line 15)".into()],
+        )];
+        let section = build_code_section(&hints);
+        assert!(section.contains("1 entry points identified (1 with linked SAST findings)"));
+        assert!(section.contains("SAST: [critical] semgrep: SQL Injection (line 15)"));
+    }
+
+    #[test]
+    fn code_section_counts_entries_with_vulns() {
+        let hints = vec![
+            make_code_hint("GET /a", "a.rs", vec!["vuln1".into()]),
+            make_code_hint("GET /b", "b.rs", vec![]),
+            make_code_hint("GET /c", "c.rs", vec!["vuln2".into(), "vuln3".into()]),
+        ];
+        let section = build_code_section(&hints);
+        assert!(section.contains("3 entry points identified (2 with linked SAST findings)"));
+    }
+
+    #[test]
+    fn code_section_truncates_at_20() {
+        let hints: Vec<CodeContextHint> = (0..25)
+            .map(|i| make_code_hint(&format!("GET /ep{i}"), &format!("f{i}.rs"), vec![]))
+            .collect();
+        let section = build_code_section(&hints);
+        assert!(section.contains("GET /ep19"));
+        assert!(!section.contains("GET /ep20"));
+    }
+}