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

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 inspects cookies set by a target for security issues.
pub struct CookieAnalyzerTool {
    http: reqwest::Client,
}

/// Parsed attributes from a Set-Cookie header.
#[derive(Debug)]
struct ParsedCookie {
    name: String,
    #[allow(dead_code)]
    value: String,
    secure: bool,
    http_only: bool,
    same_site: Option<String>,
    domain: Option<String>,
    path: Option<String>,
    raw: String,
}

impl CookieAnalyzerTool {
    pub fn new(http: reqwest::Client) -> Self {
        Self { http }
    }

    /// Parse a Set-Cookie header value into a structured representation.
    fn parse_set_cookie(header: &str) -> ParsedCookie {
        let raw = header.to_string();
        let parts: Vec<&str> = header.split(';').collect();

        let (name, value) = if let Some(kv) = parts.first() {
            let mut kv_split = kv.splitn(2, '=');
            let k = kv_split.next().unwrap_or("").trim().to_string();
            let v = kv_split.next().unwrap_or("").trim().to_string();
            (k, v)
        } else {
            (String::new(), String::new())
        };

        let mut secure = false;
        let mut http_only = false;
        let mut same_site = None;
        let mut domain = None;
        let mut path = None;

        for part in parts.iter().skip(1) {
            let trimmed = part.trim().to_lowercase();
            if trimmed == "secure" {
                secure = true;
            } else if trimmed == "httponly" {
                http_only = true;
            } else if let Some(ss) = trimmed.strip_prefix("samesite=") {
                same_site = Some(ss.trim().to_string());
            } else if let Some(d) = trimmed.strip_prefix("domain=") {
                domain = Some(d.trim().to_string());
            } else if let Some(p) = trimmed.strip_prefix("path=") {
                path = Some(p.trim().to_string());
            }
        }

        ParsedCookie {
            name,
            value,
            secure,
            http_only,
            same_site,
            domain,
            path,
            raw,
        }
    }

    /// Heuristic: does this cookie name suggest it's a session / auth cookie?
    fn is_sensitive_cookie(name: &str) -> bool {
        let lower = name.to_lowercase();
        lower.contains("session")
            || lower.contains("sess")
            || lower.contains("token")
            || lower.contains("auth")
            || lower.contains("jwt")
            || lower.contains("csrf")
            || lower.contains("sid")
            || lower == "connect.sid"
            || lower == "phpsessid"
            || lower == "jsessionid"
            || lower == "asp.net_sessionid"
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse_simple_cookie() {
        let cookie = CookieAnalyzerTool::parse_set_cookie("session_id=abc123");
        assert_eq!(cookie.name, "session_id");
        assert_eq!(cookie.value, "abc123");
        assert!(!cookie.secure);
        assert!(!cookie.http_only);
        assert!(cookie.same_site.is_none());
        assert!(cookie.domain.is_none());
        assert!(cookie.path.is_none());
    }

    #[test]
    fn parse_cookie_with_all_attributes() {
        let raw = "token=xyz; Secure; HttpOnly; SameSite=Strict; Domain=.example.com; Path=/api";
        let cookie = CookieAnalyzerTool::parse_set_cookie(raw);
        assert_eq!(cookie.name, "token");
        assert_eq!(cookie.value, "xyz");
        assert!(cookie.secure);
        assert!(cookie.http_only);
        assert_eq!(cookie.same_site.as_deref(), Some("strict"));
        assert_eq!(cookie.domain.as_deref(), Some(".example.com"));
        assert_eq!(cookie.path.as_deref(), Some("/api"));
        assert_eq!(cookie.raw, raw);
    }

    #[test]
    fn parse_cookie_samesite_none() {
        let cookie = CookieAnalyzerTool::parse_set_cookie("id=1; SameSite=None; Secure");
        assert_eq!(cookie.same_site.as_deref(), Some("none"));
        assert!(cookie.secure);
    }

    #[test]
    fn parse_cookie_with_equals_in_value() {
        let cookie = CookieAnalyzerTool::parse_set_cookie("data=a=b=c; HttpOnly");
        assert_eq!(cookie.name, "data");
        assert_eq!(cookie.value, "a=b=c");
        assert!(cookie.http_only);
    }

    #[test]
    fn is_sensitive_cookie_known_names() {
        assert!(CookieAnalyzerTool::is_sensitive_cookie("session_id"));
        assert!(CookieAnalyzerTool::is_sensitive_cookie("PHPSESSID"));
        assert!(CookieAnalyzerTool::is_sensitive_cookie("JSESSIONID"));
        assert!(CookieAnalyzerTool::is_sensitive_cookie("connect.sid"));
        assert!(CookieAnalyzerTool::is_sensitive_cookie("asp.net_sessionid"));
        assert!(CookieAnalyzerTool::is_sensitive_cookie("auth_token"));
        assert!(CookieAnalyzerTool::is_sensitive_cookie("jwt_access"));
        assert!(CookieAnalyzerTool::is_sensitive_cookie("csrf_token"));
        assert!(CookieAnalyzerTool::is_sensitive_cookie("my_sess_cookie"));
        assert!(CookieAnalyzerTool::is_sensitive_cookie("SID"));
    }

    #[test]
    fn is_sensitive_cookie_non_sensitive() {
        assert!(!CookieAnalyzerTool::is_sensitive_cookie("theme"));
        assert!(!CookieAnalyzerTool::is_sensitive_cookie("language"));
        assert!(!CookieAnalyzerTool::is_sensitive_cookie("_ga"));
        assert!(!CookieAnalyzerTool::is_sensitive_cookie("tracking"));
    }

    #[test]
    fn parse_empty_cookie_header() {
        let cookie = CookieAnalyzerTool::parse_set_cookie("");
        assert_eq!(cookie.name, "");
        assert_eq!(cookie.value, "");
    }
}

impl PentestTool for CookieAnalyzerTool {
    fn name(&self) -> &str {
        "cookie_analyzer"
    }

