compliance-scanner-agent/compliance-dast/src/tools/dns_checker.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 tokio::net::lookup_host;
use tracing::{info, warn};

/// Tool that checks DNS configuration for security issues.
///
/// Resolves A, AAAA, MX, TXT, CNAME, NS records using the system resolver
/// via `tokio::net::lookup_host` and `std::net::ToSocketAddrs`. For TXT-based
/// records (SPF, DMARC, CAA, DNSSEC) it uses a simple TXT query via the
/// `tokio::process::Command` wrapper around `dig` where available.
pub struct DnsCheckerTool;

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

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

    /// Run a `dig` query and return the answer lines.
    async fn dig_query(domain: &str, record_type: &str) -> Result<Vec<String>, CoreError> {
        let output = tokio::process::Command::new("dig")
            .args(["+short", record_type, domain])
            .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().to_string())
            .filter(|l| !l.is_empty())
            .collect();
        Ok(lines)
    }

    /// Resolve A/AAAA records using tokio lookup.
    async fn resolve_addresses(domain: &str) -> Result<(Vec<String>, Vec<String>), CoreError> {
        let mut ipv4 = Vec::new();
        let mut ipv6 = Vec::new();

        let addr_str = format!("{domain}:443");
        match lookup_host(&addr_str).await {
            Ok(addrs) => {
                for addr in addrs {
                    match addr {
                        std::net::SocketAddr::V4(v4) => ipv4.push(v4.ip().to_string()),
                        std::net::SocketAddr::V6(v6) => ipv6.push(v6.ip().to_string()),
                    }
                }
            }
            Err(e) => {
                return Err(CoreError::Dast(format!(
                    "DNS resolution failed for {domain}: {e}"
                )));
            }
        }

        Ok((ipv4, ipv6))
    }
}

impl PentestTool for DnsCheckerTool {
    fn name(&self) -> &str {
        "dns_checker"
    }

    fn description(&self) -> &str {
        "Checks DNS configuration for a domain. Resolves A, AAAA, MX, TXT, CNAME, NS records. \
         Checks for DNSSEC, CAA records, and potential subdomain takeover via dangling CNAME/NS."
    }

