compliance-scanner-agent/compliance-dast/src/tools/dmarc_checker.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, warn};

/// Tool that checks email security configuration (DMARC and SPF records).
pub struct DmarcCheckerTool;

impl Default for DmarcCheckerTool {
    fn default() -> Self {
        Self::new()
    }
}

impl DmarcCheckerTool {
    pub fn new() -> Self {
        Self
    }

    /// Query TXT records for a given name using `dig`.
    async fn query_txt(name: &str) -> Result<Vec<String>, CoreError> {
        let output = tokio::process::Command::new("dig")
            .args(["+short", "TXT", name])
            .output()
            .await
            .map_err(|e| CoreError::Dast(format!("dig command failed: {e}")))?;

        let stdout = String::from_utf8_lossy(&output.stdout);
        let lines: Vec<String> = stdout
            .lines()
            .map(|l| l.trim().trim_matches('"').to_string())
            .filter(|l| !l.is_empty())
            .collect();
        Ok(lines)
    }

    /// Parse a DMARC record string and return the policy value.
    fn parse_dmarc_policy(record: &str) -> Option<String> {
        for part in record.split(';') {
            let part = part.trim();
            if let Some(val) = part.strip_prefix("p=") {
                return Some(val.trim().to_lowercase());
            }
        }
        None
    }

    /// Parse DMARC record for sub-domain policy (sp=).
    fn parse_dmarc_subdomain_policy(record: &str) -> Option<String> {
        for part in record.split(';') {
            let part = part.trim();
            if let Some(val) = part.strip_prefix("sp=") {
                return Some(val.trim().to_lowercase());
            }
        }
        None
    }

    /// Parse DMARC record for reporting URI (rua=).
    fn parse_dmarc_rua(record: &str) -> Option<String> {
        for part in record.split(';') {
            let part = part.trim();
            if let Some(val) = part.strip_prefix("rua=") {
                return Some(val.trim().to_string());
            }
        }
        None
    }

    /// Check if an SPF record is present and parse the policy.
    fn is_spf_record(record: &str) -> bool {
        record.starts_with("v=spf1")
    }

    /// Evaluate SPF record strength.
    fn spf_uses_soft_fail(record: &str) -> bool {
        record.contains("~all")
    }

    fn spf_allows_all(record: &str) -> bool {
        record.contains("+all")
    }
}

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

    #[test]
    fn parse_dmarc_policy_reject() {
        let record = "v=DMARC1; p=reject; rua=mailto:dmarc@example.com";
        assert_eq!(
            DmarcCheckerTool::parse_dmarc_policy(record),
            Some("reject".to_string())
        );
    }

    #[test]
    fn parse_dmarc_policy_none() {
        let record = "v=DMARC1; p=none";
        assert_eq!(
            DmarcCheckerTool::parse_dmarc_policy(record),
            Some("none".to_string())
        );
    }

    #[test]
    fn parse_dmarc_policy_quarantine() {
        let record = "v=DMARC1; p=quarantine; sp=none";
        assert_eq!(
            DmarcCheckerTool::parse_dmarc_policy(record),
            Some("quarantine".to_string())
        );
    }

    #[test]
    fn parse_dmarc_policy_missing() {
        let record = "v=DMARC1; rua=mailto:test@example.com";
        assert_eq!(DmarcCheckerTool::parse_dmarc_policy(record), None);
    }

    #[test]
    fn parse_dmarc_subdomain_policy() {
        let record = "v=DMARC1; p=reject; sp=quarantine";
        assert_eq!(
            DmarcCheckerTool::parse_dmarc_subdomain_policy(record),
            Some("quarantine".to_string())
        );
    }

    #[test]
    fn parse_dmarc_subdomain_policy_missing() {
        let record = "v=DMARC1; p=reject";
        assert_eq!(DmarcCheckerTool::parse_dmarc_subdomain_policy(record), None);
    }

    #[test]
    fn parse_dmarc_rua_present() {
        let record = "v=DMARC1; p=reject; rua=mailto:dmarc@example.com";
        assert_eq!(
            DmarcCheckerTool::parse_dmarc_rua(record),
            Some("mailto:dmarc@example.com".to_string())
        );
    }

    #[test]
    fn parse_dmarc_rua_missing() {
        let record = "v=DMARC1; p=none";
        assert_eq!(DmarcCheckerTool::parse_dmarc_rua(record), None);
    }

    #[test]
    fn is_spf_record_valid() {
        assert!(DmarcCheckerTool::is_spf_record(
            "v=spf1 include:_spf.google.com -all"
        ));
        assert!(DmarcCheckerTool::is_spf_record("v=spf1 -all"));
    }

    #[test]
    fn is_spf_record_invalid() {
        assert!(!DmarcCheckerTool::is_spf_record("v=DMARC1; p=reject"));
        assert!(!DmarcCheckerTool::is_spf_record("some random txt record"));
    }

    #[test]
    fn spf_soft_fail_detection() {
        assert!(DmarcCheckerTool::spf_uses_soft_fail(
            "v=spf1 include:_spf.google.com ~all"
        ));
        assert!(!DmarcCheckerTool::spf_uses_soft_fail(
            "v=spf1 include:_spf.google.com -all"
        ));
    }

    #[test]
    fn spf_allows_all_detection() {
        assert!(DmarcCheckerTool::spf_allows_all("v=spf1 +all"));
        assert!(!DmarcCheckerTool::spf_allows_all("v=spf1 -all"));
        assert!(!DmarcCheckerTool::spf_allows_all("v=spf1 ~all"));
    }
}

