proxychains-masq 0.1.5

use std::{
    collections::HashMap,
    io::Write,
    net::{IpAddr, Ipv4Addr, Ipv6Addr},
    sync::{
        atomic::{AtomicUsize, Ordering},
        RwLock,
    },
    time::Duration,
};

use anyhow::{bail, Context, Result};
use rand::{rngs::OsRng, seq::SliceRandom};
use tokio::{net::TcpStream, time::timeout};

use crate::proxy::{http, https, raw, socks4, socks5, BoxStream, Target};

// ─── Public types re-exported from config ─────────────────────────────────────

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ChainType {
    Strict,
    Dynamic,
    Random,
    RoundRobin,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProxyType {
    Socks4,
    Socks5,
    Http,
    /// HTTP CONNECT over TLS.  Whether to verify the proxy certificate is
    /// controlled by [`ChainConfig::tls_skip_verify`].
    Https,
    Raw,
}

/// A single proxy in the chain.
#[derive(Debug, Clone)]
pub struct ProxyEntry {
    pub proxy_type: ProxyType,
    pub addr: IpAddr,
    pub port: u16,
    pub username: Option<String>,
    pub password: Option<String>,
}

/// A localnet exclusion rule.
#[derive(Debug, Clone)]
pub struct LocalNet {
    pub addr: IpAddr,
    pub mask_v4: Option<Ipv4Addr>,
    pub prefix_v6: Option<u8>,
    pub port: Option<u16>,
}

/// A DNAT rewrite rule.
#[derive(Debug, Clone)]
pub struct DnatRule {
    pub orig_addr: Ipv4Addr,
    pub orig_port: Option<u16>,
    pub new_addr: Ipv4Addr,
    pub new_port: Option<u16>,
}

/// Configuration for [`ChainEngine`].
#[derive(Debug, Clone)]
pub struct ChainConfig {
    pub proxies: Vec<ProxyEntry>,
    pub chain_type: ChainType,
    /// Number of proxies to use per connection (for Random / RoundRobin).
    pub chain_len: usize,
    /// How many times to retry with a new chain selection after total failure.
    ///
    /// Only applies to `Random` and `RoundRobin` chains, which select a subset
    /// of proxies each attempt.  `Strict` and `Dynamic` already try every proxy
    /// in the list, so retrying would repeat the same sequence.
    pub chain_retries: usize,
    pub connect_timeout: Duration,
    pub localnets: Vec<LocalNet>,
    pub dnats: Vec<DnatRule>,
    /// When `true`, TLS certificate validation is skipped for all HTTPS proxies.
    ///
    /// Controlled by the `--tls-skip-verify` CLI flag.  Defaults to `false`.
    pub tls_skip_verify: bool,
    /// Number of failures before a proxy is excluded from chain selection.
    ///
    /// A proxy that fails (TCP connect or protocol handshake) this many times
    /// across all connections is marked dead and skipped.  Set to `0` to
    /// disable the dead-proxy filter entirely.  Default: 3.
    pub proxy_dead_threshold: usize,
}

impl Default for ChainConfig {
    fn default() -> Self {
        ChainConfig {
            proxies: Vec::new(),
            chain_type: ChainType::Dynamic,
            chain_len: 1,
            chain_retries: 3,
            connect_timeout: Duration::from_secs(10),
            localnets: Vec::new(),
            dnats: Vec::new(),
            tls_skip_verify: false,
            proxy_dead_threshold: 3,
        }
    }
}

// ─── ChainEngine ─────────────────────────────────────────────────────────────

/// Routes outbound TCP connections through a configurable chain of proxies.
pub struct ChainEngine {
    config: ChainConfig,
    /// Shared offset for round-robin mode.
    rr_offset: AtomicUsize,
    /// Cumulative failure counts keyed by `(addr, port)`.  Once a proxy
    /// reaches `config.proxy_dead_threshold`, it is excluded from selection.
    failure_counts: RwLock<HashMap<(IpAddr, u16), usize>>,
}

impl ChainEngine {
    /// Create a new engine from `config`.
    pub fn new(config: ChainConfig) -> Self {
        ChainEngine {
            config,
            rr_offset: AtomicUsize::new(0),
            failure_counts: RwLock::new(HashMap::new()),
        }
    }

    /// Return the current round-robin offset without advancing it.
    ///
    /// Useful for inspection and benchmarking; the value may change
    /// concurrently and is only a snapshot.
    pub fn rr_peek_offset(&self) -> usize {
        self.rr_offset.load(Ordering::Relaxed)
    }

    /// Open a connection to `target` through the proxy chain.
    ///
    /// If the target matches a `localnet` rule, a direct connection is made.
    /// DNAT rewrites are applied before localnet and proxy selection.
    ///
    /// For `Random` and `RoundRobin` chains, the selection is retried up to
    /// `chain_retries` additional times (with a fresh random/rotating slice
    /// each attempt) before giving up.
    ///
    /// Returns a [`BoxStream`] so that TLS-wrapped hops (HTTPS proxies) and
    /// plain TCP hops share the same return type.
    pub async fn connect(&self, target: Target) -> Result<BoxStream> {
        let target = self.apply_dnat(target);

        if self.is_localnet(&target) {
            return self.direct_connect(&target).await;
        }

        // Strict / Dynamic iterate all proxies internally, so a second attempt
        // would repeat the same sequence — no benefit.
        let max_attempts = match self.config.chain_type {
            ChainType::Strict | ChainType::Dynamic => 1,
            ChainType::Random | ChainType::RoundRobin => 1 + self.config.chain_retries,
        };

        let mut last_err = anyhow::anyhow!("no proxies configured");
        for attempt in 0..max_attempts {
            let result = match self.config.chain_type {
                ChainType::Strict => self.connect_strict(target.clone()).await,
                ChainType::Dynamic => self.connect_dynamic(target.clone()).await,
                ChainType::Random => self.connect_random(target.clone()).await,
                ChainType::RoundRobin => self.connect_round_robin(target.clone()).await,
            };
            match result {
                Ok(s) => return Ok(s),
                Err(e) => {
                    if attempt + 1 < max_attempts {
                        tracing::debug!(
                            "chain attempt {} failed ({e:#}), retrying with new selection",
                            attempt + 1
                        );
                    }
                    last_err = e;
                }
            }
        }
        Err(last_err)
    }

    // ── Chain modes ───────────────────────────────────────────────────────────

    async fn connect_strict(&self, target: Target) -> Result<BoxStream> {
        // All proxies must be online — bypass the dead-proxy pre-filter so that
        // a historically-marked proxy is still attempted (and causes failure if
        // it truly is down).
        let refs: Vec<&ProxyEntry> = self.config.proxies.iter().collect();
        if refs.is_empty() {
            bail!("strict_chain: no proxies configured");
        }
        // Single fixed attempt through every proxy in order — no fallback.
        self.try_proxy_chain(&refs, target, "strict")
            .await
            .context("strict_chain")
    }

    async fn connect_dynamic(&self, target: Target) -> Result<BoxStream> {
        let refs = self.live_proxy_refs();
        if refs.is_empty() {
            bail!("dynamic_chain: no proxies configured");
        }
        // All live proxies are chained in order.  When the first-hop TCP connect
        // fails the anchor advances by one, skipping that dead proxy while keeping
        // every remaining proxy in the chain.  At least one proxy must be reachable.
        let mut last_err = anyhow::anyhow!("all proxies unreachable");
        for anchor in 0..refs.len() {
            let window = &refs[anchor..];
            match self.try_proxy_chain(window, target.clone(), "dynamic").await {
                Ok(s) => return Ok(s),
                Err(e) => last_err = e,
            }
        }
        Err(last_err).context("dynamic_chain")
    }

    async fn connect_random(&self, target: Target) -> Result<BoxStream> {
        let chain_len = self.config.chain_len.max(1);
        let pool = self.live_proxy_refs();
        if pool.len() < chain_len {
            bail!(
                "random_chain: need {chain_len} proxies, only {} available (live)",
                pool.len()
            );
        }
        // OsRng reads directly from the OS entropy source (getrandom) on every
        // call, guaranteeing different selections across runs and across
        // connections within a run.  It is also Send, so the future remains
        // Send without any drop-before-await gymnastics.
        let mut selected = pool;
        selected.shuffle(&mut OsRng);
        selected.truncate(chain_len);
        self.try_proxy_chain(&selected, target, "random")
            .await
            .context("random_chain")
    }

    async fn connect_round_robin(&self, target: Target) -> Result<BoxStream> {
        let chain_len = self.config.chain_len.max(1);
        let pool = self.live_proxy_refs();
        if pool.is_empty() {
            bail!("round_robin_chain: no proxies");
        }
        let n = pool.len();
        let offset = self.rr_offset.fetch_add(chain_len, Ordering::SeqCst) % n;
        let selected: Vec<&ProxyEntry> = (0..chain_len).map(|i| pool[(offset + i) % n]).collect();
        self.try_proxy_chain(&selected, target, "round-robin")
            .await
            .context("round_robin_chain")
    }

    // ── Helpers ───────────────────────────────────────────────────────────────

    /// Apply DNAT rewrite rules to `target`.
    ///
    /// Returns the rewritten target, or `target` unchanged if no rule matches.
    pub fn apply_dnat(&self, target: Target) -> Target {
        if let Target::Ip(IpAddr::V4(ip), port) = &target {
            for rule in &self.config.dnats {
                if rule.orig_addr == *ip {
                    if let Some(orig_port) = rule.orig_port {
                        if orig_port != *port {
                            continue;
                        }
                    }
                    let new_port = rule.new_port.unwrap_or(*port);
                    return Target::Ip(IpAddr::V4(rule.new_addr), new_port);
                }
            }
        }
        target
    }

    /// Check whether `target` matches any localnet exclusion rule.
    ///
    /// Returns `true` when the connection should bypass the proxy chain and
    /// connect directly.
    pub fn is_localnet(&self, target: &Target) -> bool {
        let (ip, port) = match target {
            Target::Ip(ip, p) => (Some(*ip), *p),
            Target::Host(_, p) => (None, *p),
        };
        let Some(ip) = ip else { return false };

        for ln in &self.config.localnets {
            if let Some(p) = ln.port {
                if p != port {
                    continue;
                }
            }
            match (ip, ln.addr) {
                (IpAddr::V4(tip), IpAddr::V4(laddr)) => {
                    if let Some(mask) = ln.mask_v4 {
                        let t = u32::from(tip);
                        let l = u32::from(laddr);
                        let m = u32::from(mask);
                        if (t & m) == (l & m) {
                            return true;
                        }
                    }
                }
                (IpAddr::V6(tip), IpAddr::V6(laddr)) => {
                    if let Some(prefix) = ln.prefix_v6 {
                        if ipv6_match(tip, laddr, prefix) {
                            return true;
                        }
                    }
                }
                _ => {}
            }
        }
        false
    }

    /// Return references to all proxies that have not been marked dead.
    ///
    /// Falls back to the full proxy list if every proxy is dead, so that
    /// callers always have at least one candidate to try.
    pub fn live_proxy_refs(&self) -> Vec<&ProxyEntry> {
        let threshold = self.config.proxy_dead_threshold;
        if threshold == 0 {
            return self.config.proxies.iter().collect();
        }
        let counts = self
            .failure_counts
            .read()
            .expect("failure_counts RwLock poisoned");
        let live: Vec<&ProxyEntry> = self
            .config
            .proxies
            .iter()
            .filter(|p| counts.get(&(p.addr, p.port)).copied().unwrap_or(0) < threshold)
            .collect();
        if live.is_empty() {
            // All proxies exceeded the threshold; reset by returning the full list
            // so the engine can keep working rather than silently failing forever.
            tracing::debug!("all proxies marked dead — using full list as fallback");
            self.config.proxies.iter().collect()
        } else {
            live
        }
    }

    /// Increment the failure counter for `proxy`.  No-op when `proxy_dead_threshold == 0`.
    ///
    /// Callers may use this to pre-populate the dead-proxy state (e.g. from
    /// persistent storage between process restarts) or to manually evict a
    /// known-bad proxy.
    pub fn record_failure(&self, proxy: &ProxyEntry) {
        if self.config.proxy_dead_threshold == 0 {
            return;
        }
        let mut counts = self
            .failure_counts
            .write()
            .expect("failure_counts RwLock poisoned");
        *counts.entry((proxy.addr, proxy.port)).or_insert(0) += 1;
        let new_count = counts[&(proxy.addr, proxy.port)];
        if new_count >= self.config.proxy_dead_threshold {
            tracing::debug!(
                "proxy {}:{} marked dead after {new_count} failures",
                proxy.addr,
                proxy.port
            );
        }
    }

    /// Attempt a single proxy chain through a pre-selected `window` of proxies.
    ///
    /// Prints each hop to stderr immediately before attempting it, so the user
    /// sees incremental progress in real time.  On failure the error marker is
    /// appended to the same line before the newline is flushed.
    async fn try_proxy_chain(
        &self,
        window: &[&ProxyEntry],
        target: Target,
        label: &str,
    ) -> Result<BoxStream> {
        debug_assert!(!window.is_empty());

        // Print label and first hop before the TCP connect attempt.
        eprint!(
            "[proxychains-tun] {label} - {}:{}",
            window[0].addr, window[0].port
        );
        let _ = std::io::stderr().flush();

        let stream = match self.tcp_connect(window[0].addr, window[0].port).await {
            Ok(s) => s,
            Err(e) => {
                eprintln!(" <--socket error or timeout!");
                tracing::debug!(
                    "{label} tcp connect to {}:{} failed: {e:#}",
                    window[0].addr, window[0].port
                );
                self.record_failure(window[0]);
                return Err(e);
            }
        };

        match chain_from(stream, window, target, self.config.tls_skip_verify).await {
            Ok(s) => Ok(s),
            Err(e) => {
                tracing::debug!("{label} handshake error: {e:#}");
                self.record_failure(window[0]);
                Err(e)
            }
        }
    }

    async fn tcp_connect(&self, addr: IpAddr, port: u16) -> Result<BoxStream> {
        let stream = timeout(
            self.config.connect_timeout,
            TcpStream::connect((addr, port)),
        )
        .await
        .context("tcp connect timed out")?
        .context("tcp connect failed")?;
        Ok(Box::new(stream))
    }

    async fn direct_connect(&self, target: &Target) -> Result<BoxStream> {
        let stream = match target {
            Target::Ip(ip, port) => timeout(
                self.config.connect_timeout,
                TcpStream::connect((*ip, *port)),
            )
            .await
            .context("direct connect timed out")?
            .context("direct connect failed")?,
            Target::Host(h, p) => timeout(
                self.config.connect_timeout,
                TcpStream::connect(format!("{h}:{p}").as_str()),
            )
            .await
            .context("direct connect timed out")?
            .context("direct connect failed")?,
        };
        Ok(Box::new(stream))
    }

}

fn ipv6_match(a: Ipv6Addr, b: Ipv6Addr, prefix: u8) -> bool {
    let a = a.octets();
    let b = b.octets();
    let full = (prefix / 8) as usize;
    let rem = prefix % 8;
    if a[..full] != b[..full] {
        return false;
    }
    if rem > 0 {
        let mask = 0xFFu8 << (8 - rem);
        if (a[full] & mask) != (b[full] & mask) {
            return false;
        }
    }
    true
}

/// Build a [`Target`] pointing to `proxy`'s address (used as intermediate hop).
fn hop_target(proxy: &ProxyEntry) -> Target {
    Target::Ip(proxy.addr, proxy.port)
}

/// Perform the appropriate protocol handshake for `prev_proxy` to connect to `next_target`.
async fn handshake(
    stream: BoxStream,
    prev_proxy: &ProxyEntry,
    next_target: Target,
    tls_skip_verify: bool,
) -> Result<BoxStream> {
    let user = prev_proxy.username.as_deref();
    let pass = prev_proxy.password.as_deref();
    match prev_proxy.proxy_type {
        ProxyType::Socks4 => socks4::connect(stream, &next_target, user).await,
        ProxyType::Socks5 => socks5::connect(stream, &next_target, user, pass).await,
        ProxyType::Http => http::connect(stream, &next_target, user, pass).await,
        ProxyType::Https => {
            https::connect(stream, &next_target, user, pass, prev_proxy.addr, tls_skip_verify)
                .await
        }
        ProxyType::Raw => raw::connect(stream, &next_target).await,
    }
}

/// Drive a pre-selected proxy window to `target`.
///
/// `window[0]` is already TCP-connected and its address was printed by the
/// caller.  This function prints each subsequent hop to stderr *before*
/// attempting the handshake, then appends `OK` or `<--socket error or
/// timeout!` and a newline once the outcome is known.
async fn chain_from(
    mut stream: BoxStream,
    window: &[&ProxyEntry],
    target: Target,
    tls_skip_verify: bool,
) -> Result<BoxStream> {
    for i in 0..window.len() - 1 {
        // Print the next hop before attempting the handshake to reach it.
        eprint!(" - {}:{}", window[i + 1].addr, window[i + 1].port);
        let _ = std::io::stderr().flush();
        match handshake(stream, window[i], hop_target(window[i + 1]), tls_skip_verify).await {
            Ok(s) => stream = s,
            Err(e) => {
                eprintln!(" <--socket error or timeout!");
                return Err(e);
            }
        }
    }

    // Final handshake: the last proxy connects us to the actual target.
    eprint!(" - {target}");
    let _ = std::io::stderr().flush();
    match handshake(stream, window[window.len() - 1], target, tls_skip_verify).await {
        Ok(s) => {
            eprintln!(" OK");
            Ok(s)
        }
        Err(e) => {
            eprintln!(" <--socket error or timeout!");
            Err(e)
        }
    }
}