borer-core 0.5.9

use std::fs::File;
use std::io::Write;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

use dashmap::DashSet;
use log::{debug, warn};
use socks5_proto::Address;

const DELETED_PREFIX: &str = "# deleted: ";

#[derive(Debug)]
pub struct AddressList {
    entries: Arc<DashSet<String>>,
    file_path: PathBuf,
    add_lock: Mutex<()>,
}

pub type ProxyList = AddressList;
pub type DirectList = AddressList;

impl AddressList {
    pub fn new<P: AsRef<Path>>(file_path: P) -> Self {
        let file_path = file_path.as_ref().to_path_buf();
        let entries = Arc::new(DashSet::new());

        let list = Self {
            entries,
            file_path,
            add_lock: Mutex::new(()),
        };
        list.load_from_file();
        list
    }

    fn load_from_file(&self) {
        if !self.file_path.exists() {
            debug!(
                "Address list file {:?} does not exist, starting with empty list",
                self.file_path
            );
            return;
        }

        let content = match std::fs::read_to_string(&self.file_path) {
            Ok(c) => c,
            Err(e) => {
                warn!(
                    "Failed to read address list file {:?}: {}",
                    self.file_path, e
                );
                return;
            }
        };

        let mut deleted_entries = std::collections::HashSet::new();
        let mut needs_compact = false;

        for line in content.lines() {
            let trimmed = line.trim();
            if let Some(entry) = trimmed.strip_prefix(DELETED_PREFIX) {
                deleted_entries.insert(entry);
                needs_compact = true;
            }
        }

        for line in content.lines() {
            let trimmed = line.trim();
            if !trimmed.is_empty()
                && !trimmed.starts_with('#')
                && !deleted_entries.contains(trimmed)
            {
                self.entries.insert(trimmed.to_string());
            }
        }

        debug!(
            "Loaded {} entries from {:?}",
            self.entries.len(),
            self.file_path
        );

        if needs_compact {
            let compacted: String = content
                .lines()
                .filter(|line| {
                    let trimmed = line.trim();
                    !trimmed.starts_with(DELETED_PREFIX) && !deleted_entries.contains(trimmed)
                })
                .collect::<Vec<_>>()
                .join("\n");

            if let Err(e) = std::fs::write(&self.file_path, compacted) {
                warn!("Failed to compact address list file: {}", e);
            } else {
                debug!("Compacted address list file {:?}", self.file_path);
            }
        }
    }

    pub fn contains_address(&self, addr: &Address) -> bool {
        match addr {
            Address::DomainAddress(domain, _) => {
                let domain = String::from_utf8_lossy(domain);
                self.contains_domain(&domain)
            }
            Address::SocketAddress(socket_addr) => self.contains_ip(&socket_addr.ip()),
        }
    }

    fn contains_domain(&self, domain: &str) -> bool {
        if self.entries.contains(domain) {
            return true;
        }

        for entry in self.entries.iter() {
            let entry_str = entry.as_str();
            if let Some(base_domain) = entry_str.strip_prefix('.') {
                if domain == base_domain || domain.ends_with(entry_str) {
                    return true;
                }
            }
        }

        let mut current = domain;
        while let Some(idx) = current.find('.') {
            let suffix = &current[idx..];
            if self.entries.contains(suffix) {
                return true;
            }
            current = &current[idx + 1..];
            if current.is_empty() {
                break;
            }
        }

        false
    }

    fn contains_ip(&self, ip: &IpAddr) -> bool {
        let ip_str = ip.to_string();
        if self.entries.contains(&ip_str) {
            return true;
        }

        match ip {
            IpAddr::V4(ipv4) => self.contains_ipv4(ipv4),
            IpAddr::V6(ipv6) => self.contains_ipv6(ipv6),
        }
    }

