refactor: modularize codebase and add 404 unit tests (#13)
CI / Format (push) Successful in 4s
CI / Clippy (push) Successful in 4m19s
CI / Security Audit (push) Successful in 1m44s
CI / Tests (push) Successful in 5m15s
CI / Detect Changes (push) Successful in 5s
CI / Deploy Agent (push) Successful in 2s
CI / Deploy Dashboard (push) Successful in 2s
CI / Deploy Docs (push) Has been skipped
CI / Deploy MCP (push) Successful in 2s

This commit was merged in pull request #13.
This commit is contained in:
2026-03-13 08:03:45 +00:00
parent acc5b86aa4
commit 3bb690e5bb
89 changed files with 11884 additions and 6046 deletions
+150
View File
@@ -0,0 +1,150 @@
use futures_util::StreamExt;
use mongodb::bson::doc;
use compliance_core::models::dast::DastTarget;
use compliance_core::models::finding::Finding;
use compliance_core::models::pentest::CodeContextHint;
use compliance_core::models::sbom::SbomEntry;
use super::orchestrator::PentestOrchestrator;
impl PentestOrchestrator {
/// Fetch SAST findings, SBOM entries (with CVEs), and code graph entry points
/// for the repo linked to this DAST target.
pub(crate) async fn gather_repo_context(
&self,
target: &DastTarget,
) -> (Vec<Finding>, Vec<SbomEntry>, Vec<CodeContextHint>) {
let Some(repo_id) = &target.repo_id else {
return (Vec::new(), Vec::new(), Vec::new());
};
let sast_findings = self.fetch_sast_findings(repo_id).await;
let sbom_entries = self.fetch_vulnerable_sbom(repo_id).await;
let code_context = self.fetch_code_context(repo_id, &sast_findings).await;
tracing::info!(
repo_id,
sast_findings = sast_findings.len(),
vulnerable_deps = sbom_entries.len(),
code_hints = code_context.len(),
"Gathered code-awareness context for pentest"
);
(sast_findings, sbom_entries, code_context)
}
/// Fetch open/triaged SAST findings for the repo (not false positives or resolved)
async fn fetch_sast_findings(&self, repo_id: &str) -> Vec<Finding> {
let cursor = self
.db
.findings()
.find(doc! {
"repo_id": repo_id,
"status": { "$in": ["open", "triaged"] },
})
.sort(doc! { "severity": -1 })
.limit(100)
.await;
match cursor {
Ok(mut c) => {
let mut results = Vec::new();
while let Some(Ok(f)) = c.next().await {
results.push(f);
}
results
}
Err(e) => {
tracing::warn!("Failed to fetch SAST findings for pentest: {e}");
Vec::new()
}
}
}
/// Fetch SBOM entries that have known vulnerabilities
async fn fetch_vulnerable_sbom(&self, repo_id: &str) -> Vec<SbomEntry> {
let cursor = self
.db
.sbom_entries()
.find(doc! {
"repo_id": repo_id,
"known_vulnerabilities": { "$exists": true, "$ne": [] },
})
.limit(50)
.await;
match cursor {
Ok(mut c) => {
let mut results = Vec::new();
while let Some(Ok(e)) = c.next().await {
results.push(e);
}
results
}
Err(e) => {
tracing::warn!("Failed to fetch vulnerable SBOM entries: {e}");
Vec::new()
}
}
}
/// Build CodeContextHint objects from the code knowledge graph.
/// Maps entry points to their source files and links SAST findings.
async fn fetch_code_context(
&self,
repo_id: &str,
sast_findings: &[Finding],
) -> Vec<CodeContextHint> {
// Get entry point nodes from the code graph
let cursor = self
.db
.graph_nodes()
.find(doc! {
"repo_id": repo_id,
"is_entry_point": true,
})
.limit(50)
.await;
let nodes = match cursor {
Ok(mut c) => {
let mut results = Vec::new();
while let Some(Ok(n)) = c.next().await {
results.push(n);
}
results
}
Err(_) => return Vec::new(),
};
// Build hints by matching graph nodes to SAST findings by file path
nodes
.into_iter()
.map(|node| {
// Find SAST findings in the same file
let linked_vulns: Vec<String> = sast_findings
.iter()
.filter(|f| f.file_path.as_deref() == Some(&node.file_path))
.map(|f| {
format!(
"[{}] {}: {} (line {})",
f.severity,
f.scanner,
f.title,
f.line_number.unwrap_or(0)
)
})
.collect();
CodeContextHint {
endpoint_pattern: node.qualified_name.clone(),
handler_function: node.name.clone(),
file_path: node.file_path.clone(),
code_snippet: String::new(), // Could fetch from embeddings
known_vulnerabilities: linked_vulns,
}
})
.collect()
}
}
+2
View File
@@ -1,4 +1,6 @@
mod context;
pub mod orchestrator;
mod prompt_builder;
pub mod report;
pub use orchestrator::PentestOrchestrator;
+44 -376
View File
@@ -1,31 +1,27 @@
use std::sync::Arc;
use std::time::Duration;
use futures_util::StreamExt;
use mongodb::bson::doc;
use tokio::sync::broadcast;
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 compliance_core::traits::pentest_tool::PentestToolContext;
use compliance_dast::ToolRegistry;
use crate::database::Database;
use crate::llm::client::{
ChatMessage, LlmResponse, ToolCallRequest, ToolCallRequestFunction, ToolDefinition,
use crate::llm::{
ChatMessage, LlmClient, LlmResponse, ToolCallRequest, ToolCallRequestFunction, ToolDefinition,
};
use crate::llm::LlmClient;
/// Maximum duration for a single pentest session before timeout
const SESSION_TIMEOUT: Duration = Duration::from_secs(30 * 60); // 30 minutes
pub struct PentestOrchestrator {
tool_registry: ToolRegistry,
llm: Arc<LlmClient>,
db: Database,
event_tx: broadcast::Sender<PentestEvent>,
pub(crate) tool_registry: ToolRegistry,
pub(crate) llm: Arc<LlmClient>,
pub(crate) db: Database,
pub(crate) event_tx: broadcast::Sender<PentestEvent>,
}
impl PentestOrchestrator {
@@ -39,10 +35,12 @@ impl PentestOrchestrator {
}
}
#[allow(dead_code)]
pub fn subscribe(&self) -> broadcast::Receiver<PentestEvent> {
self.event_tx.subscribe()
}
#[allow(dead_code)]
pub fn event_sender(&self) -> broadcast::Sender<PentestEvent> {
self.event_tx.clone()
}
@@ -111,18 +109,20 @@ impl PentestOrchestrator {
target: &DastTarget,
initial_message: &str,
) -> Result<(), crate::error::AgentError> {
let session_id = session
.id
.map(|oid| oid.to_hex())
.unwrap_or_default();
let session_id = session.id.map(|oid| oid.to_hex()).unwrap_or_default();
// Gather code-awareness context from linked repo
let (sast_findings, sbom_entries, code_context) =
self.gather_repo_context(target).await;
let (sast_findings, sbom_entries, code_context) = self.gather_repo_context(target).await;
// Build system prompt with code context
let system_prompt = self
.build_system_prompt(session, target, &sast_findings, &sbom_entries, &code_context)
.build_system_prompt(
session,
target,
&sast_findings,
&sbom_entries,
&code_context,
)
.await;
// Build tool definitions for LLM
@@ -182,8 +182,7 @@ impl PentestOrchestrator {
match response {
LlmResponse::Content(content) => {
let msg =
PentestMessage::assistant(session_id.clone(), content.clone());
let msg = PentestMessage::assistant(session_id.clone(), content.clone());
let _ = self.db.pentest_messages().insert_one(&msg).await;
let _ = self.event_tx.send(PentestEvent::Message {
content: content.clone(),
@@ -213,7 +212,10 @@ impl PentestOrchestrator {
}
break;
}
LlmResponse::ToolCalls { calls: tool_calls, reasoning } => {
LlmResponse::ToolCalls {
calls: tool_calls,
reasoning,
} => {
let tc_requests: Vec<ToolCallRequest> = tool_calls
.iter()
.map(|tc| ToolCallRequest {
@@ -221,15 +223,18 @@ impl PentestOrchestrator {
r#type: "function".to_string(),
function: ToolCallRequestFunction {
name: tc.name.clone(),
arguments: serde_json::to_string(&tc.arguments)
.unwrap_or_default(),
arguments: serde_json::to_string(&tc.arguments).unwrap_or_default(),
},
})
.collect();
messages.push(ChatMessage {
role: "assistant".to_string(),
content: if reasoning.is_empty() { None } else { Some(reasoning.clone()) },
content: if reasoning.is_empty() {
None
} else {
Some(reasoning.clone())
},
tool_calls: Some(tc_requests),
tool_call_id: None,
});
@@ -274,24 +279,30 @@ impl PentestOrchestrator {
let insert_result =
self.db.dast_findings().insert_one(&finding).await;
if let Ok(res) = &insert_result {
finding_ids.push(res.inserted_id.as_object_id().map(|oid| oid.to_hex()).unwrap_or_default());
}
let _ =
self.event_tx.send(PentestEvent::Finding {
finding_id: finding
.id
finding_ids.push(
res.inserted_id
.as_object_id()
.map(|oid| oid.to_hex())
.unwrap_or_default(),
title: finding.title.clone(),
severity: finding.severity.to_string(),
});
);
}
let _ = self.event_tx.send(PentestEvent::Finding {
finding_id: finding
.id
.map(|oid| oid.to_hex())
.unwrap_or_default(),
title: finding.title.clone(),
severity: finding.severity.to_string(),
});
}
// Compute risk score based on findings severity
let risk_score: Option<u8> = if findings_count > 0 {
Some(std::cmp::min(
100,
(findings_count as u8).saturating_mul(15).saturating_add(20),
(findings_count as u8)
.saturating_mul(15)
.saturating_add(20),
))
} else {
None
@@ -415,347 +426,4 @@ impl PentestOrchestrator {
Ok(())
}
// ── Code-Awareness: Gather context from linked repo ─────────
/// Fetch SAST findings, SBOM entries (with CVEs), and code graph entry points
/// for the repo linked to this DAST target.
async fn gather_repo_context(
&self,
target: &DastTarget,
) -> (Vec<Finding>, Vec<SbomEntry>, Vec<CodeContextHint>) {
let Some(repo_id) = &target.repo_id else {
return (Vec::new(), Vec::new(), Vec::new());
};
let sast_findings = self.fetch_sast_findings(repo_id).await;
let sbom_entries = self.fetch_vulnerable_sbom(repo_id).await;
let code_context = self.fetch_code_context(repo_id, &sast_findings).await;
tracing::info!(
repo_id,
sast_findings = sast_findings.len(),
vulnerable_deps = sbom_entries.len(),
code_hints = code_context.len(),
"Gathered code-awareness context for pentest"
);
(sast_findings, sbom_entries, code_context)
}
/// Fetch open/triaged SAST findings for the repo (not false positives or resolved)
async fn fetch_sast_findings(&self, repo_id: &str) -> Vec<Finding> {
let cursor = self
.db
.findings()
.find(doc! {
"repo_id": repo_id,
"status": { "$in": ["open", "triaged"] },
})
.sort(doc! { "severity": -1 })
.limit(100)
.await;
match cursor {
Ok(mut c) => {
let mut results = Vec::new();
while let Some(Ok(f)) = c.next().await {
results.push(f);
}
results
}
Err(e) => {
tracing::warn!("Failed to fetch SAST findings for pentest: {e}");
Vec::new()
}
}
}
/// Fetch SBOM entries that have known vulnerabilities
async fn fetch_vulnerable_sbom(&self, repo_id: &str) -> Vec<SbomEntry> {
let cursor = self
.db
.sbom_entries()
.find(doc! {
"repo_id": repo_id,
"known_vulnerabilities": { "$exists": true, "$ne": [] },
})
.limit(50)
.await;
match cursor {
Ok(mut c) => {
let mut results = Vec::new();
while let Some(Ok(e)) = c.next().await {
results.push(e);
}
results
}
Err(e) => {
tracing::warn!("Failed to fetch vulnerable SBOM entries: {e}");
Vec::new()
}
}
}
/// Build CodeContextHint objects from the code knowledge graph.
/// Maps entry points to their source files and links SAST findings.
async fn fetch_code_context(
&self,
repo_id: &str,
sast_findings: &[Finding],
) -> Vec<CodeContextHint> {
// Get entry point nodes from the code graph
let cursor = self
.db
.graph_nodes()
.find(doc! {
"repo_id": repo_id,
"is_entry_point": true,
})
.limit(50)
.await;
let nodes = match cursor {
Ok(mut c) => {
let mut results = Vec::new();
while let Some(Ok(n)) = c.next().await {
results.push(n);
}
results
}
Err(_) => return Vec::new(),
};
// Build hints by matching graph nodes to SAST findings by file path
nodes
.into_iter()
.map(|node| {
// Find SAST findings in the same file
let linked_vulns: Vec<String> = sast_findings
.iter()
.filter(|f| {
f.file_path.as_deref() == Some(&node.file_path)
})
.map(|f| {
format!(
"[{}] {}: {} (line {})",
f.severity,
f.scanner,
f.title,
f.line_number.unwrap_or(0)
)
})
.collect();
CodeContextHint {
endpoint_pattern: node.qualified_name.clone(),
handler_function: node.name.clone(),
file_path: node.file_path.clone(),
code_snippet: String::new(), // Could fetch from embeddings
known_vulnerabilities: linked_vulns,
}
})
.collect()
}
// ── System Prompt Builder ───────────────────────────────────
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 strategy_guidance = match session.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 SAST findings section
let sast_section = if sast_findings.is_empty() {
String::from("No SAST findings available for this target.")
} else {
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 SBOM/CVE section
let sbom_section = if sbom_entries.is_empty() {
String::from("No vulnerable dependencies identified.")
} else {
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 code context section
let code_section = if code_context.is_empty() {
String::from("No code knowledge graph available for this target.")
} else {
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
};
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"),
)
}
}
@@ -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"));
}
}
@@ -0,0 +1,43 @@
use std::io::{Cursor, Write};
use zip::write::SimpleFileOptions;
use zip::AesMode;
use super::ReportContext;
pub(super) fn build_zip(
ctx: &ReportContext,
password: &str,
html: &str,
pdf: &[u8],
) -> Result<Vec<u8>, zip::result::ZipError> {
let buf = Cursor::new(Vec::new());
let mut zip = zip::ZipWriter::new(buf);
let options = SimpleFileOptions::default()
.compression_method(zip::CompressionMethod::Deflated)
.with_aes_encryption(AesMode::Aes256, password);
// report.pdf (primary)
zip.start_file("report.pdf", options)?;
zip.write_all(pdf)?;
// report.html (fallback)
zip.start_file("report.html", options)?;
zip.write_all(html.as_bytes())?;
// findings.json
let findings_json =
serde_json::to_string_pretty(&ctx.findings).unwrap_or_else(|_| "[]".to_string());
zip.start_file("findings.json", options)?;
zip.write_all(findings_json.as_bytes())?;
// attack-chain.json
let chain_json =
serde_json::to_string_pretty(&ctx.attack_chain).unwrap_or_else(|_| "[]".to_string());
zip.start_file("attack-chain.json", options)?;
zip.write_all(chain_json.as_bytes())?;
let cursor = zip.finish()?;
Ok(cursor.into_inner())
}
@@ -1,193 +1,50 @@
use std::io::{Cursor, Write};
use compliance_core::models::dast::DastFinding;
use compliance_core::models::pentest::{AttackChainNode, PentestSession};
use sha2::{Digest, Sha256};
use zip::write::SimpleFileOptions;
use zip::AesMode;
use compliance_core::models::pentest::AttackChainNode;
/// Report archive with metadata
pub struct ReportArchive {
/// The password-protected ZIP bytes
pub archive: Vec<u8>,
/// SHA-256 hex digest of the archive
pub sha256: String,
}
use super::ReportContext;
/// Report context gathered from the database
pub struct ReportContext {
pub session: PentestSession,
pub target_name: String,
pub target_url: String,
pub findings: Vec<DastFinding>,
pub attack_chain: Vec<AttackChainNode>,
pub requester_name: String,
pub requester_email: String,
}
/// Generate a password-protected ZIP archive containing the pentest report.
///
/// The archive contains:
/// - `report.pdf` — Professional pentest report (PDF)
/// - `report.html` — HTML source (fallback)
/// - `findings.json` — Raw findings data
/// - `attack-chain.json` — Attack chain timeline
///
/// Files are encrypted with AES-256 inside the ZIP (standard WinZip AES format,
/// supported by 7-Zip, WinRAR, macOS Archive Utility, etc.).
pub async fn generate_encrypted_report(
ctx: &ReportContext,
password: &str,
) -> Result<ReportArchive, String> {
let html = build_html_report(ctx);
// Convert HTML to PDF via headless Chrome
let pdf_bytes = html_to_pdf(&html).await?;
let zip_bytes = build_zip(ctx, password, &html, &pdf_bytes)
.map_err(|e| format!("Failed to create archive: {e}"))?;
let mut hasher = Sha256::new();
hasher.update(&zip_bytes);
let sha256 = hex::encode(hasher.finalize());
Ok(ReportArchive { archive: zip_bytes, sha256 })
}
/// Convert HTML string to PDF bytes using headless Chrome/Chromium.
async fn html_to_pdf(html: &str) -> Result<Vec<u8>, String> {
let tmp_dir = std::env::temp_dir();
let run_id = uuid::Uuid::new_v4().to_string();
let html_path = tmp_dir.join(format!("pentest-report-{run_id}.html"));
let pdf_path = tmp_dir.join(format!("pentest-report-{run_id}.pdf"));
// Write HTML to temp file
std::fs::write(&html_path, html)
.map_err(|e| format!("Failed to write temp HTML: {e}"))?;
// Find Chrome/Chromium binary
let chrome_bin = find_chrome_binary()
.ok_or_else(|| "Chrome/Chromium not found. Install google-chrome or chromium to generate PDF reports.".to_string())?;
tracing::info!(chrome = %chrome_bin, "Generating PDF report via headless Chrome");
let html_url = format!("file://{}", html_path.display());
let output = tokio::process::Command::new(&chrome_bin)
.args([
"--headless",
"--disable-gpu",
"--no-sandbox",
"--disable-software-rasterizer",
"--run-all-compositor-stages-before-draw",
"--disable-dev-shm-usage",
&format!("--print-to-pdf={}", pdf_path.display()),
"--no-pdf-header-footer",
&html_url,
])
.output()
.await
.map_err(|e| format!("Failed to run Chrome: {e}"))?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
// Clean up temp files
let _ = std::fs::remove_file(&html_path);
let _ = std::fs::remove_file(&pdf_path);
return Err(format!("Chrome PDF generation failed: {stderr}"));
}
let pdf_bytes = std::fs::read(&pdf_path)
.map_err(|e| format!("Failed to read generated PDF: {e}"))?;
// Clean up temp files
let _ = std::fs::remove_file(&html_path);
let _ = std::fs::remove_file(&pdf_path);
if pdf_bytes.is_empty() {
return Err("Chrome produced an empty PDF".to_string());
}
tracing::info!(size_kb = pdf_bytes.len() / 1024, "PDF report generated");
Ok(pdf_bytes)
}
/// Search for Chrome/Chromium binary on the system.
fn find_chrome_binary() -> Option<String> {
let candidates = [
"google-chrome-stable",
"google-chrome",
"chromium-browser",
"chromium",
];
for name in &candidates {
if let Ok(output) = std::process::Command::new("which").arg(name).output() {
if output.status.success() {
let path = String::from_utf8_lossy(&output.stdout).trim().to_string();
if !path.is_empty() {
return Some(path);
}
}
}
}
None
}
fn build_zip(
ctx: &ReportContext,
password: &str,
html: &str,
pdf: &[u8],
) -> Result<Vec<u8>, zip::result::ZipError> {
let buf = Cursor::new(Vec::new());
let mut zip = zip::ZipWriter::new(buf);
let options = SimpleFileOptions::default()
.compression_method(zip::CompressionMethod::Deflated)
.with_aes_encryption(AesMode::Aes256, password);
// report.pdf (primary)
zip.start_file("report.pdf", options.clone())?;
zip.write_all(pdf)?;
// report.html (fallback)
zip.start_file("report.html", options.clone())?;
zip.write_all(html.as_bytes())?;
// findings.json
let findings_json =
serde_json::to_string_pretty(&ctx.findings).unwrap_or_else(|_| "[]".to_string());
zip.start_file("findings.json", options.clone())?;
zip.write_all(findings_json.as_bytes())?;
// attack-chain.json
let chain_json =
serde_json::to_string_pretty(&ctx.attack_chain).unwrap_or_else(|_| "[]".to_string());
zip.start_file("attack-chain.json", options)?;
zip.write_all(chain_json.as_bytes())?;
let cursor = zip.finish()?;
Ok(cursor.into_inner())
}
fn build_html_report(ctx: &ReportContext) -> String {
#[allow(clippy::format_in_format_args)]
pub(super) fn build_html_report(ctx: &ReportContext) -> String {
let session = &ctx.session;
let session_id = session
.id
.map(|oid| oid.to_hex())
.unwrap_or_else(|| "-".to_string());
let date_str = session.started_at.format("%B %d, %Y at %H:%M UTC").to_string();
let date_str = session
.started_at
.format("%B %d, %Y at %H:%M UTC")
.to_string();
let date_short = session.started_at.format("%B %d, %Y").to_string();
let completed_str = session
.completed_at
.map(|d| d.format("%B %d, %Y at %H:%M UTC").to_string())
.unwrap_or_else(|| "In Progress".to_string());
let critical = ctx.findings.iter().filter(|f| f.severity.to_string() == "critical").count();
let high = ctx.findings.iter().filter(|f| f.severity.to_string() == "high").count();
let medium = ctx.findings.iter().filter(|f| f.severity.to_string() == "medium").count();
let low = ctx.findings.iter().filter(|f| f.severity.to_string() == "low").count();
let info = ctx.findings.iter().filter(|f| f.severity.to_string() == "info").count();
let critical = ctx
.findings
.iter()
.filter(|f| f.severity.to_string() == "critical")
.count();
let high = ctx
.findings
.iter()
.filter(|f| f.severity.to_string() == "high")
.count();
let medium = ctx
.findings
.iter()
.filter(|f| f.severity.to_string() == "medium")
.count();
let low = ctx
.findings
.iter()
.filter(|f| f.severity.to_string() == "low")
.count();
let info = ctx
.findings
.iter()
.filter(|f| f.severity.to_string() == "info")
.count();
let exploitable = ctx.findings.iter().filter(|f| f.exploitable).count();
let total = ctx.findings.len();
@@ -212,10 +69,8 @@ fn build_html_report(ctx: &ReportContext) -> String {
};
// Risk score 0-100
let risk_score: usize = std::cmp::min(
100,
critical * 25 + high * 15 + medium * 8 + low * 3 + info * 1,
);
let risk_score: usize =
std::cmp::min(100, critical * 25 + high * 15 + medium * 8 + low * 3 + info);
// Collect unique tool names used
let tool_names: Vec<String> = {
@@ -247,7 +102,8 @@ fn build_html_report(ctx: &ReportContext) -> String {
if high > 0 {
bar.push_str(&format!(
r#"<div class="sev-bar-seg sev-bar-high" style="width:{}%"><span>{}</span></div>"#,
std::cmp::max(high_pct, 4), high
std::cmp::max(high_pct, 4),
high
));
}
if medium > 0 {
@@ -259,22 +115,38 @@ fn build_html_report(ctx: &ReportContext) -> String {
if low > 0 {
bar.push_str(&format!(
r#"<div class="sev-bar-seg sev-bar-low" style="width:{}%"><span>{}</span></div>"#,
std::cmp::max(low_pct, 4), low
std::cmp::max(low_pct, 4),
low
));
}
if info > 0 {
bar.push_str(&format!(
r#"<div class="sev-bar-seg sev-bar-info" style="width:{}%"><span>{}</span></div>"#,
std::cmp::max(info_pct, 4), info
std::cmp::max(info_pct, 4),
info
));
}
bar.push_str("</div>");
bar.push_str(r#"<div class="sev-bar-legend">"#);
if critical > 0 { bar.push_str(r#"<span><i class="sev-dot" style="background:#991b1b"></i> Critical</span>"#); }
if high > 0 { bar.push_str(r#"<span><i class="sev-dot" style="background:#c2410c"></i> High</span>"#); }
if medium > 0 { bar.push_str(r#"<span><i class="sev-dot" style="background:#a16207"></i> Medium</span>"#); }
if low > 0 { bar.push_str(r#"<span><i class="sev-dot" style="background:#1d4ed8"></i> Low</span>"#); }
if info > 0 { bar.push_str(r#"<span><i class="sev-dot" style="background:#4b5563"></i> Info</span>"#); }
if critical > 0 {
bar.push_str(
r#"<span><i class="sev-dot" style="background:#991b1b"></i> Critical</span>"#,
);
}
if high > 0 {
bar.push_str(r#"<span><i class="sev-dot" style="background:#c2410c"></i> High</span>"#);
}
if medium > 0 {
bar.push_str(
r#"<span><i class="sev-dot" style="background:#a16207"></i> Medium</span>"#,
);
}
if low > 0 {
bar.push_str(r#"<span><i class="sev-dot" style="background:#1d4ed8"></i> Low</span>"#);
}
if info > 0 {
bar.push_str(r#"<span><i class="sev-dot" style="background:#4b5563"></i> Info</span>"#);
}
bar.push_str("</div>");
bar
} else {
@@ -322,7 +194,12 @@ fn build_html_report(ctx: &ReportContext) -> String {
let param_row = f
.parameter
.as_deref()
.map(|p| format!("<tr><td>Parameter</td><td><code>{}</code></td></tr>", html_escape(p)))
.map(|p| {
format!(
"<tr><td>Parameter</td><td><code>{}</code></td></tr>",
html_escape(p)
)
})
.unwrap_or_default();
let remediation = f
.remediation
@@ -332,7 +209,9 @@ fn build_html_report(ctx: &ReportContext) -> String {
let evidence_html = if f.evidence.is_empty() {
String::new()
} else {
let mut eh = String::from(r#"<div class="evidence-block"><div class="evidence-title">Evidence</div><table class="evidence-table"><thead><tr><th>Request</th><th>Status</th><th>Details</th></tr></thead><tbody>"#);
let mut eh = String::from(
r#"<div class="evidence-block"><div class="evidence-title">Evidence</div><table class="evidence-table"><thead><tr><th>Request</th><th>Status</th><th>Details</th></tr></thead><tbody>"#,
);
for ev in &f.evidence {
let payload_info = ev
.payload
@@ -346,7 +225,7 @@ fn build_html_report(ctx: &ReportContext) -> String {
ev.response_status,
ev.response_snippet
.as_deref()
.map(|s| html_escape(s))
.map(html_escape)
.unwrap_or_default(),
payload_info,
));
@@ -402,7 +281,8 @@ fn build_html_report(ctx: &ReportContext) -> String {
let mut chain_html = String::new();
if !ctx.attack_chain.is_empty() {
// Compute phases via BFS from root nodes
let mut phase_map: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
let mut phase_map: std::collections::HashMap<String, usize> =
std::collections::HashMap::new();
let mut queue: std::collections::VecDeque<String> = std::collections::VecDeque::new();
for node in &ctx.attack_chain {
@@ -438,7 +318,13 @@ fn build_html_report(ctx: &ReportContext) -> String {
// Group nodes by phase
let max_phase = phase_map.values().copied().max().unwrap_or(0);
let phase_labels = ["Reconnaissance", "Enumeration", "Exploitation", "Validation", "Post-Exploitation"];
let phase_labels = [
"Reconnaissance",
"Enumeration",
"Exploitation",
"Validation",
"Post-Exploitation",
];
for phase_idx in 0..=max_phase {
let phase_nodes: Vec<&AttackChainNode> = ctx
@@ -485,15 +371,28 @@ fn build_html_report(ctx: &ReportContext) -> String {
format!(
r#"<span class="step-findings">{} finding{}</span>"#,
node.findings_produced.len(),
if node.findings_produced.len() == 1 { "" } else { "s" },
if node.findings_produced.len() == 1 {
""
} else {
"s"
},
)
} else {
String::new()
};
let risk_badge = node.risk_score.map(|r| {
let risk_class = if r >= 70 { "risk-high" } else if r >= 40 { "risk-med" } else { "risk-low" };
format!(r#"<span class="step-risk {risk_class}">Risk: {r}</span>"#)
}).unwrap_or_default();
let risk_badge = node
.risk_score
.map(|r| {
let risk_class = if r >= 70 {
"risk-high"
} else if r >= 40 {
"risk-med"
} else {
"risk-low"
};
format!(r#"<span class="step-risk {risk_class}">Risk: {r}</span>"#)
})
.unwrap_or_default();
let reasoning_html = if node.llm_reasoning.is_empty() {
String::new()
@@ -547,10 +446,20 @@ fn build_html_report(ctx: &ReportContext) -> String {
let toc_findings_sub = if !ctx.findings.is_empty() {
let mut sub = String::new();
let mut fnum = 0usize;
for (si, &sev_key) in severity_order.iter().enumerate() {
let count = ctx.findings.iter().filter(|f| f.severity.to_string() == sev_key).count();
if count == 0 { continue; }
for f in ctx.findings.iter().filter(|f| f.severity.to_string() == sev_key) {
for &sev_key in severity_order.iter() {
let count = ctx
.findings
.iter()
.filter(|f| f.severity.to_string() == sev_key)
.count();
if count == 0 {
continue;
}
for f in ctx
.findings
.iter()
.filter(|f| f.severity.to_string() == sev_key)
{
fnum += 1;
sub.push_str(&format!(
r#"<div class="toc-sub">F-{:03} — {}</div>"#,
@@ -1577,19 +1486,49 @@ table.tools-table td:first-child {{
fn tool_category(tool_name: &str) -> &'static str {
let name = tool_name.to_lowercase();
if name.contains("nmap") || name.contains("port") { return "Network Reconnaissance"; }
if name.contains("nikto") || name.contains("header") { return "Web Server Analysis"; }
if name.contains("zap") || name.contains("spider") || name.contains("crawl") { return "Web Application Scanning"; }
if name.contains("sqlmap") || name.contains("sqli") || name.contains("sql") { return "SQL Injection Testing"; }
if name.contains("xss") || name.contains("cross-site") { return "Cross-Site Scripting Testing"; }
if name.contains("dir") || name.contains("brute") || name.contains("fuzz") || name.contains("gobuster") { return "Directory Enumeration"; }
if name.contains("ssl") || name.contains("tls") || name.contains("cert") { return "SSL/TLS Analysis"; }
if name.contains("api") || name.contains("endpoint") { return "API Security Testing"; }
if name.contains("auth") || name.contains("login") || name.contains("credential") { return "Authentication Testing"; }
if name.contains("cors") { return "CORS Testing"; }
if name.contains("csrf") { return "CSRF Testing"; }
if name.contains("nuclei") || name.contains("template") { return "Vulnerability Scanning"; }
if name.contains("whatweb") || name.contains("tech") || name.contains("wappalyzer") { return "Technology Fingerprinting"; }
if name.contains("nmap") || name.contains("port") {
return "Network Reconnaissance";
}
if name.contains("nikto") || name.contains("header") {
return "Web Server Analysis";
}
if name.contains("zap") || name.contains("spider") || name.contains("crawl") {
return "Web Application Scanning";
}
if name.contains("sqlmap") || name.contains("sqli") || name.contains("sql") {
return "SQL Injection Testing";
}
if name.contains("xss") || name.contains("cross-site") {
return "Cross-Site Scripting Testing";
}
if name.contains("dir")
|| name.contains("brute")
|| name.contains("fuzz")
|| name.contains("gobuster")
{
return "Directory Enumeration";
}
if name.contains("ssl") || name.contains("tls") || name.contains("cert") {
return "SSL/TLS Analysis";
}
if name.contains("api") || name.contains("endpoint") {
return "API Security Testing";
}
if name.contains("auth") || name.contains("login") || name.contains("credential") {
return "Authentication Testing";
}
if name.contains("cors") {
return "CORS Testing";
}
if name.contains("csrf") {
return "CSRF Testing";
}
if name.contains("nuclei") || name.contains("template") {
return "Vulnerability Scanning";
}
if name.contains("whatweb") || name.contains("tech") || name.contains("wappalyzer") {
return "Technology Fingerprinting";
}
"Security Testing"
}
@@ -1599,3 +1538,314 @@ fn html_escape(s: &str) -> String {
.replace('>', "&gt;")
.replace('"', "&quot;")
}
#[cfg(test)]
mod tests {
use super::*;
use compliance_core::models::dast::{DastFinding, DastVulnType};
use compliance_core::models::finding::Severity;
use compliance_core::models::pentest::{
AttackChainNode, AttackNodeStatus, PentestSession, PentestStrategy,
};
// ── html_escape ──────────────────────────────────────────────────
#[test]
fn html_escape_handles_ampersand() {
assert_eq!(html_escape("a & b"), "a &amp; b");
}
#[test]
fn html_escape_handles_angle_brackets() {
assert_eq!(html_escape("<script>"), "&lt;script&gt;");
}
#[test]
fn html_escape_handles_quotes() {
assert_eq!(html_escape(r#"key="val""#), "key=&quot;val&quot;");
}
#[test]
fn html_escape_handles_all_special_chars() {
assert_eq!(
html_escape(r#"<a href="x">&y</a>"#),
"&lt;a href=&quot;x&quot;&gt;&amp;y&lt;/a&gt;"
);
}
#[test]
fn html_escape_no_change_for_plain_text() {
assert_eq!(html_escape("hello world"), "hello world");
}
#[test]
fn html_escape_empty_string() {
assert_eq!(html_escape(""), "");
}
// ── tool_category ────────────────────────────────────────────────
#[test]
fn tool_category_nmap() {
assert_eq!(tool_category("nmap_scan"), "Network Reconnaissance");
}
#[test]
fn tool_category_port_scanner() {
assert_eq!(tool_category("port_scanner"), "Network Reconnaissance");
}
#[test]
fn tool_category_nikto() {
assert_eq!(tool_category("nikto"), "Web Server Analysis");
}
#[test]
fn tool_category_header_check() {
assert_eq!(
tool_category("security_header_check"),
"Web Server Analysis"
);
}
#[test]
fn tool_category_zap_spider() {
assert_eq!(tool_category("zap_spider"), "Web Application Scanning");
}
#[test]
fn tool_category_sqlmap() {
assert_eq!(tool_category("sqlmap"), "SQL Injection Testing");
}
#[test]
fn tool_category_xss_scanner() {
assert_eq!(tool_category("xss_scanner"), "Cross-Site Scripting Testing");
}
#[test]
fn tool_category_dir_bruteforce() {
assert_eq!(tool_category("dir_bruteforce"), "Directory Enumeration");
}
#[test]
fn tool_category_gobuster() {
assert_eq!(tool_category("gobuster"), "Directory Enumeration");
}
#[test]
fn tool_category_ssl_check() {
assert_eq!(tool_category("ssl_check"), "SSL/TLS Analysis");
}
#[test]
fn tool_category_tls_scan() {
assert_eq!(tool_category("tls_scan"), "SSL/TLS Analysis");
}
#[test]
fn tool_category_api_test() {
assert_eq!(tool_category("api_endpoint_test"), "API Security Testing");
}
#[test]
fn tool_category_auth_bypass() {
assert_eq!(tool_category("auth_bypass_check"), "Authentication Testing");
}
#[test]
fn tool_category_cors() {
assert_eq!(tool_category("cors_check"), "CORS Testing");
}
#[test]
fn tool_category_csrf() {
assert_eq!(tool_category("csrf_scanner"), "CSRF Testing");
}
#[test]
fn tool_category_nuclei() {
assert_eq!(tool_category("nuclei"), "Vulnerability Scanning");
}
#[test]
fn tool_category_whatweb() {
assert_eq!(tool_category("whatweb"), "Technology Fingerprinting");
}
#[test]
fn tool_category_unknown_defaults_to_security_testing() {
assert_eq!(tool_category("custom_tool"), "Security Testing");
}
#[test]
fn tool_category_is_case_insensitive() {
assert_eq!(tool_category("NMAP_Scanner"), "Network Reconnaissance");
assert_eq!(tool_category("SQLMap"), "SQL Injection Testing");
}
// ── build_html_report ────────────────────────────────────────────
fn make_session(strategy: PentestStrategy) -> PentestSession {
let mut s = PentestSession::new("target-1".into(), strategy);
s.tool_invocations = 5;
s.tool_successes = 4;
s.findings_count = 2;
s.exploitable_count = 1;
s
}
fn make_finding(severity: Severity, title: &str, exploitable: bool) -> DastFinding {
let mut f = DastFinding::new(
"run-1".into(),
"target-1".into(),
DastVulnType::Xss,
title.into(),
"description".into(),
severity,
"https://example.com/test".into(),
"GET".into(),
);
f.exploitable = exploitable;
f
}
fn make_attack_node(tool_name: &str) -> AttackChainNode {
let mut node = AttackChainNode::new(
"session-1".into(),
"node-1".into(),
tool_name.into(),
serde_json::json!({}),
"Testing this tool".into(),
);
node.status = AttackNodeStatus::Completed;
node
}
fn make_report_context(
findings: Vec<DastFinding>,
chain: Vec<AttackChainNode>,
) -> ReportContext {
ReportContext {
session: make_session(PentestStrategy::Comprehensive),
target_name: "Test App".into(),
target_url: "https://example.com".into(),
findings,
attack_chain: chain,
requester_name: "Alice".into(),
requester_email: "alice@example.com".into(),
}
}
#[test]
fn report_contains_target_info() {
let ctx = make_report_context(vec![], vec![]);
let html = build_html_report(&ctx);
assert!(html.contains("Test App"));
assert!(html.contains("https://example.com"));
}
#[test]
fn report_contains_requester_info() {
let ctx = make_report_context(vec![], vec![]);
let html = build_html_report(&ctx);
assert!(html.contains("Alice"));
assert!(html.contains("alice@example.com"));
}
#[test]
fn report_shows_informational_risk_when_no_findings() {
let ctx = make_report_context(vec![], vec![]);
let html = build_html_report(&ctx);
assert!(html.contains("INFORMATIONAL"));
}
#[test]
fn report_shows_critical_risk_with_critical_finding() {
let findings = vec![make_finding(Severity::Critical, "Critical XSS", true)];
let ctx = make_report_context(findings, vec![]);
let html = build_html_report(&ctx);
assert!(html.contains("CRITICAL"));
}
#[test]
fn report_shows_high_risk_without_critical() {
let findings = vec![make_finding(Severity::High, "High SQLi", false)];
let ctx = make_report_context(findings, vec![]);
let html = build_html_report(&ctx);
// Should show HIGH, not CRITICAL
assert!(html.contains("HIGH"));
}
#[test]
fn report_shows_medium_risk_level() {
let findings = vec![make_finding(Severity::Medium, "Medium Issue", false)];
let ctx = make_report_context(findings, vec![]);
let html = build_html_report(&ctx);
assert!(html.contains("MEDIUM"));
}
#[test]
fn report_includes_finding_title() {
let findings = vec![make_finding(
Severity::High,
"Reflected XSS in /search",
true,
)];
let ctx = make_report_context(findings, vec![]);
let html = build_html_report(&ctx);
assert!(html.contains("Reflected XSS in /search"));
}
#[test]
fn report_shows_exploitable_badge() {
let findings = vec![make_finding(Severity::Critical, "SQLi", true)];
let ctx = make_report_context(findings, vec![]);
let html = build_html_report(&ctx);
// The report should mark exploitable findings
assert!(html.contains("EXPLOITABLE"));
}
#[test]
fn report_includes_attack_chain_tool_names() {
let chain = vec![make_attack_node("nmap_scan"), make_attack_node("sqlmap")];
let ctx = make_report_context(vec![], chain);
let html = build_html_report(&ctx);
assert!(html.contains("nmap_scan"));
assert!(html.contains("sqlmap"));
}
#[test]
fn report_is_valid_html_structure() {
let ctx = make_report_context(vec![], vec![]);
let html = build_html_report(&ctx);
assert!(html.contains("<!DOCTYPE html>") || html.contains("<html"));
assert!(html.contains("</html>"));
}
#[test]
fn report_strategy_appears() {
let ctx = make_report_context(vec![], vec![]);
let html = build_html_report(&ctx);
// PentestStrategy::Comprehensive => "comprehensive"
assert!(html.contains("comprehensive") || html.contains("Comprehensive"));
}
#[test]
fn report_finding_count_is_correct() {
let findings = vec![
make_finding(Severity::Critical, "F1", true),
make_finding(Severity::High, "F2", false),
make_finding(Severity::Low, "F3", false),
];
let ctx = make_report_context(findings, vec![]);
let html = build_html_report(&ctx);
// The total count "3" should appear somewhere
assert!(
html.contains(">3<")
|| html.contains(">3 ")
|| html.contains("3 findings")
|| html.contains("3 Total")
);
}
}
@@ -0,0 +1,58 @@
mod archive;
mod html;
mod pdf;
use compliance_core::models::dast::DastFinding;
use compliance_core::models::pentest::{AttackChainNode, PentestSession};
use sha2::{Digest, Sha256};
/// Report archive with metadata
pub struct ReportArchive {
/// The password-protected ZIP bytes
pub archive: Vec<u8>,
/// SHA-256 hex digest of the archive
pub sha256: String,
}
/// Report context gathered from the database
pub struct ReportContext {
pub session: PentestSession,
pub target_name: String,
pub target_url: String,
pub findings: Vec<DastFinding>,
pub attack_chain: Vec<AttackChainNode>,
pub requester_name: String,
pub requester_email: String,
}
/// Generate a password-protected ZIP archive containing the pentest report.
///
/// The archive contains:
/// - `report.pdf` — Professional pentest report (PDF)
/// - `report.html` — HTML source (fallback)
/// - `findings.json` — Raw findings data
/// - `attack-chain.json` — Attack chain timeline
///
/// Files are encrypted with AES-256 inside the ZIP (standard WinZip AES format,
/// supported by 7-Zip, WinRAR, macOS Archive Utility, etc.).
pub async fn generate_encrypted_report(
ctx: &ReportContext,
password: &str,
) -> Result<ReportArchive, String> {
let html = html::build_html_report(ctx);
// Convert HTML to PDF via headless Chrome
let pdf_bytes = pdf::html_to_pdf(&html).await?;
let zip_bytes = archive::build_zip(ctx, password, &html, &pdf_bytes)
.map_err(|e| format!("Failed to create archive: {e}"))?;
let mut hasher = Sha256::new();
hasher.update(&zip_bytes);
let sha256 = hex::encode(hasher.finalize());
Ok(ReportArchive {
archive: zip_bytes,
sha256,
})
}
@@ -0,0 +1,79 @@
/// Convert HTML string to PDF bytes using headless Chrome/Chromium.
pub(super) async fn html_to_pdf(html: &str) -> Result<Vec<u8>, String> {
let tmp_dir = std::env::temp_dir();
let run_id = uuid::Uuid::new_v4().to_string();
let html_path = tmp_dir.join(format!("pentest-report-{run_id}.html"));
let pdf_path = tmp_dir.join(format!("pentest-report-{run_id}.pdf"));
// Write HTML to temp file
std::fs::write(&html_path, html).map_err(|e| format!("Failed to write temp HTML: {e}"))?;
// Find Chrome/Chromium binary
let chrome_bin = find_chrome_binary().ok_or_else(|| {
"Chrome/Chromium not found. Install google-chrome or chromium to generate PDF reports."
.to_string()
})?;
tracing::info!(chrome = %chrome_bin, "Generating PDF report via headless Chrome");
let html_url = format!("file://{}", html_path.display());
let output = tokio::process::Command::new(&chrome_bin)
.args([
"--headless",
"--disable-gpu",
"--no-sandbox",
"--disable-software-rasterizer",
"--run-all-compositor-stages-before-draw",
"--disable-dev-shm-usage",
&format!("--print-to-pdf={}", pdf_path.display()),
"--no-pdf-header-footer",
&html_url,
])
.output()
.await
.map_err(|e| format!("Failed to run Chrome: {e}"))?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
// Clean up temp files
let _ = std::fs::remove_file(&html_path);
let _ = std::fs::remove_file(&pdf_path);
return Err(format!("Chrome PDF generation failed: {stderr}"));
}
let pdf_bytes =
std::fs::read(&pdf_path).map_err(|e| format!("Failed to read generated PDF: {e}"))?;
// Clean up temp files
let _ = std::fs::remove_file(&html_path);
let _ = std::fs::remove_file(&pdf_path);
if pdf_bytes.is_empty() {
return Err("Chrome produced an empty PDF".to_string());
}
tracing::info!(size_kb = pdf_bytes.len() / 1024, "PDF report generated");
Ok(pdf_bytes)
}
/// Search for Chrome/Chromium binary on the system.
fn find_chrome_binary() -> Option<String> {
let candidates = [
"google-chrome-stable",
"google-chrome",
"chromium-browser",
"chromium",
];
for name in &candidates {
if let Ok(output) = std::process::Command::new("which").arg(name).output() {
if output.status.success() {
let path = String::from_utf8_lossy(&output.stdout).trim().to_string();
if !path.is_empty() {
return Some(path);
}
}
}
}
None
}