// 「網路流量改走你的電腦」: the bot's computer can send its new connections out through this Mac
// (the original: 可以改走你的本機電腦，而不是雲端 — 「僅對新連線生效」, and an admin can turn it off;
// it runs an egress tunnel, Ap.egressConnectProxy 8791). Here: a small HTTP/HTTPS proxy on this
// Mac that the box's commands use through HTTP(S)_PROXY.

use std::io::{Read, Write};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::sync::atomic::{AtomicBool, AtomicU16, AtomicU64, Ordering};
use std::time::Duration;

/// The original's egress CONNECT proxy port; a free one nearby is used when it's taken.
pub const PORT: u16 = 8791;
const HEAD_MAX: usize = 64 * 1024;

static RUNNING: AtomicBool = AtomicBool::new(false);
static BOUND: AtomicU16 = AtomicU16::new(0);
pub static CONNECTIONS: AtomicU64 = AtomicU64::new(0);

/// Start the proxy once (on 127.0.0.1). Returns its port.
pub fn start() -> Result<u16, String> {
    if RUNNING.load(Ordering::SeqCst) {
        return Ok(BOUND.load(Ordering::SeqCst));
    }
    // macOS (Docker Desktop / OrbStack) forwards host.docker.internal to this machine's
    // 127.0.0.1. Linux Docker reaches the host through the bridge: bind loopback plus the
    // bridge's own gateway address (detected, `bridge_gateways`), never 0.0.0.0 — an
    // accidentally-open proxy on the LAN is an open door, and `allowed_peer` stays as the
    // second wall (a container's address changes with the network, a bound socket does not).
    let addrs = bind_addrs();
    let mut listeners = Vec::new();
    for p in PORT..PORT + 20 {
        listeners = addrs.iter().filter_map(|a| TcpListener::bind(((*a).to_string(), p)).ok()).collect();
        if listeners.len() == addrs.len() && !listeners.is_empty() {
            break;
        }
        listeners.clear();
    }
    if listeners.is_empty() {
        return Err("找不到可用的連接埠".into());
    }
    let port = listeners[0].local_addr().map_err(|e| e.to_string())?.port();
    BOUND.store(port, Ordering::SeqCst);
    RUNNING.store(true, Ordering::SeqCst);
    for l in listeners {
        std::thread::spawn(move || {
            for conn in l.incoming().flatten() {
                if !conn.peer_addr().map(|a| allowed_peer(a.ip())).unwrap_or(false) {
                    let _ = conn.shutdown(Shutdown::Both);
                    continue;
                }
                std::thread::spawn(move || {
                    let _ = handle(conn);
                });
            }
        });
    }
    Ok(port)
}

/// Where to listen: loopback, and on Linux each Docker bridge's gateway address (the address
/// containers reach us on). A host with no `ip` command keeps only loopback — the box's
// commands then fail the proxy and the person is told to turn the setting off, which beats
/// silently serving the whole LAN.
pub fn bind_addrs() -> Vec<std::net::IpAddr> {
    let mut v: Vec<std::net::IpAddr> = vec!["127.0.0.1".parse().unwrap()];
    if cfg!(target_os = "linux") {
        let out = std::process::Command::new("ip").args(["-o", "-4", "addr", "show"]).output().map(|o| String::from_utf8_lossy(&o.stdout).to_string()).unwrap_or_default();
        for (iface, ip) in parse_ipv4_addrs(&out) {
            if bridge_iface(&iface) && !v.iter().any(|a| a.to_string() == ip.to_string()) {
                v.push(std::net::IpAddr::V4(ip));
            }
        }
    }
    v
}

/// Docker's own bridges only: `docker0`, `docker_gwbridge`, and the per-network
/// `br-<12 hex>` — not user bridges like a router's `br-lan`.
pub fn bridge_iface(name: &str) -> bool {
    name == "docker0" || name == "docker_gwbridge"
        || (name.len() == 15 && name.starts_with("br-") && name[3..].bytes().all(|b| b.is_ascii_hexdigit()))
}

/// `ip -o -4 addr show` → (interface, address) pairs. Lines look like
/// `3: docker0    INET  172.17.0.1/16 brd 172.17.0.255 scope global docker0`.
pub fn parse_ipv4_addrs(out: &str) -> Vec<(String, std::net::Ipv4Addr)> {
    let mut v = Vec::new();
    for line in out.lines() {
        let mut w = line.split_whitespace();
        let _idx = w.next();
        let Some(iface) = w.next() else { continue };
        if w.next() != Some("INET") {
            continue;
        }
        let Some(cidr) = w.next() else { continue };
        if let Ok(ip) = cidr.split('/').next().unwrap_or("").parse::<std::net::Ipv4Addr>() {
            v.push((iface.trim_matches(':').to_string(), ip));
        }
    }
    v
}