    fn contains_ipv4(&self, ip: &Ipv4Addr) -> bool {
        let octets = ip.octets();

        for prefix_len in 8..=32 {
            let mask = !0u32 << (32 - prefix_len);
            let network = u32::from_be_bytes(octets) & mask;
            let cidr = format!("{}/{}", Ipv4Addr::from(network.to_be_bytes()), prefix_len);
            if self.entries.contains(&cidr) {
                return true;
            }
        }

        false
    }

    fn contains_ipv6(&self, ip: &Ipv6Addr) -> bool {
        let segments = ip.segments();

        for prefix_len in 8..=128 {
            let full_segments = prefix_len / 16;
            let remaining_bits = prefix_len % 16;

            let mut network_bytes = [0u8; 16];
            for i in 0..full_segments as usize {
                network_bytes[i * 2] = (segments[i] >> 8) as u8;
                network_bytes[i * 2 + 1] = segments[i] as u8;
            }

            if remaining_bits > 0 && full_segments < 8 {
                let mask = !0u16 << (16 - remaining_bits);
                let masked = segments[full_segments as usize] & mask;
                network_bytes[full_segments as usize * 2] = (masked >> 8) as u8;
                network_bytes[full_segments as usize * 2 + 1] = masked as u8;
            }

            let network = Ipv6Addr::from(network_bytes);
            let cidr = format!("{}/{}", network, prefix_len);
            if self.entries.contains(&cidr) {
                return true;
            }
        }

        false
    }

    pub fn add_address(&self, addr: &Address) {
        let _guard = self.add_lock.lock().unwrap();

        let entry = match addr {
            Address::DomainAddress(domain, _) => String::from_utf8_lossy(domain).to_string(),
            Address::SocketAddress(socket_addr) => socket_addr.ip().to_string(),
        };
        if self.entries.contains(&entry) {
            return;
        }

        self.entries.insert(entry.clone());
        self.append_to_file(&entry);
    }

    pub fn remove_address(&self, addr: &Address) {
        let _guard = self.add_lock.lock().unwrap();

        let entry = match addr {
            Address::DomainAddress(domain, _) => String::from_utf8_lossy(domain).to_string(),
            Address::SocketAddress(socket_addr) => socket_addr.ip().to_string(),
        };

        if self.entries.remove(&entry).is_some() {
            self.mark_deleted_in_file(&entry);
            debug!("Removed entry {} from {:?}", entry, self.file_path);
        }
    }

    fn append_to_file(&self, entry: &str) {
        match File::options()
            .create(true)
            .append(true)
            .open(&self.file_path)
        {
            Ok(mut file) => {
                if let Err(e) = writeln!(file, "{}", entry) {
                    warn!("Failed to write entry {} to file: {}", entry, e);
                } else {
                    debug!("Added entry {} to {:?}", entry, self.file_path);
                }
            }
            Err(e) => {
                warn!("Failed to open address list file for appending: {}", e);
            }
        }
    }

    fn mark_deleted_in_file(&self, entry: &str) {
        match File::options()
            .create(true)
            .append(true)
            .open(&self.file_path)
        {
            Ok(mut file) => {
                if let Err(e) = writeln!(file, "{}{}", DELETED_PREFIX, entry) {
                    warn!("Failed to mark entry {} as deleted: {}", entry, e);
                }
            }
            Err(e) => {
                warn!("Failed to open file for marking deletion: {}", e);
            }
        }
    }

    pub fn len(&self) -> usize {
        self.entries.len()
    }

    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use std::net::SocketAddr;
    use tempfile::NamedTempFile;

    #[test]
    fn test_domain_matching_exact() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(temp_file.path(), "example.com\n").unwrap();

        let list = AddressList::new(temp_file.path());

