//! 邮件投递安全：轻量 SPF 校验、DKIM 头检测与 DMARC 策略发现。

use hickory_resolver::{config::ResolverConfig, config::ResolverOpts, TokioAsyncResolver};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::sync::OnceLock;

static RESOLVER: OnceLock<TokioAsyncResolver> = std::sync::OnceLock::new();

/// 共享 DNS 解析器（SPF / PTR 复用）。
pub fn shared_resolver() -> &'static TokioAsyncResolver {
    RESOLVER.get_or_init(|| {
        TokioAsyncResolver::tokio(ResolverConfig::default(), ResolverOpts::default())
    })
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SpfResult {
    Pass,
    Fail,
    SoftFail,
    Neutral,
    None,
}

impl SpfResult {
    pub fn as_str(&self) -> &'static str {
        match self {
            SpfResult::Pass => "pass",
            SpfResult::Fail => "fail",
            SpfResult::SoftFail => "softfail",
            SpfResult::Neutral => "neutral",
            SpfResult::None => "none",
        }
    }
}

/// SPF 校验（简化版：仅 ip4/ip6/all 机制）。
pub async fn check_spf(ip: IpAddr, helo_or_domain: &str, mail_from: &str) -> SpfResult {
    let domain = extract_domain(mail_from).unwrap_or(helo_or_domain);
    if domain.is_empty() {
        return SpfResult::None;
    }
    let resolver = shared_resolver();
    let txt = match resolver.txt_lookup(domain).await {
        Ok(t) => t,
        Err(_) => return SpfResult::None,
    };
    // 拼接所有 TXT 记录文本，定位 v=spf1 记录
    let mut record_str: Option<String> = None;
    for r in txt.into_iter() {
        let s = r.to_string();
        if s.contains("v=spf1") {
            record_str = Some(s);
            break;
        }
    }
    let record = match record_str {
        Some(r) => r,
        None => return SpfResult::None,
    };
    let cleaned = record.replace('"', " ");
    let idx = match cleaned.find("v=spf1") {
        Some(i) => i,
        None => return SpfResult::None,
    };
    let spf_part = &cleaned[idx..];
    for term in spf_part.split_whitespace().skip(1) {
        let (qual, mech) = split_qual(term);
        let matched = match mech {
            "all" => true,
            m if m.starts_with("ip4:") => match_ip4(ip, &m[4..]),
            m if m.starts_with("ip6:") => match_ip6(ip, &m[4..]),
            _ => false, // a/mx/include/redirect 尚未解析
        };
        if matched {
            return match qual {
                '+' => SpfResult::Pass,
                '-' => SpfResult::Fail,
                '~' => SpfResult::SoftFail,
                '?' => SpfResult::Neutral,
                _ => SpfResult::Pass,
            };
        }
    }
    SpfResult::Neutral
}

/// 检测 DKIM-Signature 头是否存在（轻量；不验证签名本身，真实验签留 mail-auth）。
pub fn detect_dkim_signature(raw: &[u8]) -> bool {
    let Ok(s) = std::str::from_utf8(raw) else {
        return false;
    };
    s.lines().take_while(|l| !l.is_empty()).any(|l| {
        l.trim_start()
            .to_ascii_lowercase()
            .starts_with("dkim-signature:")
    })
}

/// 查询发件域 DMARC 策略（只表示策略存在，不表示 DMARC 验证通过）。
pub async fn query_dmarc(resolver: &TokioAsyncResolver, domain: &str) -> &'static str {
    let dmarc_domain = format!("_dmarc.{domain}");
    match resolver.txt_lookup(&dmarc_domain).await {
        Ok(t) => {
            if t.into_iter().next().is_some() {
                "policy_present"
            } else {
                "none"
            }
        }
        Err(_) => "none",
    }
}

fn extract_domain(addr: &str) -> Option<&str> {
    let a = addr.trim_matches(|c| c == '<' || c == '>' || c == ' ');
    a.rsplit_once('@').map(|(_, d)| d)
}

fn split_qual(term: &str) -> (char, &str) {
    match term.chars().next() {
        Some(c @ ('+' | '-' | '~' | '?')) => (c, &term[1..]),
        _ => ('+', term),
    }
}

fn match_ip4(ip: IpAddr, spec: &str) -> bool {
    let ip = match ip {
        IpAddr::V4(v) => v,
        _ => return false,
    };
    let (addr_str, prefix) = match spec.split_once('/') {
        Some((a, p)) => (a, p.parse::<u8>().unwrap_or(32)),
        None => (spec, 32),
    };
    let net: Ipv4Addr = match addr_str.parse() {
        Ok(a) => a,
        Err(_) => return false,
    };
    if prefix == 0 {
        return true;
    }
    let mask: u32 = if prefix >= 32 {
        !0u32
    } else {
        (!0u32) << (32 - prefix)
    };
    (u32::from(ip) & mask) == (u32::from(net) & mask)
}

fn match_ip6(ip: IpAddr, spec: &str) -> bool {
    let ip = match ip {
        IpAddr::V6(v) => v,
        _ => return false,
    };
    let (addr_str, prefix) = match spec.split_once('/') {
        Some((a, p)) => (a, p.parse::<u8>().unwrap_or(128)),
        None => (spec, 128),
    };
    let net: Ipv6Addr = match addr_str.parse() {
        Ok(a) => a,
        Err(_) => return false,
    };
    if prefix == 0 {
        return true;
    }
    let mask: u128 = if prefix >= 128 {
        !0u128
    } else {
        (!0u128) << (128 - prefix)
    };
    (u128::from(ip) & mask) == (u128::from(net) & mask)
}

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

    #[test]
    fn ip4_cidr_matching() {
        let ip = IpAddr::V4(Ipv4Addr::from_str("192.168.1.5").unwrap());
        assert!(match_ip4(ip, "192.168.1.0/24"));
        assert!(match_ip4(ip, "192.168.1.5"));
        assert!(!match_ip4(ip, "10.0.0.0/8"));
        assert!(match_ip4(ip, "0.0.0.0/0"));
    }

    #[test]
    fn qualifier_parsing() {
        assert_eq!(split_qual("-all"), ('-', "all"));
        assert_eq!(split_qual("~all"), ('~', "all"));
        assert_eq!(split_qual("ip4:1.2.3.4"), ('+', "ip4:1.2.3.4"));
    }
}
