//! Passive TLS ClientHello observer for the Blaze redirector.
//!
//! ```text
//!   FIFA ──> tls-observe ──(raw bytes, unmodified)──> Python redirector
//!               │
//!               └── parses and reports the ClientHello
//! ```
//!
//! # Why this exists
//!
//! Every observed redirector handshake selected `AES256-GCM-SHA384` — TLS 1.2
//! with **static RSA key exchange**. `rustls` supports only forward-secret
//! (EC)DHE suites, so if that is all the client offers, rustls is ruled out for
//! a Rust redirector and an OpenSSL-backed stack is required.
//!
//! But **the selected cipher does not reveal what was offered.** OpenSSL's
//! server follows the client's preference order by default, so preferring
//! static RSA does not prove ECDHE was unavailable. Choosing a TLS stack on
//! that inference would be exactly the kind of guess this project keeps
//! refusing to make. So: read the actual ClientHello.
//!
//! # Passive by construction
//!
//! Bytes are forwarded verbatim in both directions and nothing is injected,
//! rewritten or delayed beyond a parse of the first record. The handshake is
//! still terminated by the untouched Python redirector, so a FIFA session runs
//! exactly as it otherwise would. If parsing fails, the proxy still relays —
//! observation must never be able to break the path it is observing.
//!
//! ```text
//!   tls-observe --listen 0.0.0.0:42227 --upstream 127.0.0.1:42127
//! ```

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

// ---------------------------------------------------------------- TLS names

/// Cipher suites we care about naming. Not exhaustive: unknown values are
/// reported as hex so nothing is silently dropped.
fn cipher_name(id: u16) -> &'static str {
    match id {
        0x0004 => "TLS_RSA_WITH_RC4_128_MD5",
        0x0005 => "TLS_RSA_WITH_RC4_128_SHA",
        0x000A => "TLS_RSA_WITH_3DES_EDE_CBC_SHA",
        0x002F => "TLS_RSA_WITH_AES_128_CBC_SHA",
        0x0035 => "TLS_RSA_WITH_AES_256_CBC_SHA",
        0x003C => "TLS_RSA_WITH_AES_128_CBC_SHA256",
        0x003D => "TLS_RSA_WITH_AES_256_CBC_SHA256",
        0x009C => "TLS_RSA_WITH_AES_128_GCM_SHA256",
        0x009D => "TLS_RSA_WITH_AES_256_GCM_SHA384",
        0x0033 => "TLS_DHE_RSA_WITH_AES_128_CBC_SHA",
        0x0039 => "TLS_DHE_RSA_WITH_AES_256_CBC_SHA",
        0x009E => "TLS_DHE_RSA_WITH_AES_128_GCM_SHA256",
        0x009F => "TLS_DHE_RSA_WITH_AES_256_GCM_SHA384",
        0xC013 => "TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA",
        0xC014 => "TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA",
        0xC027 => "TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256",
        0xC028 => "TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA384",
        0xC02F => "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256",
        0xC030 => "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384",
        0xC009 => "TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA",
        0xC00A => "TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA",
        0xC02B => "TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256",
        0xC02C => "TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384",
        0xCCA8 => "TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256",
        0x1301 => "TLS_AES_128_GCM_SHA256 (1.3)",
        0x1302 => "TLS_AES_256_GCM_SHA384 (1.3)",
        0x1303 => "TLS_CHACHA20_POLY1305_SHA256 (1.3)",
        0x00FF => "TLS_EMPTY_RENEGOTIATION_INFO_SCSV",
        _ => "",
    }
}

/// Does this suite use an ephemeral (forward-secret) key exchange?
///
/// This is the whole decision: rustls offers ECDHE/DHE suites only.
fn is_forward_secret(id: u16) -> bool {
    let n = cipher_name(id);
    n.contains("ECDHE") || n.contains("DHE_") || n.ends_with("(1.3)")
}

fn tls_version_name(v: u16) -> &'static str {
    match v {
        0x0300 => "SSL 3.0",
        0x0301 => "TLS 1.0",
        0x0302 => "TLS 1.1",
        0x0303 => "TLS 1.2",
        0x0304 => "TLS 1.3",
        _ => "unknown",
    }
}