    fn input_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "domain": {
                    "type": "string",
                    "description": "The domain to check (e.g., 'example.com')"
                },
                "subdomains": {
                    "type": "array",
                    "description": "Optional list of subdomains to also check (e.g., ['www', 'api', 'mail'])",
                    "items": { "type": "string" }
                }
            },
            "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 subdomains: Vec<String> = input
                .get("subdomains")
                .and_then(|v| v.as_array())
                .map(|arr| {
                    arr.iter()
                        .filter_map(|v| v.as_str().map(String::from))
                        .collect()
                })
                .unwrap_or_default();

            let mut findings = Vec::new();
            let mut dns_data: HashMap<String, serde_json::Value> = HashMap::new();

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

            // --- A / AAAA records ---
            match Self::resolve_addresses(domain).await {
                Ok((ipv4, ipv6)) => {
                    dns_data.insert("a_records".to_string(), json!(ipv4));
                    dns_data.insert("aaaa_records".to_string(), json!(ipv6));
                }
                Err(e) => {
                    dns_data.insert("a_records_error".to_string(), json!(e.to_string()));
                }
            }

            // --- MX records ---
            match Self::dig_query(domain, "MX").await {
                Ok(mx) => {
                    dns_data.insert("mx_records".to_string(), json!(mx));
                }
                Err(e) => {
                    dns_data.insert("mx_records_error".to_string(), json!(e.to_string()));
                }
            }

            // --- NS records ---
            let ns_records = match Self::dig_query(domain, "NS").await {
                Ok(ns) => {
                    dns_data.insert("ns_records".to_string(), json!(ns));
                    ns
                }
                Err(e) => {
                    dns_data.insert("ns_records_error".to_string(), json!(e.to_string()));
                    Vec::new()
                }
            };

            // --- TXT records ---
            match Self::dig_query(domain, "TXT").await {
                Ok(txt) => {
                    dns_data.insert("txt_records".to_string(), json!(txt));
                }
                Err(e) => {
                    dns_data.insert("txt_records_error".to_string(), json!(e.to_string()));
                }
            }

            // --- CNAME records (for subdomains) ---
            let mut cname_data: HashMap<String, Vec<String>> = HashMap::new();
            let mut domains_to_check = vec![domain.to_string()];
            for sub in &subdomains {
                domains_to_check.push(format!("{sub}.{domain}"));
            }

            for fqdn in &domains_to_check {
                match Self::dig_query(fqdn, "CNAME").await {
                    Ok(cnames) if !cnames.is_empty() => {
                        // Check for dangling CNAME
                        for cname in &cnames {
                            let cname_clean = cname.trim_end_matches('.');
                            let check_addr = format!("{cname_clean}:443");
                            let is_dangling = lookup_host(&check_addr).await.is_err();
                            if is_dangling {
                                let evidence = DastEvidence {
                                    request_method: "DNS".to_string(),
                                    request_url: fqdn.clone(),
                                    request_headers: None,
                                    request_body: None,
                                    response_status: 0,
                                    response_headers: None,
                                    response_snippet: Some(format!(
                                        "CNAME {fqdn} -> {cname} (target does not resolve)"
                                    )),
                                    screenshot_path: None,
                                    payload: None,
                                    response_time_ms: None,
                                };

                                let mut finding = DastFinding::new(
                                String::new(),
                                target_id.clone(),
                                DastVulnType::DnsMisconfiguration,
                                format!("Dangling CNAME on {fqdn}"),
                                format!(
                                    "The subdomain {fqdn} has a CNAME record pointing to {cname} which does not resolve. \
                                     This may allow subdomain takeover if an attacker can claim the target hostname."
                                ),
                                Severity::High,
                                fqdn.clone(),
                                "DNS".to_string(),
                            );
                                finding.cwe = Some("CWE-923".to_string());
                                finding.evidence = vec![evidence];
                                finding.remediation = Some(
                                "Remove dangling CNAME records or ensure the target hostname is \
                                 properly configured and resolvable."
                                    .to_string(),
                            );
                                findings.push(finding);
                                warn!(
                                    fqdn,
                                    cname, "Dangling CNAME detected - potential subdomain takeover"
                                );
                            }
                        }
                        cname_data.insert(fqdn.clone(), cnames);
                    }
                    _ => {}
                }
            }
            if !cname_data.is_empty() {
                dns_data.insert("cname_records".to_string(), json!(cname_data));
            }

            // --- CAA records ---
            match Self::dig_query(domain, "CAA").await {
                Ok(caa) => {
                    if caa.is_empty() {
                        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 CAA records found".to_string()),
                            screenshot_path: None,
                            payload: None,
                            response_time_ms: None,
                        };

                        let mut finding = DastFinding::new(
                        String::new(),
                        target_id.clone(),
                        DastVulnType::DnsMisconfiguration,
                        format!("Missing CAA records for {domain}"),
                        format!(
                            "No CAA (Certificate Authority Authorization) records are set for {domain}. \
                             Without CAA records, any certificate authority can issue certificates for this domain."
                        ),
                        Severity::Low,
                        domain.to_string(),
                        "DNS".to_string(),
                    );
                        finding.cwe = Some("CWE-295".to_string());
                        finding.evidence = vec![evidence];
                        finding.remediation = Some(
                        "Add CAA DNS records to restrict which certificate authorities can issue \
                         certificates for your domain. Example: '0 issue \"letsencrypt.org\"'."
                            .to_string(),
                    );
                        findings.push(finding);
                    }
                    dns_data.insert("caa_records".to_string(), json!(caa));
                }
                Err(e) => {
                    dns_data.insert("caa_records_error".to_string(), json!(e.to_string()));
                }
            }

            // --- DNSSEC check ---
            let dnssec_output = tokio::process::Command::new("dig")
                .args(["+dnssec", "+short", "DNSKEY", domain])
                .output()
                .await;

            match dnssec_output {
                Ok(output) => {
                    let stdout = String::from_utf8_lossy(&output.stdout);
                    let has_dnssec = !stdout.trim().is_empty();
                    dns_data.insert("dnssec_enabled".to_string(), json!(has_dnssec));

                    if !has_dnssec {
                        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 DNSKEY records found - DNSSEC not enabled".to_string(),
                            ),
                            screenshot_path: None,
                            payload: None,
                            response_time_ms: None,
                        };

                        let mut finding = DastFinding::new(
                            String::new(),
                            target_id.clone(),
                            DastVulnType::DnsMisconfiguration,
                            format!("DNSSEC not enabled for {domain}"),
                            format!(
                            "DNSSEC is not enabled for {domain}. Without DNSSEC, DNS responses \
                             can be spoofed, allowing man-in-the-middle attacks."
                        ),
                            Severity::Medium,
                            domain.to_string(),
                            "DNS".to_string(),
                        );
                        finding.cwe = Some("CWE-350".to_string());
                        finding.evidence = vec![evidence];
                        finding.remediation = Some(
                            "Enable DNSSEC for your domain by configuring DNSKEY and DS records \
                         with your DNS provider and domain registrar."
                                .to_string(),
                        );
                        findings.push(finding);
                    }
                }
                Err(_) => {
                    dns_data.insert("dnssec_check_error".to_string(), json!("dig not available"));
                }
            }

            // --- Check NS records for dangling ---
            for ns in &ns_records {
                let ns_clean = ns.trim_end_matches('.');
                let check_addr = format!("{ns_clean}:53");
                if lookup_host(&check_addr).await.is_err() {
                    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(format!("NS record {ns} does not resolve")),
                        screenshot_path: None,
                        payload: None,
                        response_time_ms: None,
                    };

                    let mut finding = DastFinding::new(
                    String::new(),
                    target_id.clone(),
                    DastVulnType::DnsMisconfiguration,
                    format!("Dangling NS record for {domain}"),
                    format!(
                        "The NS record {ns} for {domain} does not resolve. \
                         This could allow domain takeover if an attacker can claim the nameserver hostname."
                    ),
                    Severity::Critical,
                    domain.to_string(),
                    "DNS".to_string(),
                );
                    finding.cwe = Some("CWE-923".to_string());
                    finding.evidence = vec![evidence];
                    finding.remediation = Some(
                        "Remove dangling NS records or ensure the nameserver hostname is properly \
                     configured. Dangling NS records can lead to full domain takeover."
                            .to_string(),
                    );
                    findings.push(finding);
                    warn!(
                        domain,
                        ns, "Dangling NS record detected - potential domain takeover"
                    );
                }
            }

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

            Ok(PentestToolResult {
                summary: if count > 0 {
                    format!("Found {count} DNS configuration issues for {domain}.")
                } else {
                    format!("No DNS configuration issues found for {domain}.")
                },
                findings,
                data: json!(dns_data),
            })
        })
    }
}