smartmta/rust/crates/mailer-security/src/ip_reputation.rs

use hickory_resolver::TokioResolver;
use serde::{Deserialize, Serialize};
use std::net::{IpAddr, Ipv4Addr};

use crate::error::Result;

/// Default DNSBL servers to check, same as the TypeScript IPReputationChecker.
pub const DEFAULT_DNSBL_SERVERS: &[&str] = &[
    "zen.spamhaus.org",
    "bl.spamcop.net",
    "b.barracudacentral.org",
    "spam.dnsbl.sorbs.net",
    "dnsbl.sorbs.net",
    "cbl.abuseat.org",
    "xbl.spamhaus.org",
    "pbl.spamhaus.org",
    "dnsbl-1.uceprotect.net",
    "psbl.surriel.com",
];

/// Result of a DNSBL check.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DnsblResult {
    /// IP address that was checked.
    pub ip: String,
    /// Number of DNSBL servers that list this IP.
    pub listed_count: usize,
    /// Names of DNSBL servers that list this IP.
    pub listed_on: Vec<String>,
    /// Total number of DNSBL servers checked.
    pub total_checked: usize,
}

/// Result of a full IP reputation check.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReputationResult {
    /// Reputation score: 0 (worst) to 100 (best).
    pub score: u8,
    /// Whether the IP is considered spam source.
    pub is_spam: bool,
    /// IP address that was checked.
    pub ip: String,
    /// DNSBL results.
    pub dnsbl: DnsblResult,
    /// Heuristic IP type classification.
    pub ip_type: IpType,
}

/// Heuristic IP type classification.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum IpType {
    Residential,
    Datacenter,
    Proxy,
    Tor,
    Vpn,
    Unknown,
}

/// Risk level based on reputation score.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RiskLevel {
    /// Score < 20
    High,
    /// Score 20-49
    Medium,
    /// Score 50-79
    Low,
    /// Score >= 80
    Trusted,
}

/// Get the risk level for a reputation score.
pub fn risk_level(score: u8) -> RiskLevel {
    match score {
        0..=19 => RiskLevel::High,
        20..=49 => RiskLevel::Medium,
        50..=79 => RiskLevel::Low,
        _ => RiskLevel::Trusted,
    }
}

/// Check an IP against DNSBL servers.
///
/// * `ip` - The IP address to check (must be IPv4)
/// * `dnsbl_servers` - DNSBL servers to query (use `DEFAULT_DNSBL_SERVERS` for defaults)
/// * `resolver` - DNS resolver to use
pub async fn check_dnsbl(
    ip: IpAddr,
    dnsbl_servers: &[&str],
    resolver: &TokioResolver,
) -> Result<DnsblResult> {
    let ipv4 = match ip {
        IpAddr::V4(v4) => v4,
        IpAddr::V6(_) => {
            // IPv6 DNSBL is less common; return clean result
            return Ok(DnsblResult {
                ip: ip.to_string(),
                listed_count: 0,
                listed_on: Vec::new(),
                total_checked: 0,
            });
        }
    };

    let reversed = reverse_ipv4(ipv4);
    let total = dnsbl_servers.len();

    // Query all DNSBL servers in parallel
    let mut handles = Vec::with_capacity(total);
    for &server in dnsbl_servers {
        let query = format!("{}.{}", reversed, server);
        let resolver = resolver.clone();
        let server_name = server.to_string();
        handles.push(tokio::spawn(async move {
            match resolver.lookup_ip(&query).await {
                Ok(_) => Some(server_name), // IP is listed
                Err(_) => None,             // IP is not listed (NXDOMAIN)
            }
        }));
    }

    let mut listed_on = Vec::new();
    for handle in handles {
        match handle.await {
            Ok(Some(server)) => listed_on.push(server),
            _ => {}
        }
    }

    Ok(DnsblResult {
        ip: ip.to_string(),
        listed_count: listed_on.len(),
        listed_on,
        total_checked: total,
    })
}