fn extension_name(id: u16) -> &'static str {
    match id {
        0 => "server_name",
        5 => "status_request",
        10 => "supported_groups",
        11 => "ec_point_formats",
        13 => "signature_algorithms",
        16 => "ALPN",
        23 => "extended_master_secret",
        35 => "session_ticket",
        43 => "supported_versions",
        45 => "psk_key_exchange_modes",
        51 => "key_share",
        65281 => "renegotiation_info",
        _ => "",
    }
}

// ------------------------------------------------------------- parsing

#[derive(Debug, Default)]
struct ClientHello {
    record_version: u16,
    client_version: u16,
    cipher_suites: Vec<u16>,
    extensions: Vec<u16>,
    server_name: Option<String>,
    supported_versions: Vec<u16>,
    supported_groups: Vec<u16>,
}

fn be16(b: &[u8], i: usize) -> Option<u16> {
    Some(u16::from_be_bytes([*b.get(i)?, *b.get(i + 1)?]))
}

/// Parse a ClientHello from the first TLS record.
///
/// Returns `None` rather than erroring: an unparseable hello must still be
/// relayed, so the caller treats this as "nothing learned", never as a failure.
fn parse_client_hello(buf: &[u8]) -> Option<ClientHello> {
    // TLS record: type(1) version(2) length(2)
    if *buf.first()? != 0x16 {
        return None; // not a handshake record
    }
    let mut h = ClientHello {
        record_version: be16(buf, 1)?,
        ..Default::default()
    };

    // Handshake: type(1) length(3) then the body.
    let hs = buf.get(5..)?;
    if *hs.first()? != 0x01 {
        return None; // not a ClientHello
    }
    let mut i = 4; // skip handshake header
    h.client_version = be16(hs, i)?;
    i += 2;
    i += 32; // random
    let sid_len = *hs.get(i)? as usize;
    i += 1 + sid_len;

    let cs_len = be16(hs, i)? as usize;
    i += 2;
    for k in (0..cs_len).step_by(2) {
        if let Some(c) = be16(hs, i + k) {
            h.cipher_suites.push(c);
        }
    }
    i += cs_len;

    let comp_len = *hs.get(i)? as usize;
    i += 1 + comp_len;

    // Extensions are optional (SSL3/TLS1.0 clients may omit them).
    let Some(ext_total) = be16(hs, i) else {
        return Some(h);
    };
    i += 2;
    let end = i + ext_total as usize;
    while i + 4 <= end.min(hs.len()) {
        let etype = be16(hs, i)?;
        let elen = be16(hs, i + 2)? as usize;
        let body = hs.get(i + 4..i + 4 + elen).unwrap_or(&[]);
        h.extensions.push(etype);
        match etype {
            0 => {
                // server_name: list(2) type(1) len(2) host
                if body.len() > 5 {
                    let n = be16(body, 3).unwrap_or(0) as usize;
                    if let Some(s) = body.get(5..5 + n) {
                        h.server_name = Some(String::from_utf8_lossy(s).into_owned());
                    }
                }
            }
            43 => {
                if let Some(&n) = body.first() {
                    for k in (0..n as usize).step_by(2) {
                        if let Some(v) = be16(body, 1 + k) {
                            h.supported_versions.push(v);
                        }
                    }
                }
            }
            10 => {
                let n = be16(body, 0).unwrap_or(0) as usize;
                for k in (0..n).step_by(2) {
                    if let Some(g) = be16(body, 2 + k) {
                        h.supported_groups.push(g);
                    }
                }
            }
            _ => {}
        }
        i += 4 + elen;
    }
    Some(h)
}

