compliance-scanner-agent/compliance-agent/src/pentest/prompt_builder.rs

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;

/// Attempt to decrypt a field; if decryption fails, return the original value
/// (which may be plaintext from before encryption was enabled).
fn decrypt_field(value: &str) -> String {
    super::crypto::decrypt(value).unwrap_or_else(|| value.to_string())
}

/// Build additional prompt sections from PentestConfig when present.
fn build_config_sections(config: &PentestConfig) -> String {
    let mut sections = String::new();

    // Authentication section
    match config.auth.mode {
        AuthMode::Manual => {
            sections.push_str("\n## Authentication\n");
            sections.push_str("- **Mode**: Manual credentials\n");
            if let Some(ref u) = config.auth.username {
                let decrypted = decrypt_field(u);
                sections.push_str(&format!("- **Username**: {decrypted}\n"));
            }
            if let Some(ref p) = config.auth.password {
                let decrypted = decrypt_field(p);
                sections.push_str(&format!("- **Password**: {decrypted}\n"));
            }
            sections.push_str(
                "Use these credentials to log in before testing authenticated endpoints.\n",
            );
        }
        AuthMode::AutoRegister => {
            sections.push_str("\n## Authentication\n");
            sections.push_str("- **Mode**: Auto-register\n");
            if let Some(ref url) = config.auth.registration_url {
                sections.push_str(&format!("- **Registration URL**: {url}\n"));
            } else {
                sections.push_str(
                    "- **Registration URL**: Not provided — use Playwright to discover the registration page.\n",
                );
            }
            if let Some(ref email) = config.auth.verification_email {
                sections.push_str(&format!(
                    "- **Verification Email**: Use plus-addressing from `{email}` \
                     (e.g. `{base}+{{session_id}}@{domain}`) for email verification. \
                     The system will poll the IMAP mailbox for verification links.\n",
                    base = email.split('@').next().unwrap_or(email),
                    domain = email.split('@').nth(1).unwrap_or("example.com"),
                ));
            }
            sections.push_str(
                "Register a new test account using the registration page, then use it for testing.\n",
            );
        }
        AuthMode::None => {}
    }

    // Custom headers
    if !config.custom_headers.is_empty() {
        sections.push_str("\n## Custom HTTP Headers\n");
        sections.push_str("Include these headers in all HTTP requests:\n");
        for (k, v) in &config.custom_headers {
            sections.push_str(&format!("- `{k}: {v}`\n"));
        }
    }

    // Scope exclusions
    if !config.scope_exclusions.is_empty() {
        sections.push_str("\n## Scope Exclusions\n");
        sections.push_str("Do NOT test the following paths:\n");
        for path in &config.scope_exclusions {
            sections.push_str(&format!("- `{path}`\n"));
        }
    }

    // Git context
    if config.git_repo_url.is_some() || config.branch.is_some() || config.commit_hash.is_some() {
        sections.push_str("\n## Git Context\n");
        if let Some(ref url) = config.git_repo_url {
            sections.push_str(&format!("- **Repository**: {url}\n"));
        }
        if let Some(ref branch) = config.branch {
            sections.push_str(&format!("- **Branch**: {branch}\n"));
        }
        if let Some(ref commit) = config.commit_hash {
            sections.push_str(&format!("- **Commit**: {commit}\n"));
        }
    }

    // Environment
    sections.push_str(&format!(
        "\n## Environment\n- **Target environment**: {}\n",
        config.environment
    ));

    sections
}

/// 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);
        let config_sections = session
            .config
            .as_ref()
            .map(build_config_sections)
            .unwrap_or_default();

        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}
{config_sections}
## 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. If the target requires authentication (auto-register mode), use the browser tool to:
   a. Navigate to the target — it will redirect to the login page.
   b. Click the "Register" link to reach the registration form.
   c. Fill all form fields (username, email with plus-addressing, password, name) one by one.
   d. Click submit. If a Terms & Conditions page appears, accept it.
   e. After registration, use the browser to navigate through the application pages.
   f. **Take a screenshot after each major page** for evidence in the report.
5. Use the browser tool to explore the authenticated application — navigate to each section,
   use get_content to understand the page structure, and take screenshots.
6. Based on SAST findings, prioritize testing endpoints where vulnerabilities were found in code.
7. For each vulnerability type found in SAST, use the corresponding DAST tool to verify exploitability.
8. If vulnerable dependencies are listed, try to trigger known CVE conditions against the running application.
9. Test rate limiting on critical endpoints (login, API).
10. Check for console.log leakage in frontend JavaScript.
11. Analyze tool results and chain findings — if one vulnerability enables others, explore the chain.
12. When testing is complete, provide a structured summary with severity and remediation.
13. Always explain your reasoning before invoking each tool.
14. When done, say "Testing complete" followed by a final summary.

## Browser Tool Usage
- The browser tab **persists** between calls — cookies and login state are preserved.
- After navigate, the response includes `elements` (links, inputs, buttons on the page).
- Use `get_content` to see forms, links, buttons, headings, and page text.
- Use `click` with CSS selectors to interact (e.g., `a:text('Register')`, `input[type='submit']`).
- Use `fill` with selector + value to fill form fields (e.g., `input[name='email']`).
- **Take screenshots** (`action: screenshot`) after important actions for evidence.
- For SPA apps: a 200 HTTP status does NOT mean the page is accessible — check the actual
  page content with the browser tool to verify if it shows real data or a login redirect.

## 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"));
    }
}