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
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:
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,6 @@
|
||||
mod context;
|
||||
pub mod orchestrator;
|
||||
mod prompt_builder;
|
||||
pub mod report;
|
||||
|
||||
pub use orchestrator::PentestOrchestrator;
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
+465
-215
@@ -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('>', ">")
|
||||
.replace('"', """)
|
||||
}
|
||||
|
||||
#[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 & b");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn html_escape_handles_angle_brackets() {
|
||||
assert_eq!(html_escape("<script>"), "<script>");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn html_escape_handles_quotes() {
|
||||
assert_eq!(html_escape(r#"key="val""#), "key="val"");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn html_escape_handles_all_special_chars() {
|
||||
assert_eq!(
|
||||
html_escape(r#"<a href="x">&y</a>"#),
|
||||
"<a href="x">&y</a>"
|
||||
);
|
||||
}
|
||||
|
||||
#[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
|
||||
}
|
||||
Reference in New Issue
Block a user