fn report(conn: u64, peer: &str, h: &ClientHello) -> String {
    let mut s = String::new();
    s.push_str(&format!("=== ClientHello #{conn} from {peer} ===\n"));
    s.push_str(&format!(
        "  record version : 0x{:04x} {}\n",
        h.record_version,
        tls_version_name(h.record_version)
    ));
    s.push_str(&format!(
        "  client version : 0x{:04x} {}\n",
        h.client_version,
        tls_version_name(h.client_version)
    ));
    if !h.supported_versions.is_empty() {
        let v: Vec<String> = h
            .supported_versions
            .iter()
            .map(|v| tls_version_name(*v).to_string())
            .collect();
        s.push_str(&format!("  supported_versions: {}\n", v.join(", ")));
    }
    if let Some(sni) = &h.server_name {
        s.push_str(&format!("  SNI            : {sni}\n"));
    }
    s.push_str(&format!("  cipher suites  : {}\n", h.cipher_suites.len()));
    for c in &h.cipher_suites {
        let name = cipher_name(*c);
        s.push_str(&format!(
            "    0x{:04x}  {:<45} {}\n",
            c,
            if name.is_empty() { "<unknown>" } else { name },
            if is_forward_secret(*c) {
                "[forward-secret]"
            } else {
                ""
            }
        ));
    }
    let fs: Vec<u16> = h
        .cipher_suites
        .iter()
        .copied()
        .filter(|c| is_forward_secret(*c))
        .collect();
    s.push_str(&format!(
        "  extensions     : {}\n",
        h.extensions
            .iter()
            .map(|e| {
                let n = extension_name(*e);
                if n.is_empty() {
                    format!("{e}")
                } else {
                    n.to_string()
                }
            })
            .collect::<Vec<_>>()
            .join(", ")
    ));
    s.push('\n');
    s.push_str("  VERDICT:\n");
    if fs.is_empty() {
        s.push_str("    NO forward-secret suite offered.\n");
        s.push_str("    => rustls is RULED OUT for the redirector; an OpenSSL-backed\n");
        s.push_str("       stack (native-tls / openssl) is required.\n");
    } else {
        s.push_str(&format!(
            "    {} forward-secret suite(s) OFFERED:\n",
            fs.len()
        ));
        for c in &fs {
            s.push_str(&format!("      0x{:04x}  {}\n", c, cipher_name(*c)));
        }
        s.push_str("    => rustls is VIABLE in principle. Confirm with a real handshake\n");
        s.push_str("       against a rustls listener before deciding.\n");
    }
    s
}

// --------------------------------------------------------------- proxy