    fn description(&self) -> &str {
        "Analyzes cookies set by a target URL. Checks for Secure, HttpOnly, SameSite attributes \
         and overly broad Domain/Path settings. Focuses on session and authentication cookies."
    }

    fn input_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "url": {
                    "type": "string",
                    "description": "URL to fetch and analyze cookies from"
                },
                "login_url": {
                    "type": "string",
                    "description": "Optional login URL to also check (may set auth cookies)"
                }
            },
            "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 login_url = input.get("login_url").and_then(|v| v.as_str());

            let target_id = context
                .target
                .id
                .map(|oid| oid.to_hex())
                .unwrap_or_else(|| "unknown".to_string());

            let mut findings = Vec::new();
            let mut cookie_data = Vec::new();

            // Collect Set-Cookie headers from the main URL and optional login URL
            let urls_to_check: Vec<&str> = std::iter::once(url).chain(login_url).collect();

            for check_url in &urls_to_check {
                // Use a client that does NOT follow redirects so we catch cookies on redirect responses
                let no_redirect_client = reqwest::Client::builder()
                    .danger_accept_invalid_certs(true)
                    .redirect(reqwest::redirect::Policy::none())
                    .timeout(std::time::Duration::from_secs(15))
                    .build()
                    .map_err(|e| CoreError::Dast(format!("Client build error: {e}")))?;

                let response = match no_redirect_client.get(*check_url).send().await {
                    Ok(r) => r,
                    Err(_e) => {
                        // Try with the main client that follows redirects
                        match self.http.get(*check_url).send().await {
                            Ok(r) => r,
                            Err(_) => continue,
                        }
                    }
                };

                let status = response.status().as_u16();
                let set_cookie_headers: Vec<String> = response
                    .headers()
                    .get_all("set-cookie")
                    .iter()
                    .filter_map(|v| v.to_str().ok().map(String::from))
                    .collect();

                for raw_cookie in &set_cookie_headers {
                    let cookie = Self::parse_set_cookie(raw_cookie);
                    let is_sensitive = Self::is_sensitive_cookie(&cookie.name);
                    let is_https = check_url.starts_with("https://");

                    let cookie_info = json!({
                        "name": cookie.name,
                        "secure": cookie.secure,
                        "http_only": cookie.http_only,
                        "same_site": cookie.same_site,
                        "domain": cookie.domain,
                        "path": cookie.path,
                        "is_sensitive": is_sensitive,
                        "url": check_url,
                    });
                    cookie_data.push(cookie_info);

                    // Check: missing Secure flag
                    if !cookie.secure && (is_https || is_sensitive) {
                        let severity = if is_sensitive {
                            Severity::High
                        } else {
                            Severity::Medium
                        };

                        let evidence = DastEvidence {
                            request_method: "GET".to_string(),
                            request_url: check_url.to_string(),
                            request_headers: None,
                            request_body: None,
                            response_status: status,
                            response_headers: None,
                            response_snippet: Some(cookie.raw.clone()),
                            screenshot_path: None,
                            payload: None,
                            response_time_ms: None,
                        };

                        let mut finding = DastFinding::new(
                        String::new(),
                        target_id.clone(),
                        DastVulnType::CookieSecurity,
                        format!("Cookie '{}' missing Secure flag", cookie.name),
                        format!(
                            "The cookie '{}' does not have the Secure attribute set. \
                             Without this flag, the cookie can be transmitted over unencrypted HTTP connections.",
                            cookie.name
                        ),
                        severity,
                        check_url.to_string(),
                        "GET".to_string(),
                    );
                        finding.cwe = Some("CWE-614".to_string());
                        finding.evidence = vec![evidence];
                        finding.remediation = Some(
                            "Add the 'Secure' attribute to the Set-Cookie header to ensure the \
                         cookie is only sent over HTTPS connections."
                                .to_string(),
                        );
                        findings.push(finding);
                    }

                    // Check: missing HttpOnly flag on sensitive cookies
                    if !cookie.http_only && is_sensitive {
                        let evidence = DastEvidence {
                            request_method: "GET".to_string(),
                            request_url: check_url.to_string(),
                            request_headers: None,
                            request_body: None,
                            response_status: status,
                            response_headers: None,
                            response_snippet: Some(cookie.raw.clone()),
                            screenshot_path: None,
                            payload: None,
                            response_time_ms: None,
                        };

                        let mut finding = DastFinding::new(
                        String::new(),
                        target_id.clone(),
                        DastVulnType::CookieSecurity,
                        format!("Cookie '{}' missing HttpOnly flag", cookie.name),
                        format!(
                            "The session/auth cookie '{}' does not have the HttpOnly attribute. \
                             This makes it accessible to JavaScript, increasing the impact of XSS attacks.",
                            cookie.name
                        ),
                        Severity::High,
                        check_url.to_string(),
                        "GET".to_string(),
                    );
                        finding.cwe = Some("CWE-1004".to_string());
                        finding.evidence = vec![evidence];
                        finding.remediation = Some(
                            "Add the 'HttpOnly' attribute to the Set-Cookie header to prevent \
                         JavaScript access to the cookie."
                                .to_string(),
                        );
                        findings.push(finding);
                    }

                    // Check: missing or weak SameSite
                    if is_sensitive {
                        let weak_same_site = match &cookie.same_site {
                            None => true,
                            Some(ss) => ss == "none",
                        };

                        if weak_same_site {
                            let evidence = DastEvidence {
                                request_method: "GET".to_string(),
                                request_url: check_url.to_string(),
                                request_headers: None,
                                request_body: None,
                                response_status: status,
                                response_headers: None,
                                response_snippet: Some(cookie.raw.clone()),
                                screenshot_path: None,
                                payload: None,
                                response_time_ms: None,
                            };

                            let desc = if cookie.same_site.is_none() {
                                format!(
                                "The session/auth cookie '{}' does not have a SameSite attribute. \
                                 This may allow cross-site request forgery (CSRF) attacks.",
                                cookie.name
                            )
                            } else {
                                format!(
                                "The session/auth cookie '{}' has SameSite=None, which allows it \
                                 to be sent in cross-site requests, enabling CSRF attacks.",
                                cookie.name
                            )
                            };

                            let mut finding = DastFinding::new(
                                String::new(),
                                target_id.clone(),
                                DastVulnType::CookieSecurity,
                                format!("Cookie '{}' missing or weak SameSite", cookie.name),
                                desc,
                                Severity::Medium,
                                check_url.to_string(),
                                "GET".to_string(),
                            );
                            finding.cwe = Some("CWE-1275".to_string());
                            finding.evidence = vec![evidence];
                            finding.remediation = Some(
                                "Set 'SameSite=Strict' or 'SameSite=Lax' on session/auth cookies \
                             to prevent cross-site request inclusion."
                                    .to_string(),
                            );
                            findings.push(finding);
                        }
                    }

                    // Check: overly broad domain
                    if let Some(ref domain) = cookie.domain {
                        // A domain starting with a dot applies to all subdomains
                        let dot_domain = domain.starts_with('.');
                        // Count domain parts - if only 2 parts (e.g., .example.com), it's broad
                        let parts: Vec<&str> = domain.trim_start_matches('.').split('.').collect();
                        if dot_domain && parts.len() <= 2 && is_sensitive {
                            let evidence = DastEvidence {
                                request_method: "GET".to_string(),
                                request_url: check_url.to_string(),
                                request_headers: None,
                                request_body: None,
                                response_status: status,
                                response_headers: None,
                                response_snippet: Some(cookie.raw.clone()),
                                screenshot_path: None,
                                payload: None,
                                response_time_ms: None,
                            };

                            let mut finding = DastFinding::new(
                                String::new(),
                                target_id.clone(),
                                DastVulnType::CookieSecurity,
                                format!("Cookie '{}' has overly broad domain", cookie.name),
                                format!(
                                    "The cookie '{}' is scoped to domain '{}' which includes all \
                                 subdomains. If any subdomain is compromised, the attacker can \
                                 access this cookie.",
                                    cookie.name, domain
                                ),
                                Severity::Low,
                                check_url.to_string(),
                                "GET".to_string(),
                            );
                            finding.cwe = Some("CWE-1004".to_string());
                            finding.evidence = vec![evidence];
                            finding.remediation = Some(
                            "Restrict the cookie domain to the specific subdomain that needs it \
                             rather than the entire parent domain."
                                .to_string(),
                        );
                            findings.push(finding);
                        }
                    }
                }
            }

            let count = findings.len();
            info!(url, findings = count, "Cookie analysis complete");

            Ok(PentestToolResult {
                summary: if count > 0 {
                    format!("Found {count} cookie security issues.")
                } else if cookie_data.is_empty() {
                    "No cookies were set by the target.".to_string()
                } else {
                    "All cookies have proper security attributes.".to_string()
                },
                findings,
                data: json!({
                    "cookies": cookie_data,
                    "total_cookies": cookie_data.len(),
                }),
            })
        })
    }
}