refactor: modularize codebase and add 404 unit tests (#13)

compliance-dast/src/tools/cookie_analyzer.rs

@@ -14,6 +14,7 @@ pub struct CookieAnalyzerTool {
 #[derive(Debug)]
 struct ParsedCookie {
     name: String,
+    #[allow(dead_code)]
     value: String,
     secure: bool,
     http_only: bool,
@@ -92,6 +93,81 @@ impl CookieAnalyzerTool {
     }
 }
 
+#[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"
@@ -123,96 +199,96 @@ impl PentestTool for CookieAnalyzerTool {
         &'a self,
         input: serde_json::Value,
         context: &'a PentestToolContext,
-    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<PentestToolResult, CoreError>> + Send + 'a>> {
+    ) -> 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 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 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 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();
+            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.into_iter())
-            .collect();
+            // 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}")))?;
+            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 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();
+                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://");
+                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);
+                    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
-                    };
+                    // 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 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(
+                        let mut finding = DastFinding::new(
                         String::new(),
                         target_id.clone(),
                         DastVulnType::CookieSecurity,
@@ -226,32 +302,32 @@ impl PentestTool for CookieAnalyzerTool {
                         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 \
+                        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);
-                }
+                                .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,
-                    };
+                    // 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(
+                        let mut finding = DastFinding::new(
                         String::new(),
                         target_id.clone(),
                         DastVulnType::CookieSecurity,
@@ -265,137 +341,137 @@ impl PentestTool for CookieAnalyzerTool {
                         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 \
+                        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);
-                }
+                                .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,
+                    // Check: missing or weak SameSite
+                    if is_sensitive {
+                        let weak_same_site = match &cookie.same_site {
+                            None => true,
+                            Some(ss) => ss == "none",
                         };
 
-                        let desc = if cookie.same_site.is_none() {
-                            format!(
+                        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!(
+                            } 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 \
+                            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);
+                                    .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,
-                        };
+                    // 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 \
+                            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(
+                                    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);
+                            findings.push(finding);
+                        }
                     }
                 }
             }
-        }
 
-        let count = findings.len();
-        info!(url, findings = count, "Cookie analysis complete");
+            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(),
-            }),
-        })
+            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(),
+                }),
+            })
         })
     }
 }