feat: AI-driven automated penetration testing (#12)

compliance-dast/src/tools/security_headers.rs

@@ -0,0 +1,300 @@
+use std::collections::HashMap;
+
+use compliance_core::error::CoreError;
+use compliance_core::models::dast::{DastEvidence, DastFinding, DastVulnType};
+use compliance_core::models::Severity;
+use compliance_core::traits::pentest_tool::{PentestTool, PentestToolContext, PentestToolResult};
+use serde_json::json;
+use tracing::info;
+
+/// Tool that checks for the presence and correctness of security headers.
+pub struct SecurityHeadersTool {
+    http: reqwest::Client,
+}
+
+/// A security header we expect to be present and its metadata.
+struct ExpectedHeader {
+    name: &'static str,
+    description: &'static str,
+    severity: Severity,
+    cwe: &'static str,
+    remediation: &'static str,
+    /// If present, the value must contain one of these substrings to be considered valid.
+    valid_values: Option<Vec<&'static str>>,
+}
+
+impl SecurityHeadersTool {
+    pub fn new(http: reqwest::Client) -> Self {
+        Self { http }
+    }
+
+    fn expected_headers() -> Vec<ExpectedHeader> {
+        vec![
+            ExpectedHeader {
+                name: "strict-transport-security",
+                description: "HTTP Strict Transport Security (HSTS) forces browsers to use HTTPS",
+                severity: Severity::Medium,
+                cwe: "CWE-319",
+                remediation: "Add 'Strict-Transport-Security: max-age=31536000; includeSubDomains' header.",
+                valid_values: None,
+            },
+            ExpectedHeader {
+                name: "x-content-type-options",
+                description: "Prevents MIME type sniffing",
+                severity: Severity::Low,
+                cwe: "CWE-16",
+                remediation: "Add 'X-Content-Type-Options: nosniff' header.",
+                valid_values: Some(vec!["nosniff"]),
+            },
+            ExpectedHeader {
+                name: "x-frame-options",
+                description: "Prevents clickjacking by controlling iframe embedding",
+                severity: Severity::Medium,
+                cwe: "CWE-1021",
+                remediation: "Add 'X-Frame-Options: DENY' or 'X-Frame-Options: SAMEORIGIN' header.",
+                valid_values: Some(vec!["deny", "sameorigin"]),
+            },
+            ExpectedHeader {
+                name: "x-xss-protection",
+                description: "Enables browser XSS filtering (legacy but still recommended)",
+                severity: Severity::Low,
+                cwe: "CWE-79",
+                remediation: "Add 'X-XSS-Protection: 1; mode=block' header.",
+                valid_values: None,
+            },
+            ExpectedHeader {
+                name: "referrer-policy",
+                description: "Controls how much referrer information is shared",
+                severity: Severity::Low,
+                cwe: "CWE-200",
+                remediation: "Add 'Referrer-Policy: strict-origin-when-cross-origin' or 'no-referrer' header.",
+                valid_values: None,
+            },
+            ExpectedHeader {
+                name: "permissions-policy",
+                description: "Controls browser feature access (camera, microphone, geolocation, etc.)",
+                severity: Severity::Low,
+                cwe: "CWE-16",
+                remediation: "Add a Permissions-Policy header to restrict browser feature access. \
+                              Example: 'Permissions-Policy: camera=(), microphone=(), geolocation=()'.",
+                valid_values: None,
+            },
+        ]
+    }
+}
+
+impl PentestTool for SecurityHeadersTool {
+    fn name(&self) -> &str {
+        "security_headers"
+    }
+
+    fn description(&self) -> &str {
+        "Checks a URL for the presence and correctness of security headers: HSTS, \
+         X-Content-Type-Options, X-Frame-Options, X-XSS-Protection, Referrer-Policy, \
+         and Permissions-Policy."
+    }
+
+    fn input_schema(&self) -> serde_json::Value {
+        json!({
+            "type": "object",
+            "properties": {
+                "url": {
+                    "type": "string",
+                    "description": "URL to check security headers for"
+                }
+            },
+            "required": ["url"]
+        })
+    }
+
+    fn execute<'a>(
+        &'a self,
+        input: serde_json::Value,
+        context: &'a PentestToolContext,
+    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<PentestToolResult, CoreError>> + Send + 'a>> {
+        Box::pin(async move {
+        let url = input
+            .get("url")
+            .and_then(|v| v.as_str())
+            .ok_or_else(|| CoreError::Dast("Missing required 'url' parameter".to_string()))?;
+
+        let target_id = context
+            .target
+            .id
+            .map(|oid| oid.to_hex())
+            .unwrap_or_else(|| "unknown".to_string());
+
+        let response = self
+            .http
+            .get(url)
+            .send()
+            .await
+            .map_err(|e| CoreError::Dast(format!("Failed to fetch {url}: {e}")))?;
+
+        let status = response.status().as_u16();
+        let response_headers: HashMap<String, String> = response
+            .headers()
+            .iter()
+            .map(|(k, v)| (k.to_string().to_lowercase(), v.to_str().unwrap_or("").to_string()))
+            .collect();
+
+        let mut findings = Vec::new();
+        let mut header_results: HashMap<String, serde_json::Value> = HashMap::new();
+
+        for expected in Self::expected_headers() {
+            let header_value = response_headers.get(expected.name);
+
+            match header_value {
+                Some(value) => {
+                    let mut is_valid = true;
+                    if let Some(ref valid) = expected.valid_values {
+                        let lower = value.to_lowercase();
+                        is_valid = valid.iter().any(|v| lower.contains(v));
+                    }
+
+                    header_results.insert(
+                        expected.name.to_string(),
+                        json!({
+                            "present": true,
+                            "value": value,
+                            "valid": is_valid,
+                        }),
+                    );
+
+                    if !is_valid {
+                        let evidence = DastEvidence {
+                            request_method: "GET".to_string(),
+                            request_url: url.to_string(),
+                            request_headers: None,
+                            request_body: None,
+                            response_status: status,
+                            response_headers: Some(response_headers.clone()),
+                            response_snippet: Some(format!("{}: {}", expected.name, value)),
+                            screenshot_path: None,
+                            payload: None,
+                            response_time_ms: None,
+                        };
+
+                        let mut finding = DastFinding::new(
+                            String::new(),
+                            target_id.clone(),
+                            DastVulnType::SecurityHeaderMissing,
+                            format!("Invalid {} header value", expected.name),
+                            format!(
+                                "The {} header is present but has an invalid or weak value: '{}'. \
+                                 {}",
+                                expected.name, value, expected.description
+                            ),
+                            expected.severity.clone(),
+                            url.to_string(),
+                            "GET".to_string(),
+                        );
+                        finding.cwe = Some(expected.cwe.to_string());
+                        finding.evidence = vec![evidence];
+                        finding.remediation = Some(expected.remediation.to_string());
+                        findings.push(finding);
+                    }
+                }
+                None => {
+                    header_results.insert(
+                        expected.name.to_string(),
+                        json!({
+                            "present": false,
+                            "value": null,
+                            "valid": false,
+                        }),
+                    );
+
+                    let evidence = DastEvidence {
+                        request_method: "GET".to_string(),
+                        request_url: url.to_string(),
+                        request_headers: None,
+                        request_body: None,
+                        response_status: status,
+                        response_headers: Some(response_headers.clone()),
+                        response_snippet: Some(format!("{} header is missing", expected.name)),
+                        screenshot_path: None,
+                        payload: None,
+                        response_time_ms: None,
+                    };
+
+                    let mut finding = DastFinding::new(
+                        String::new(),
+                        target_id.clone(),
+                        DastVulnType::SecurityHeaderMissing,
+                        format!("Missing {} header", expected.name),
+                        format!(
+                            "The {} header is not present in the response. {}",
+                            expected.name, expected.description
+                        ),
+                        expected.severity.clone(),
+                        url.to_string(),
+                        "GET".to_string(),
+                    );
+                    finding.cwe = Some(expected.cwe.to_string());
+                    finding.evidence = vec![evidence];
+                    finding.remediation = Some(expected.remediation.to_string());
+                    findings.push(finding);
+                }
+            }
+        }
+
+        // Also check for information disclosure headers
+        let disclosure_headers = ["server", "x-powered-by", "x-aspnet-version", "x-aspnetmvc-version"];
+        for h in &disclosure_headers {
+            if let Some(value) = response_headers.get(*h) {
+                header_results.insert(
+                    format!("{h}_disclosure"),
+                    json!({ "present": true, "value": value }),
+                );
+
+                let evidence = DastEvidence {
+                    request_method: "GET".to_string(),
+                    request_url: url.to_string(),
+                    request_headers: None,
+                    request_body: None,
+                    response_status: status,
+                    response_headers: Some(response_headers.clone()),
+                    response_snippet: Some(format!("{h}: {value}")),
+                    screenshot_path: None,
+                    payload: None,
+                    response_time_ms: None,
+                };
+
+                let mut finding = DastFinding::new(
+                    String::new(),
+                    target_id.clone(),
+                    DastVulnType::SecurityHeaderMissing,
+                    format!("Information disclosure via {h} header"),
+                    format!(
+                        "The {h} header exposes server technology information: '{value}'. \
+                         This helps attackers fingerprint the server and find known vulnerabilities."
+                    ),
+                    Severity::Info,
+                    url.to_string(),
+                    "GET".to_string(),
+                );
+                finding.cwe = Some("CWE-200".to_string());
+                finding.evidence = vec![evidence];
+                finding.remediation = Some(format!(
+                    "Remove or suppress the {h} header in your server configuration."
+                ));
+                findings.push(finding);
+            }
+        }
+
+        let count = findings.len();
+        info!(url, findings = count, "Security headers check complete");
+
+        Ok(PentestToolResult {
+            summary: if count > 0 {
+                format!("Found {count} security header issues for {url}.")
+            } else {
+                format!("All checked security headers are present and valid for {url}.")
+            },
+            findings,
+            data: json!(header_results),
+        })
+        })
+    }
+}