        assert!(list.contains_domain("example.com"));
        assert!(!list.contains_domain("other.com"));
    }

    #[test]
    fn test_domain_matching_wildcard() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(temp_file.path(), ".google.com\n").unwrap();

        let list = AddressList::new(temp_file.path());

        assert!(list.contains_domain("www.google.com"));
        assert!(list.contains_domain("mail.google.com"));
        assert!(list.contains_domain(".google.com"));
        assert!(list.contains_domain("google.com"));
        assert!(list.contains_domain("a.b.google.com"));
        assert!(!list.contains_domain("example.com"));
        assert!(!list.contains_domain("google.com.evil.com"));
    }

    #[test]
    fn test_domain_matching_wildcard_multi_level() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(temp_file.path(), ".xx.com\n").unwrap();

        let list = AddressList::new(temp_file.path());

        assert!(list.contains_domain("xx.com"));
        assert!(list.contains_domain("a.xx.com"));
        assert!(list.contains_domain("b.xx.com"));
        assert!(list.contains_domain("c.d.xx.com"));
        assert!(!list.contains_domain("xx.com.evil.com"));
    }

    #[test]
    fn test_ip_matching_exact() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(temp_file.path(), "192.168.1.1\n").unwrap();

        let list = AddressList::new(temp_file.path());

        assert!(list.contains_ip(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1))));
        assert!(!list.contains_ip(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 2))));
    }

    #[test]
    fn test_ip_matching_cidr() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(temp_file.path(), "10.0.0.0/8\n").unwrap();

        let list = AddressList::new(temp_file.path());

        assert!(list.contains_ip(&IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1))));
        assert!(list.contains_ip(&IpAddr::V4(Ipv4Addr::new(10, 255, 255, 255))));
        assert!(!list.contains_ip(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1))));
    }

    #[test]
    fn test_ip_matching_cidr_24() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(temp_file.path(), "192.168.1.0/24\n").unwrap();

        let list = AddressList::new(temp_file.path());

        assert!(list.contains_ip(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 0))));
        assert!(list.contains_ip(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 255))));
        assert!(!list.contains_ip(&IpAddr::V4(Ipv4Addr::new(192, 168, 2, 1))));
    }

    #[test]
    fn test_address_matching_domain() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(temp_file.path(), "example.com\n").unwrap();

        let list = AddressList::new(temp_file.path());

        assert!(list.contains_address(&Address::DomainAddress(b"example.com".to_vec(), 443)));
        assert!(!list.contains_address(&Address::DomainAddress(b"other.com".to_vec(), 443)));
    }

    #[test]
    fn test_address_matching_ip() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(temp_file.path(), "192.168.1.0/24\n").unwrap();

        let list = AddressList::new(temp_file.path());

        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 100)), 443);
        assert!(list.contains_address(&Address::SocketAddress(addr)));

        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 168, 2, 100)), 443);
        assert!(!list.contains_address(&Address::SocketAddress(addr)));
    }

    #[test]
    fn test_add_domain() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(temp_file.path(), "existing.com\n").unwrap();

        let list = AddressList::new(temp_file.path());
        assert!(list.contains_domain("existing.com"));
        assert_eq!(list.len(), 1);

        list.add_address(&Address::DomainAddress(b"new.com".to_vec(), 443));

        assert!(list.contains_domain("new.com"));
        assert_eq!(list.len(), 2);

        let contents = fs::read_to_string(temp_file.path()).unwrap();
        assert!(contents.contains("new.com"));
    }

    #[test]
    fn test_add_ip() {
        let temp_file = NamedTempFile::new().unwrap();

        let list = AddressList::new(temp_file.path());
        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), 443);
        list.add_address(&Address::SocketAddress(addr));

        assert!(list.contains_ip(&IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1))));

        let contents = fs::read_to_string(temp_file.path()).unwrap();
        assert!(contents.contains("10.0.0.1"));
    }

    #[test]
    fn test_remove_domain() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(temp_file.path(), "example.com\nother.com\n").unwrap();

        let list = AddressList::new(temp_file.path());
        assert!(list.contains_domain("example.com"));
        assert_eq!(list.len(), 2);

        list.remove_address(&Address::DomainAddress(b"example.com".to_vec(), 443));

        assert!(!list.contains_domain("example.com"));
        assert!(list.contains_domain("other.com"));
        assert_eq!(list.len(), 1);

        let contents = fs::read_to_string(temp_file.path()).unwrap();
        assert!(contents.contains("# deleted: example.com"));
        assert!(contents.contains("other.com"));
    }

    #[test]
    fn test_remove_ip() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(temp_file.path(), "10.0.0.1\n10.0.0.2\n").unwrap();

        let list = AddressList::new(temp_file.path());
        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), 443);
        list.remove_address(&Address::SocketAddress(addr));

        assert!(!list.contains_ip(&IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1))));
        assert!(list.contains_ip(&IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2))));

        let contents = fs::read_to_string(temp_file.path()).unwrap();
        assert!(contents.contains("# deleted: 10.0.0.1"));
        assert!(contents.contains("10.0.0.2"));
    }

    #[test]
    fn test_remove_nonexistent() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(temp_file.path(), "existing.com\n").unwrap();

        let list = AddressList::new(temp_file.path());
        list.remove_address(&Address::DomainAddress(b"nonexistent.com".to_vec(), 443));

        assert_eq!(list.len(), 1);
        assert!(list.contains_domain("existing.com"));
    }

    #[test]
    fn test_empty_file() {
        let temp_file = NamedTempFile::new().unwrap();
        let list = AddressList::new(temp_file.path());

        assert!(list.is_empty());
        assert!(!list.contains_domain("any.com"));
    }

    #[test]
    fn test_nonexistent_file() {
        let list = AddressList::new("/nonexistent/path/list.txt");
        assert!(list.is_empty());
    }

    #[test]
    fn test_comments_and_whitespace() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(
            temp_file.path(),
            "# comment\n  trimmed.com  \n\n# another comment\n",
        )
        .unwrap();

        let list = AddressList::new(temp_file.path());

        assert!(list.contains_domain("trimmed.com"));
        assert_eq!(list.len(), 1);
    }

    #[test]
    fn test_type_aliases() {
        let temp_file = NamedTempFile::new().unwrap();

        let proxy_list: ProxyList = ProxyList::new(temp_file.path());
        assert!(proxy_list.is_empty());

        let direct_list: DirectList = DirectList::new(temp_file.path());
        assert!(direct_list.is_empty());
    }

    #[test]
    fn test_mark_deleted_in_file() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(temp_file.path(), "example.com\nother.com\n").unwrap();

        let list = AddressList::new(temp_file.path());
        list.remove_address(&Address::DomainAddress(b"example.com".to_vec(), 443));

        let contents = fs::read_to_string(temp_file.path()).unwrap();
        assert!(contents.contains("# deleted: example.com"));
        assert!(contents.contains("other.com"));
        assert!(contents.contains("example.com"));
    }

    #[test]
    fn test_compact_on_load() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(
            temp_file.path(),
            "removed.com\nactive.com\nanother.com\n# deleted: removed.com\n# deleted: another.com\n",
        )
        .unwrap();

        let list = AddressList::new(temp_file.path());

        assert!(list.contains_domain("active.com"));
        assert!(!list.contains_domain("removed.com"));
        assert!(!list.contains_domain("another.com"));
        assert_eq!(list.len(), 1);

        let contents = fs::read_to_string(temp_file.path()).unwrap();
        assert!(!contents.contains("# deleted:"));
        assert!(!contents.contains("removed.com"));
        assert!(!contents.contains("another.com"));
        assert!(contents.contains("active.com"));
    }

    #[test]
    fn test_deleted_entry_not_loaded() {
        let temp_file = NamedTempFile::new().unwrap();
        fs::write(
            temp_file.path(),
            "example.com\nother.com\n# deleted: example.com\n",
        )
        .unwrap();

        let list = AddressList::new(temp_file.path());

        assert!(!list.contains_domain("example.com"));
        assert!(list.contains_domain("other.com"));
        assert_eq!(list.len(), 1);
    }
}