fn pump(mut from: TcpStream, mut to: TcpStream) {
    let mut buf = [0u8; 32768];
    loop {
        match from.read(&mut buf) {
            Ok(0) | Err(_) => break,
            Ok(n) => {
                if to.write_all(&buf[..n]).is_err() {
                    break;
                }
                let _ = to.flush();
            }
        }
    }
    let _ = to.shutdown(std::net::Shutdown::Write);
}

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let get = |flag: &str| -> Option<String> {
        args.windows(2).find(|w| w[0] == flag).map(|w| w[1].clone())
    };
    let listen = get("--listen").unwrap_or_else(|| {
        eprintln!(
            "usage: tls-observe --listen <addr:port> --upstream <addr:port> [--report <path>]"
        );
        eprintln!();
        eprintln!("Passively observes the TLS ClientHello and relays every byte to the");
        eprintln!("upstream redirector unmodified. Never terminates TLS itself.");
        std::process::exit(2);
    });
    let upstream = get("--upstream").unwrap_or_else(|| {
        eprintln!("--upstream is required (the real redirector, e.g. 127.0.0.1:42127)");
        std::process::exit(2);
    });
    let report_path = get("--report");

    let listener = TcpListener::bind(&listen).unwrap_or_else(|e| {
        eprintln!("cannot bind {listen}: {e}");
        std::process::exit(1);
    });
    eprintln!("tls-observe: {listen} -> {upstream} (passive; TLS terminated upstream)");

    let counter = Arc::new(AtomicU64::new(0));
    for incoming in listener.incoming() {
        let Ok(mut client) = incoming else { continue };
        let id = counter.fetch_add(1, Ordering::Relaxed) + 1;
        let upstream = upstream.clone();
        let report_path = report_path.clone();

        std::thread::spawn(move || {
            let peer = client
                .peer_addr()
                .map(|a| a.to_string())
                .unwrap_or_else(|_| "<unknown>".into());
            let _ = client.set_nodelay(true);

            // Read the first chunk: the ClientHello. Observation only — these
            // bytes are forwarded verbatim regardless of what we make of them.
            let mut head = vec![0u8; 8192];
            let _ = client.set_read_timeout(Some(Duration::from_secs(15)));
            let n = match client.read(&mut head) {
                Ok(0) | Err(_) => return,
                Ok(n) => n,
            };
            head.truncate(n);
            let _ = client.set_read_timeout(None);

            match parse_client_hello(&head) {
                Some(h) => {
                    let text = report(id, &peer, &h);
                    print!("{text}");
                    use std::io::Write as _;
                    let _ = std::io::stdout().flush();
                    if let Some(p) = &report_path {
                        if let Ok(mut f) = std::fs::OpenOptions::new()
                            .create(true)
                            .append(true)
                            .open(p)
                        {
                            let _ = f.write_all(text.as_bytes());
                        }
                    }
                }
                None => {
                    eprintln!("conn {id} from {peer}: first record is not a ClientHello ({n}B)")
                }
            }

            let Ok(mut server) = TcpStream::connect(&upstream) else {
                eprintln!("conn {id}: cannot reach upstream {upstream}");
                return;
            };
            let _ = server.set_nodelay(true);
            if server.write_all(&head).is_err() {
                return;
            }
            let _ = server.flush();

            let (c2, s2) = match (client.try_clone(), server.try_clone()) {
                (Ok(a), Ok(b)) => (a, b),
                _ => return,
            };
            let up = std::thread::spawn(move || pump(client, server));
            pump(s2, c2);
            let _ = up.join();
        });
    }
}

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

    /// A minimal TLS 1.2 ClientHello offering one ECDHE and one static-RSA suite.
    fn synthetic_hello() -> Vec<u8> {
        let mut body = Vec::new();
        body.extend_from_slice(&[0x03, 0x03]); // client_version TLS 1.2
        body.extend_from_slice(&[0u8; 32]); // random
        body.push(0); // session id len
        body.extend_from_slice(&2u16.to_be_bytes().map(|b| b)); // placeholder
        let cs: [u16; 2] = [0xC030, 0x009D];
        let cs_bytes: Vec<u8> = cs.iter().flat_map(|c| c.to_be_bytes()).collect();
        let l = body.len();
        body.truncate(l - 2);
        body.extend_from_slice(&(cs_bytes.len() as u16).to_be_bytes());
        body.extend_from_slice(&cs_bytes);
        body.push(1); // compression len
        body.push(0); // null
        body.extend_from_slice(&0u16.to_be_bytes()); // no extensions

        let mut hs = vec![0x01];
        hs.extend_from_slice(&(body.len() as u32).to_be_bytes()[1..]);
        hs.extend_from_slice(&body);

        let mut rec = vec![0x16, 0x03, 0x01];
        rec.extend_from_slice(&(hs.len() as u16).to_be_bytes());
        rec.extend_from_slice(&hs);
        rec
    }

    #[test]
    fn parses_ciphers_and_classifies_forward_secrecy() {
        let h = parse_client_hello(&synthetic_hello()).expect("parses");
        assert_eq!(h.client_version, 0x0303);
        assert_eq!(h.cipher_suites, vec![0xC030, 0x009D]);
        assert!(
            is_forward_secret(0xC030),
            "ECDHE must count as forward-secret"
        );
        assert!(!is_forward_secret(0x009D), "static RSA must not");
    }

    #[test]
    fn verdict_reflects_what_was_offered() {
        let h = parse_client_hello(&synthetic_hello()).unwrap();
        let text = report(1, "test", &h);
        assert!(text.contains("rustls is VIABLE"), "{text}");

        let only_static = ClientHello {
            cipher_suites: vec![0x009D, 0x002F],
            ..Default::default()
        };
        let text = report(1, "test", &only_static);
        assert!(text.contains("RULED OUT"), "{text}");
    }

    #[test]
    fn non_tls_input_is_declined_rather_than_misparsed() {
        assert!(parse_client_hello(b"GET / HTTP/1.1\r\n\r\n").is_none());
        assert!(parse_client_hello(&[]).is_none());
        // A handshake record that is not a ClientHello.
        assert!(parse_client_hello(&[0x16, 0x03, 0x01, 0x00, 0x04, 0x02, 0, 0, 0]).is_none());
    }

    #[test]
    fn truncated_hello_does_not_panic() {
        let full = synthetic_hello();
        for cut in 1..full.len() {
            let _ = parse_client_hello(&full[..cut]);
        }
    }
}