compliance-scanner-agent/compliance-dast/src/tools/security_headers.rs

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),
            })
        })
    }
}