/// Full IP reputation check: DNSBL + heuristic classification + scoring.
pub async fn check_reputation(
    ip: IpAddr,
    dnsbl_servers: &[&str],
    resolver: &TokioResolver,
) -> Result<ReputationResult> {
    let dnsbl = check_dnsbl(ip, dnsbl_servers, resolver).await?;
    let ip_type = classify_ip(ip);

    // Scoring: start at 100
    let mut score: i16 = 100;

    // Subtract 10 per DNSBL listing
    score -= (dnsbl.listed_count as i16) * 10;

    // Subtract 30 for suspicious IP types
    match ip_type {
        IpType::Proxy | IpType::Tor | IpType::Vpn => {
            score -= 30;
        }
        _ => {}
    }

    let score = score.clamp(0, 100) as u8;
    let is_spam = score < 50;

    Ok(ReputationResult {
        score,
        is_spam,
        ip: ip.to_string(),
        dnsbl,
        ip_type,
    })
}

/// Reverse IPv4 octets for DNSBL queries: "1.2.3.4" -> "4.3.2.1".
fn reverse_ipv4(ip: Ipv4Addr) -> String {
    let octets = ip.octets();
    format!("{}.{}.{}.{}", octets[3], octets[2], octets[1], octets[0])
}

/// Heuristic IP type classification based on well-known prefix ranges.
/// Same heuristics as the TypeScript IPReputationChecker.
fn classify_ip(ip: IpAddr) -> IpType {
    let ip_str = ip.to_string();

    // Known Tor exit node prefixes
    if ip_str.starts_with("171.25.")
        || ip_str.starts_with("185.220.")
        || ip_str.starts_with("95.216.")
    {
        return IpType::Tor;
    }

    // Known VPN provider prefixes
    if ip_str.starts_with("185.156.") || ip_str.starts_with("37.120.") {
        return IpType::Vpn;
    }

    // Known proxy prefixes
    if ip_str.starts_with("34.92.") || ip_str.starts_with("34.206.") {
        return IpType::Proxy;
    }

    // Major cloud provider prefixes (datacenter)
    if ip_str.starts_with("13.")
        || ip_str.starts_with("35.")
        || ip_str.starts_with("52.")
        || ip_str.starts_with("34.")
        || ip_str.starts_with("104.")
    {
        return IpType::Datacenter;
    }

    IpType::Residential
}

/// Validate an IPv4 address string.
pub fn is_valid_ipv4(ip: &str) -> bool {
    ip.parse::<Ipv4Addr>().is_ok()
}

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

    #[test]
    fn test_reverse_ipv4() {
        let ip: Ipv4Addr = "1.2.3.4".parse().unwrap();
        assert_eq!(reverse_ipv4(ip), "4.3.2.1");

        let ip: Ipv4Addr = "192.168.1.100".parse().unwrap();
        assert_eq!(reverse_ipv4(ip), "100.1.168.192");
    }

    #[test]
    fn test_classify_ip() {
        assert_eq!(
            classify_ip("171.25.193.20".parse().unwrap()),
            IpType::Tor
        );
        assert_eq!(
            classify_ip("185.156.73.1".parse().unwrap()),
            IpType::Vpn
        );
        assert_eq!(
            classify_ip("34.92.1.1".parse().unwrap()),
            IpType::Proxy
        );
        assert_eq!(
            classify_ip("52.0.0.1".parse().unwrap()),
            IpType::Datacenter
        );
        assert_eq!(
            classify_ip("203.0.113.1".parse().unwrap()),
            IpType::Residential
        );
    }

    #[test]
    fn test_risk_level() {
        assert_eq!(risk_level(10), RiskLevel::High);
        assert_eq!(risk_level(30), RiskLevel::Medium);
        assert_eq!(risk_level(60), RiskLevel::Low);
        assert_eq!(risk_level(90), RiskLevel::Trusted);
    }

    #[test]
    fn test_is_valid_ipv4() {
        assert!(is_valid_ipv4("1.2.3.4"));
        assert!(is_valid_ipv4("255.255.255.255"));
        assert!(!is_valid_ipv4("999.999.999.999"));
        assert!(!is_valid_ipv4("not-an-ip"));
    }

    #[test]
    fn test_default_dnsbl_servers() {
        assert_eq!(DEFAULT_DNSBL_SERVERS.len(), 10);
        assert!(DEFAULT_DNSBL_SERVERS.contains(&"zen.spamhaus.org"));
    }
}