/// Who may use the proxy: this machine, and Docker's default bridge (172.17.0.0/16, where the
/// boxes run). The rest of 172.16.0.0/12 is often the corporate LAN — an open proxy there is
/// an open door, so it is not served.
pub fn allowed_peer(ip: std::net::IpAddr) -> bool {
    match ip {
        std::net::IpAddr::V4(v) => v.is_loopback() || (v.octets()[0] == 172 && v.octets()[1] == 17),
        std::net::IpAddr::V6(v) => v.is_loopback(),
    }
}

pub fn port() -> Option<u16> {
    RUNNING.load(Ordering::SeqCst).then(|| BOUND.load(Ordering::SeqCst))
}

/// Where the request goes and what to send there first.
#[derive(Debug, PartialEq)]
pub struct Route {
    pub host: String,
    pub port: u16,
    /// Bytes to send to the server first (empty for CONNECT).
    pub preamble: Vec<u8>,
    pub tunnel: bool,
}

/// Parse a proxy request head: `CONNECT host:443` or `GET http://host[:port]/path HTTP/1.1`.
pub fn route(head: &[u8]) -> Result<Route, String> {
    let text = String::from_utf8_lossy(head);
    let first = text.lines().next().unwrap_or("");
    let mut parts = first.split_whitespace();
    let (method, target, version) = (parts.next().unwrap_or(""), parts.next().unwrap_or(""), parts.next().unwrap_or("HTTP/1.1"));
    if method.eq_ignore_ascii_case("CONNECT") {
        let (h, p) = target.rsplit_once(':').ok_or("CONNECT 要有連接埠")?;
        return Ok(Route { host: h.trim_matches(['[', ']']).to_string(), port: p.parse().map_err(|_| "連接埠不對")?, preamble: vec![], tunnel: true });
    }
    let rest = target.strip_prefix("http://").ok_or("只支援 http:// 或 CONNECT")?;
    let (authority, path) = match rest.find('/') {
        Some(i) => (&rest[..i], &rest[i..]),
        None => (rest, "/"),
    };
    let (host, port) = match authority.rsplit_once(':') {
        Some((h, p)) if p.chars().all(|c| c.is_ascii_digit()) => (h.to_string(), p.parse().unwrap_or(80)),
        _ => (authority.to_string(), 80),
    };
    let mut pre = format!("{method} {path} {version}\r\n").into_bytes();
    let body_start = text.find("\r\n").map(|i| i + 2).unwrap_or(text.len());
    // Drop hop-by-hop proxy headers.
    for line in text[body_start..].split_inclusive("\r\n") {
        let l = line.to_ascii_lowercase();
        if l.starts_with("proxy-connection:") || l.starts_with("proxy-authorization:") {
            continue;
        }
        pre.extend_from_slice(line.as_bytes());
    }
    Ok(Route { host, port, preamble: pre, tunnel: false })
}