impl PentestTool for DmarcCheckerTool {
    fn name(&self) -> &str {
        "dmarc_checker"
    }

    fn description(&self) -> &str {
        "Checks email security configuration for a domain. Queries DMARC and SPF records and \
         evaluates policy strength. Reports missing or weak email authentication settings."
    }

    fn input_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "domain": {
                    "type": "string",
                    "description": "The domain to check (e.g., 'example.com')"
                }
            },
            "required": ["domain"]
        })
    }

    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 domain = input
                .get("domain")
                .and_then(|v| v.as_str())
                .ok_or_else(|| {
                    CoreError::Dast("Missing required 'domain' parameter".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 email_data = json!({});

            // ---- DMARC check ----
            let dmarc_domain = format!("_dmarc.{domain}");
            let dmarc_records = Self::query_txt(&dmarc_domain).await.unwrap_or_default();

            let dmarc_record = dmarc_records.iter().find(|r| r.starts_with("v=DMARC1"));

            match dmarc_record {
                Some(record) => {
                    email_data["dmarc_record"] = json!(record);

                    let policy = Self::parse_dmarc_policy(record);
                    let sp = Self::parse_dmarc_subdomain_policy(record);
                    let rua = Self::parse_dmarc_rua(record);

                    email_data["dmarc_policy"] = json!(policy);
                    email_data["dmarc_subdomain_policy"] = json!(sp);
                    email_data["dmarc_rua"] = json!(rua);

                    // Warn on weak policy
                    if let Some(ref p) = policy {
                        if p == "none" {
                            let evidence = DastEvidence {
                                request_method: "DNS".to_string(),
                                request_url: dmarc_domain.clone(),
                                request_headers: None,
                                request_body: None,
                                response_status: 0,
                                response_headers: None,
                                response_snippet: Some(record.clone()),
                                screenshot_path: None,
                                payload: None,
                                response_time_ms: None,
                            };

                            let mut finding = DastFinding::new(
                            String::new(),
                            target_id.clone(),
                            DastVulnType::EmailSecurity,
                            format!("Weak DMARC policy for {domain}"),
                            format!(
                                "The DMARC policy for {domain} is set to 'none', which only monitors \
                                 but does not enforce email authentication. Attackers can spoof emails \
                                 from this domain."
                            ),
                            Severity::Medium,
                            domain.to_string(),
                            "DNS".to_string(),
                        );
                            finding.cwe = Some("CWE-290".to_string());
                            finding.evidence = vec![evidence];
                            finding.remediation = Some(
                                "Upgrade the DMARC policy from 'p=none' to 'p=quarantine' or \
                             'p=reject' after verifying legitimate email flows are properly \
                             authenticated."
                                    .to_string(),
                            );
                            findings.push(finding);
                            warn!(domain, "DMARC policy is 'none'");
                        }
                    }

                    // Warn if no reporting URI
                    if rua.is_none() {
                        let evidence = DastEvidence {
                            request_method: "DNS".to_string(),
                            request_url: dmarc_domain.clone(),
                            request_headers: None,
                            request_body: None,
                            response_status: 0,
                            response_headers: None,
                            response_snippet: Some(record.clone()),
                            screenshot_path: None,
                            payload: None,
                            response_time_ms: None,
                        };

                        let mut finding = DastFinding::new(
                        String::new(),
                        target_id.clone(),
                        DastVulnType::EmailSecurity,
                        format!("DMARC without reporting for {domain}"),
                        format!(
                            "The DMARC record for {domain} does not include a reporting URI (rua=). \
                             Without reporting, you will not receive aggregate feedback about email \
                             authentication failures."
                        ),
                        Severity::Info,
                        domain.to_string(),
                        "DNS".to_string(),
                    );
                        finding.cwe = Some("CWE-778".to_string());
                        finding.evidence = vec![evidence];
                        finding.remediation = Some(
                            "Add a 'rua=' tag to your DMARC record to receive aggregate reports. \
                         Example: 'rua=mailto:dmarc-reports@example.com'."
                                .to_string(),
                        );
                        findings.push(finding);
                    }
                }
                None => {
                    email_data["dmarc_record"] = json!(null);

                    let evidence = DastEvidence {
                        request_method: "DNS".to_string(),
                        request_url: dmarc_domain.clone(),
                        request_headers: None,
                        request_body: None,
                        response_status: 0,
                        response_headers: None,
                        response_snippet: Some("No DMARC record found".to_string()),
                        screenshot_path: None,
                        payload: None,
                        response_time_ms: None,
                    };

                    let mut finding = DastFinding::new(
                        String::new(),
                        target_id.clone(),
                        DastVulnType::EmailSecurity,
                        format!("Missing DMARC record for {domain}"),
                        format!(
                            "No DMARC record was found for {domain}. Without DMARC, there is no \
                         policy to prevent email spoofing and phishing attacks using this domain."
                        ),
                        Severity::High,
                        domain.to_string(),
                        "DNS".to_string(),
                    );
                    finding.cwe = Some("CWE-290".to_string());
                    finding.evidence = vec![evidence];
                    finding.remediation = Some(
                    "Create a DMARC TXT record at _dmarc.<domain>. Start with 'v=DMARC1; p=none; \
                     rua=mailto:dmarc@example.com' and gradually move to 'p=reject'."
                        .to_string(),
                );
                    findings.push(finding);
                    warn!(domain, "No DMARC record found");
                }
            }

            // ---- SPF check ----
            let txt_records = Self::query_txt(domain).await.unwrap_or_default();
            let spf_record = txt_records.iter().find(|r| Self::is_spf_record(r));

            match spf_record {
                Some(record) => {
                    email_data["spf_record"] = json!(record);

                    if Self::spf_allows_all(record) {
                        let evidence = DastEvidence {
                            request_method: "DNS".to_string(),
                            request_url: domain.to_string(),
                            request_headers: None,
                            request_body: None,
                            response_status: 0,
                            response_headers: None,
                            response_snippet: Some(record.clone()),
                            screenshot_path: None,
                            payload: None,
                            response_time_ms: None,
                        };

                        let mut finding = DastFinding::new(
                        String::new(),
                        target_id.clone(),
                        DastVulnType::EmailSecurity,
                        format!("SPF allows all senders for {domain}"),
                        format!(
                            "The SPF record for {domain} uses '+all' which allows any server to \
                             send email on behalf of this domain, completely negating SPF protection."
                        ),
                        Severity::Critical,
                        domain.to_string(),
                        "DNS".to_string(),
                    );
                        finding.cwe = Some("CWE-290".to_string());
                        finding.evidence = vec![evidence];
                        finding.remediation = Some(
                        "Change '+all' to '-all' (hard fail) or '~all' (soft fail) in your SPF record. \
                         Only list authorized mail servers."
                            .to_string(),
                    );
                        findings.push(finding);
                    } else if Self::spf_uses_soft_fail(record) {
                        let evidence = DastEvidence {
                            request_method: "DNS".to_string(),
                            request_url: domain.to_string(),
                            request_headers: None,
                            request_body: None,
                            response_status: 0,
                            response_headers: None,
                            response_snippet: Some(record.clone()),
                            screenshot_path: None,
                            payload: None,
                            response_time_ms: None,
                        };

                        let mut finding = DastFinding::new(
                            String::new(),
                            target_id.clone(),
                            DastVulnType::EmailSecurity,
                            format!("SPF soft fail for {domain}"),
                            format!(
                                "The SPF record for {domain} uses '~all' (soft fail) instead of \
                             '-all' (hard fail). Soft fail marks unauthorized emails as suspicious \
                             but does not reject them."
                            ),
                            Severity::Low,
                            domain.to_string(),
                            "DNS".to_string(),
                        );
                        finding.cwe = Some("CWE-290".to_string());
                        finding.evidence = vec![evidence];
                        finding.remediation = Some(
                            "Consider changing '~all' to '-all' in your SPF record once you have \
                         confirmed all legitimate mail sources are listed."
                                .to_string(),
                        );
                        findings.push(finding);
                    }
                }
                None => {
                    email_data["spf_record"] = json!(null);

                    let evidence = DastEvidence {
                        request_method: "DNS".to_string(),
                        request_url: domain.to_string(),
                        request_headers: None,
                        request_body: None,
                        response_status: 0,
                        response_headers: None,
                        response_snippet: Some("No SPF record found".to_string()),
                        screenshot_path: None,
                        payload: None,
                        response_time_ms: None,
                    };

                    let mut finding = DastFinding::new(
                        String::new(),
                        target_id.clone(),
                        DastVulnType::EmailSecurity,
                        format!("Missing SPF record for {domain}"),
                        format!(
                        "No SPF record was found for {domain}. Without SPF, any server can claim \
                         to send email on behalf of this domain."
                    ),
                        Severity::High,
                        domain.to_string(),
                        "DNS".to_string(),
                    );
                    finding.cwe = Some("CWE-290".to_string());
                    finding.evidence = vec![evidence];
                    finding.remediation = Some(
                        "Create an SPF TXT record for your domain. Example: \
                     'v=spf1 include:_spf.google.com -all'."
                            .to_string(),
                    );
                    findings.push(finding);
                    warn!(domain, "No SPF record found");
                }
            }

            let count = findings.len();
            info!(domain, findings = count, "DMARC/SPF check complete");

            Ok(PentestToolResult {
                summary: if count > 0 {
                    format!("Found {count} email security issues for {domain}.")
                } else {
                    format!("Email security configuration looks good for {domain}.")
                },
                findings,
                data: email_data,
            })
        })
    }
}