fn handle(mut client: TcpStream) -> Result<(), String> {
    client.set_read_timeout(Some(Duration::from_secs(30))).ok();
    let mut buf = vec![];
    let mut chunk = [0u8; 4096];
    let head_end = loop {
        let n = client.read(&mut chunk).map_err(|e| e.to_string())?;
        if n == 0 {
            return Ok(());
        }
        buf.extend_from_slice(&chunk[..n]);
        if let Some(i) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
            break i + 4;
        }
        if buf.len() > HEAD_MAX {
            return Err("標頭太大".into());
        }
    };
    let r = match route(&buf[..head_end]) {
        Ok(r) => r,
        Err(e) => {
            let _ = client.write_all(format!("HTTP/1.1 400 Bad Request\r\nContent-Length: {}\r\n\r\n{e}", e.len()).as_bytes());
            return Err(e);
        }
    };
    // Resolve and connect with a real timeout: a plain TcpStream::connect can hang ~2 minutes
    // on a dead address, and every queued box request piles up behind it.
    let addrs: Vec<std::net::SocketAddr> = match std::net::ToSocketAddrs::to_socket_addrs(&(r.host.as_str(), r.port)) {
        Ok(a) => a.collect(),
        Err(_) => {
            let _ = client.write_all(b"HTTP/1.1 502 Bad Gateway\r\nContent-Length: 0\r\n\r\n");
            return Err(format!("連不到 {}", r.host));
        }
    };
    let mut server = match addrs.into_iter().find_map(|a| TcpStream::connect_timeout(&a, std::time::Duration::from_secs(10)).ok()) {
        Some(s) => s,
        None => {
            let _ = client.write_all(b"HTTP/1.1 502 Bad Gateway\r\nContent-Length: 0\r\n\r\n");
            return Err(format!("連不到 {}", r.host));
        }
    };
    CONNECTIONS.fetch_add(1, Ordering::Relaxed);
    // A dead peer must not pin the copy loops (and their threads) forever: an idle tunnel ends.
    let idle = std::time::Duration::from_secs(30 * 60);
    client.set_read_timeout(Some(idle)).ok();
    server.set_read_timeout(Some(idle)).ok();
    if r.tunnel {
        client.write_all(b"HTTP/1.1 200 Connection established\r\n\r\n").map_err(|e| e.to_string())?;
    } else {
        server.write_all(&r.preamble).map_err(|e| e.to_string())?;
    }
    // Anything the client sent after the head.
    if head_end < buf.len() {
        server.write_all(&buf[head_end..]).map_err(|e| e.to_string())?;
    }
    let (mut c2, mut s2) = (client.try_clone().map_err(|e| e.to_string())?, server.try_clone().map_err(|e| e.to_string())?);
    let up = std::thread::spawn(move || {
        let _ = std::io::copy(&mut c2, &mut s2);
        let _ = s2.shutdown(Shutdown::Write);
    });
    let _ = std::io::copy(&mut server, &mut client);
    let _ = client.shutdown(Shutdown::Write);
    let _ = up.join();
    Ok(())
}

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

    #[test]
    fn binds_bridge_gateways_not_the_whole_lan() {
        let out = "1: lo    INET 127.0.0.1/8 scope host lo\n3: docker0    INET 172.17.0.1/16 brd 172.17.0.255 scope global docker0\n8: br-5f3a11c2d4e6    INET 172.18.0.1/16 scope global br-5f3a11c2d4e6\n2: eth0    INET 192.168.1.30/24 brd 192.168.1.255 scope global eth0\n5: virbr0    INET 192.168.122.1/24 scope global virbr0\n";
        let addrs: Vec<_> = parse_ipv4_addrs(out).into_iter().filter(|(i, _)| bridge_iface(i)).map(|(_, ip)| ip.to_string()).collect();
        assert_eq!(addrs, ["172.17.0.1", "172.18.0.1"]); // eth0 and virbr0 are not Docker's
        assert!(bridge_iface("docker_gwbridge") && !bridge_iface("br-lan") && !bridge_iface("docker0x"));
        // The proxy's second wall matches what it binds: loopback and 172.17/16 only.
        assert!(allowed_peer("127.0.0.1".parse().unwrap()));
        assert!(allowed_peer("172.17.0.5".parse().unwrap()));
        assert!(!allowed_peer("172.18.0.5".parse().unwrap())); // a user-defined network is not served
        assert!(!allowed_peer("192.168.1.20".parse().unwrap()));
    }

    #[test]
    fn routes_connect_and_plain_http() {
        let r = route(b"CONNECT example.com:443 HTTP/1.1\r\nHost: example.com:443\r\n\r\n").unwrap();
        assert_eq!((r.host.as_str(), r.port, r.tunnel), ("example.com", 443, true));
        let r = route(b"GET http://example.com:8080/a?b=1 HTTP/1.1\r\nHost: example.com\r\nProxy-Connection: keep-alive\r\n\r\n").unwrap();
        assert_eq!((r.host.as_str(), r.port, r.tunnel), ("example.com", 8080, false));
        let pre = String::from_utf8(r.preamble).unwrap();
        assert!(pre.starts_with("GET /a?b=1 HTTP/1.1\r\n") && !pre.to_lowercase().contains("proxy-connection"));
        assert!(route(b"GET /x HTTP/1.1\r\n\r\n").is_err());
    }

    #[test]
    fn proxies_a_real_request() {
        // A tiny origin server on this machine.
        let origin = TcpListener::bind("127.0.0.1:0").unwrap();
        let op = origin.local_addr().unwrap().port();
        std::thread::spawn(move || {
            if let Ok((mut s, _)) = origin.accept() {
                let mut b = [0u8; 1024];
                let _ = s.read(&mut b);
                let _ = s.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\nConnection: close\r\n\r\nhello");
            }
        });
        let port = start().unwrap();
        let mut c = TcpStream::connect(("127.0.0.1", port)).unwrap();
        c.write_all(format!("GET http://127.0.0.1:{op}/ HTTP/1.1\r\nHost: x\r\n\r\n").as_bytes()).unwrap();
        let mut out = String::new();
        c.read_to_string(&mut out).unwrap();
        assert!(out.ends_with("hello"), "{out}");
    }

    #[test]
    fn only_this_machine_and_containers() {
        assert!(super::allowed_peer("127.0.0.1".parse().unwrap()));
        // Docker's default bridge (where the boxes run) is served…
        assert!(super::allowed_peer("172.17.0.2".parse().unwrap()));
        // …but the rest of 172.16/12 (often the corporate LAN) is not: an open proxy there
        // is an open door.
        assert!(!super::allowed_peer("172.16.0.1".parse().unwrap()));
        assert!(!super::allowed_peer("172.31.255.1".parse().unwrap()));
        assert!(!super::allowed_peer("172.32.0.1".parse().unwrap()));
        assert!(!super::allowed_peer("192.168.1.20".parse().unwrap()));
    }
}