zerogravity/src/snapshot.rs

//! HTTP/2 traffic snapshot parser.
//!
//! Parses Go's `GODEBUG=http2debug=2` output into structured,
//! human-readable traffic snapshots. Can also be used to parse
//! MITM proxy raw data dumps.

use std::collections::HashMap;
use std::io::{self, Read};

// ── Domain metadata ──────────────────────────────────────────────────────────

const DOMAIN_INFO: &[(&str, &str, &str)] = &[
    (
        "antigravity-unleash.goog",
        "Feature Flags",
        "Unleash SDK — controls A/B tests and feature rollouts",
    ),
    (
        "daily-cloudcode-pa.googleapis.com",
        "LLM API (gRPC)",
        "Primary Gemini/Claude API endpoint",
    ),
    (
        "cloudcode-pa.googleapis.com",
        "LLM API (gRPC)",
        "Production Gemini/Claude API endpoint",
    ),
    (
        "api.anthropic.com",
        "Claude API",
        "Direct Anthropic API calls",
    ),
    (
        "lh3.googleusercontent.com",
        "Profile Picture",
        "User avatar",
    ),
    ("play.googleapis.com", "Telemetry", "Google Play telemetry"),
    (
        "firebaseinstallations.googleapis.com",
        "Firebase",
        "Installation tracking",
    ),
    ("oauth2.googleapis.com", "OAuth", "Token refresh/exchange"),
    ("speech.googleapis.com", "Speech", "Voice input processing"),
    (
        "modelarmor.googleapis.com",
        "Safety",
        "Content safety/filtering",
    ),
];

fn domain_label(domain: &str) -> (&str, &str) {
    for &(d, label, desc) in DOMAIN_INFO {
        if domain == d || domain.ends_with(d) {
            return (label, desc);
        }
    }
    ("External", "")
}

// ── ANSI colors ──────────────────────────────────────────────────────────────

const BOLD: &str = "\x1b[1m";
const DIM: &str = "\x1b[2m";
const GREEN: &str = "\x1b[92m";
const YELLOW: &str = "\x1b[93m";
const CYAN: &str = "\x1b[96m";
const MAGENTA: &str = "\x1b[95m";
const NC: &str = "\x1b[0m";

// ── Parsed types ─────────────────────────────────────────────────────────────

#[derive(Debug, Clone)]
struct HttpExchange {
    authority: String,
    method: String,
    path: String,
    headers: Vec<(String, String)>,
    body: Vec<u8>,
    body_total_len: usize,
    stream_id: Option<String>,
    _direction: Direction,
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum Direction {
    Outgoing, // LS → upstream
    Incoming, // external → LS (our curl calls)
}

#[derive(Default)]
struct Snapshot {
    connections: Vec<String>,
    exchanges: Vec<HttpExchange>,
    ls_logs: Vec<String>,
}

// ── Interesting headers to show prominently ──────────────────────────────────

const INTERESTING_HEADERS: &[&str] = &[
    "authorization",
    "content-type",
    "user-agent",
    "unleash-appname",
    "unleash-instanceid",
    "unleash-sdk",
    "x-goog-api-key",
    "x-goog-api-client",
    "grpc-encoding",
    "te",
];

// ── Parser ───────────────────────────────────────────────────────────────────

impl Snapshot {
    fn parse(input: &str) -> Self {
        let mut snap = Snapshot::default();
        let mut current_headers: Vec<(String, String)> = Vec::new();
        let mut current_pseudo: HashMap<String, String> = HashMap::new();
        let mut current_direction = Direction::Outgoing;
        let mut current_stream: Option<String> = None;

        for line in input.lines() {
            let line = line.trim_end();
            if line.is_empty() {
                continue;
            }

            // LS process logs
            if (line.starts_with('I') || line.starts_with('W') || line.starts_with('E'))
                && line.len() > 4
                && line.chars().nth(1).is_some_and(|c| c.is_ascii_digit())
            {
                snap.ls_logs.push(line.to_string());
                continue;
            }
            if line.contains("maxprocs:") {
                snap.ls_logs.push(line.to_string());
                continue;
            }

            // New connection
            if let Some(target) = extract_after(line, "Transport creating client conn") {
                let target = target.trim_start_matches(|c: char| !c.is_alphanumeric());
                if let Some(to) = target.strip_prefix("to ") {
                    snap.connections.push(to.to_string());
                } else {
                    // e.g. "0x... to host:port"
                    if let Some(pos) = target.find(" to ") {
                        snap.connections.push(target[pos + 4..].to_string());
                    }
                }
                continue;
            }

            // Outgoing request headers (Transport encoding header)
            if let Some((key, val)) = extract_header(line, "Transport encoding header") {
                if key == ":method" {
                    // Finalize previous exchange
                    if current_pseudo.contains_key(":path")
                        || current_pseudo.contains_key(":method")
                    {
                        snap.finalize_exchange(
                            &current_pseudo,
                            &current_headers,
                            current_direction,
                            current_stream.clone(),
                        );
                    }
                    current_headers.clear();
                    current_pseudo.clear();
                    current_direction = Direction::Outgoing;
                    current_stream = None;
                }
                if key.starts_with(':') {
                    current_pseudo.insert(key, val);
                } else {
                    current_headers.push((key, val));
                }
                continue;
            }

            // Incoming / server-received headers
            if let Some((key, val)) = extract_header(line, "decoded hpack field header field") {
                if key == ":authority" && !line.contains("server read frame") {
                    if current_pseudo.contains_key(":path")
                        || current_pseudo.contains_key(":method")
                    {
                        snap.finalize_exchange(
                            &current_pseudo,
                            &current_headers,
                            current_direction,
                            current_stream.clone(),
                        );
                    }
                    current_headers.clear();
                    current_pseudo.clear();
                    current_direction = Direction::Incoming;
                    current_stream = None;
                }
                if key.starts_with(':') {
                    current_pseudo.insert(key, val);
                } else {
                    current_headers.push((key, val));
                }
                continue;
            }

            // HEADERS frame (extracts stream ID)
            if line.contains("wrote HEADERS") {
                if let Some(stream) = extract_stream_id(line) {
                    current_stream = Some(stream.clone());
                    if current_pseudo.contains_key(":path")
                        || current_pseudo.contains_key(":method")
                    {
                        let ex = snap.finalize_exchange(
                            &current_pseudo,
                            &current_headers,
                            current_direction,
                            Some(stream),
                        );
                        if ex.is_some() {
                            current_headers.clear();
                            current_pseudo.clear();
                        }
                    }
                }
                continue;
            }

            // DATA frames
            if (line.contains("wrote DATA")
                || line.contains("read DATA")
                || line.contains("server read frame DATA"))
                && line.contains("data=\"")
            {
                let is_outgoing =
                    line.contains("wrote DATA") || line.contains("server read frame DATA");
                if let Some(stream) = extract_stream_id(line) {
                    if let Some(data_str) = extract_data(line) {
                        let raw = decode_go_escaped(&data_str);
                        let len = extract_data_len(line).unwrap_or(raw.len());
                        // Find matching exchange by stream
                        for ex in snap.exchanges.iter_mut().rev() {
                            if ex.stream_id.as_deref() == Some(&stream) && is_outgoing {
                                ex.body.extend_from_slice(&raw);
                                ex.body_total_len = ex.body_total_len.max(len);
                                break;
                            }
                        }
                    }
                }
                continue;
            }
        }

        // Finalize remaining
        if current_pseudo.contains_key(":path") || current_pseudo.contains_key(":method") {
            snap.finalize_exchange(
                &current_pseudo,
                &current_headers,
                current_direction,
                current_stream,
            );
        }

        snap
    }

    fn finalize_exchange(
        &mut self,
        pseudo: &HashMap<String, String>,
        headers: &[(String, String)],
        direction: Direction,
        stream_id: Option<String>,
    ) -> Option<usize> {
        let method = pseudo.get(":method").cloned().unwrap_or_default();
        let path = pseudo.get(":path").cloned().unwrap_or_default();
        let authority = pseudo.get(":authority").cloned().unwrap_or_default();

        if method.is_empty() && path.is_empty() && authority.is_empty() {
            return None;
        }

        self.exchanges.push(HttpExchange {
            authority,
            method: if method.is_empty() {
                "GET".into()
            } else {
                method
            },
            path,
            headers: headers.to_vec(),
            body: Vec::new(),
            body_total_len: 0,
            stream_id,
            _direction: direction,
        });
        Some(self.exchanges.len() - 1)
    }

    // ── Rendering ────────────────────────────────────────────────────────────

    fn render(&self) -> String {
        let mut out = String::new();

        let sep = "═".repeat(70);
        let sep_thin = "─".repeat(60);
        out.push_str(&format!("\n{BOLD}{CYAN}{sep}{NC}\n"));
        out.push_str(&format!(
            "{BOLD}{CYAN}  STANDALONE LS TRAFFIC SNAPSHOT{NC}\n"
        ));
        out.push_str(&format!("{BOLD}{CYAN}{sep}{NC}\n\n"));

        // LS Logs
        if !self.ls_logs.is_empty() {
            out.push_str(&format!("{BOLD}▸ Language Server Logs{NC}\n"));
            out.push_str(&format!("{DIM}{sep_thin}{NC}\n"));
            for log in &self.ls_logs {
                out.push_str(&format!("  {DIM}{log}{NC}\n"));
            }
            out.push('\n');
        }

        // Connections
        if !self.connections.is_empty() {
            out.push_str(&format!("{BOLD}▸ Outbound Connections{NC}\n"));
            out.push_str(&format!("{DIM}{sep_thin}{NC}\n"));
            for target in &self.connections {
                let domain = target.split(':').next().unwrap_or(target);
                let (label, desc) = domain_label(domain);
                out.push_str(&format!(
                    "  {GREEN}→{NC} {BOLD}{target}{NC}  {DIM}({label}){NC}\n"
                ));
                if !desc.is_empty() {
                    out.push_str(&format!("    {DIM}{desc}{NC}\n"));
                }
            }
            out.push('\n');
        }

        // Group by domain
        let mut by_domain: Vec<(&str, Vec<&HttpExchange>)> = Vec::new();
        for ex in &self.exchanges {
            if let Some(entry) = by_domain
                .iter_mut()
                .find(|(d, _)| *d == ex.authority.as_str())
            {
                entry.1.push(ex);
            } else {
                by_domain.push((&ex.authority, vec![ex]));
            }
        }

        for (domain, exchanges) in &by_domain {
            if domain.starts_with("127.0.0.1") {
                // Skip local requests (our own curl calls to the LS)
                continue;
            }

            let (label, _desc) = domain_label(domain);
            let color = if label.contains("API") { YELLOW } else { CYAN };

            out.push_str(&format!("\n{BOLD}{sep}{NC}\n"));
            out.push_str(&format!(
                "{BOLD}{color}  {domain}{NC}  {DIM}— {label}{NC}\n"
            ));
            out.push_str(&format!("{BOLD}{sep}{NC}\n"));

            for ex in exchanges {
                let method_color = if ex.method == "GET" { GREEN } else { YELLOW };
                out.push_str(&format!(
                    "\n  {BOLD}→ {method_color}{}{NC} {}\n",
                    ex.method, ex.path
                ));

                // Interesting headers
                for (key, val) in &ex.headers {
                    let key_lower = key.to_lowercase();
                    if INTERESTING_HEADERS.contains(&key_lower.as_str()) {
                        let display_val = mask_token(&key_lower, val);
                        out.push_str(&format!("    {DIM}{key}:{NC} {display_val}\n"));
                    }
                }

                // Other headers
                for (key, val) in &ex.headers {
                    let key_lower = key.to_lowercase();
                    if !INTERESTING_HEADERS.contains(&key_lower.as_str()) {
                        out.push_str(&format!("    {DIM}{key}:{NC} {val}\n"));
                    }
                }

                // Body
                if !ex.body.is_empty() {
                    out.push_str(&render_body(&ex.body, ex.body_total_len));
                }
            }
            out.push('\n');
        }

        out
    }
}

// ── Body rendering ───────────────────────────────────────────────────────────

fn render_body(data: &[u8], total_len: usize) -> String {
    let mut out = String::new();
    let len = total_len.max(data.len());

    // Try JSON
    if let Ok(text) = std::str::from_utf8(data) {
        if let Ok(val) = serde_json::from_str::<serde_json::Value>(text) {
            let pretty = serde_json::to_string_pretty(&val).unwrap_or_default();
            out.push_str(&format!("    {BOLD}Body ({len} bytes, JSON):{NC}\n"));
            for (i, line) in pretty.lines().enumerate() {
                if i >= 40 {
                    out.push_str(&format!(
                        "      {DIM}... ({} more lines){NC}\n",
                        pretty.lines().count() - 40
                    ));
                    break;
                }
                out.push_str(&format!("      {GREEN}{line}{NC}\n"));
            }
            return out;
        }
    }

    // Try gzip → JSON
    if data.len() >= 2 && data[0] == 0x1f && data[1] == 0x8b {
        if let Ok(decompressed) = decompress_gzip(data) {
            if let Ok(text) = std::str::from_utf8(&decompressed) {
                if let Ok(val) = serde_json::from_str::<serde_json::Value>(text) {
                    let pretty = serde_json::to_string_pretty(&val).unwrap_or_default();
                    out.push_str(&format!(
                        "    {BOLD}Body ({len} bytes gzip → {} bytes, JSON):{NC}\n",
                        decompressed.len()
                    ));
                    for (i, line) in pretty.lines().enumerate() {
                        if i >= 50 {
                            out.push_str(&format!(
                                "      {DIM}... ({} more lines){NC}\n",
                                pretty.lines().count() - 50
                            ));
                            break;
                        }
                        out.push_str(&format!("      {GREEN}{line}{NC}\n"));
                    }
                    return out;
                }
                // Plain text
                out.push_str(&format!(
                    "    {BOLD}Body ({len} bytes gzip → {} bytes, text):{NC}\n",
                    decompressed.len()
                ));
                for line in text.lines().take(20) {
                    out.push_str(&format!("      {line}\n"));
                }
                return out;
            }
            // Binary gzip
            out.push_str(&format!(
                "    {BOLD}Body ({len} bytes gzip → {} bytes, binary):{NC}\n",
                decompressed.len()
            ));
            let strings = extract_strings(&decompressed);
            for s in strings.iter().take(15) {
                out.push_str(&format!("      {MAGENTA}{s}{NC}\n"));
            }
            return out;
        }
    }

    // PNG
    if data.len() >= 4 && &data[..4] == b"\x89PNG" {
        out.push_str(&format!("    {BOLD}Body ({len} bytes, PNG image){NC}\n"));
        return out;
    }

    // Protobuf / binary with string extraction
    let strings = extract_strings(data);
    if !strings.is_empty() {
        let kind = if !data.is_empty()
            && matches!(
                data[0],
                0x08 | 0x0a | 0x10 | 0x12 | 0x18 | 0x1a | 0x20 | 0x22
            ) {
            "protobuf"
        } else {
            "binary"
        };
        out.push_str(&format!("    {BOLD}Body ({len} bytes, {kind}):{NC}\n"));
        out.push_str(&format!("    {DIM}Extracted strings:{NC}\n"));
        for s in strings.iter().take(20) {
            out.push_str(&format!("      {MAGENTA}{s}{NC}\n"));
        }
        return out;
    }

    // Plain text fallback
    if let Ok(text) = std::str::from_utf8(data) {
        out.push_str(&format!("    {BOLD}Body ({len} bytes, text):{NC}\n"));
        for line in text.lines().take(10) {
            out.push_str(&format!("      {line}\n"));
        }
        return out;
    }

    out.push_str(&format!("    {BOLD}Body ({len} bytes, binary){NC}\n"));
    out
}

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

fn extract_after<'a>(line: &'a str, pattern: &str) -> Option<&'a str> {
    line.find(pattern).map(|pos| &line[pos + pattern.len()..])
}

fn extract_header(line: &str, pattern: &str) -> Option<(String, String)> {
    if !line.contains(pattern) {
        return None;
    }
    // Pattern: ... "key" = "value"
    let after = extract_after(line, pattern)?;
    let first_quote = after.find('"')?;
    let rest = &after[first_quote + 1..];
    let end_key = rest.find('"')?;
    let key = &rest[..end_key];

    let after_eq = &rest[end_key + 1..];
    let val_start = after_eq.find('"')?;
    let val_rest = &after_eq[val_start + 1..];
    let end_val = val_rest.rfind('"').unwrap_or(val_rest.len());
    let val = &val_rest[..end_val];

    Some((key.to_string(), val.to_string()))
}

fn extract_stream_id(line: &str) -> Option<String> {
    let pos = line.find("stream=")?;
    let rest = &line[pos + 7..];
    let end = rest
        .find(|c: char| !c.is_ascii_digit())
        .unwrap_or(rest.len());
    Some(rest[..end].to_string())
}

fn extract_data(line: &str) -> Option<String> {
    let pos = line.find("data=\"")?;
    let rest = &line[pos + 6..];
    // Find the closing quote — but beware of escaped quotes
    let mut end = 0;
    let chars: Vec<char> = rest.chars().collect();
    while end < chars.len() {
        if chars[end] == '"' && (end == 0 || chars[end - 1] != '\\') {
            break;
        }
        end += 1;
    }
    Some(rest[..end].to_string())
}

fn extract_data_len(line: &str) -> Option<usize> {
    let pos = line.find("len=")?;
    let rest = &line[pos + 4..];
    let end = rest
        .find(|c: char| !c.is_ascii_digit())
        .unwrap_or(rest.len());
    rest[..end].parse().ok()
}

fn decode_go_escaped(s: &str) -> Vec<u8> {
    let mut result = Vec::new();
    let bytes = s.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        if bytes[i] == b'\\' && i + 1 < bytes.len() {
            match bytes[i + 1] {
                b'x' if i + 3 < bytes.len() => {
                    if let Ok(b) = u8::from_str_radix(
                        std::str::from_utf8(&bytes[i + 2..i + 4]).unwrap_or(""),
                        16,
                    ) {
                        result.push(b);
                        i += 4;
                        continue;
                    }
                }
                b'n' => {
                    result.push(b'\n');
                    i += 2;
                    continue;
                }
                b'r' => {
                    result.push(b'\r');
                    i += 2;
                    continue;
                }
                b't' => {
                    result.push(b'\t');
                    i += 2;
                    continue;
                }
                b'\\' => {
                    result.push(b'\\');
                    i += 2;
                    continue;
                }
                b'"' => {
                    result.push(b'"');
                    i += 2;
                    continue;
                }
                _ => {}
            }
        }
        result.push(bytes[i]);
        i += 1;
    }
    result
}

fn extract_strings(data: &[u8]) -> Vec<String> {
    let mut strings = Vec::new();
    let mut current = Vec::new();
    for &b in data {
        if (32..=126).contains(&b) {
            current.push(b);
        } else {
            if current.len() >= 4 {
                if let Ok(s) = std::str::from_utf8(&current) {
                    if !strings.contains(&s.to_string()) {
                        strings.push(s.to_string());
                    }
                }
            }
            current.clear();
        }
    }
    if current.len() >= 4 {
        if let Ok(s) = std::str::from_utf8(&current) {
            if !strings.contains(&s.to_string()) {
                strings.push(s.to_string());
            }
        }
    }
    strings.truncate(30);
    strings
}

fn mask_token(key: &str, val: &str) -> String {
    if key == "authorization" && val.len() > 30 {
        if let Some(rest) = val.strip_prefix("Bearer ") {
            if rest.len() > 20 {
                return format!("Bearer {}...{}", &rest[..12], &rest[rest.len() - 8..]);
            }
        }
        if val.len() > 40 {
            return format!("{}...{}", &val[..25], &val[val.len() - 8..]);
        }
    }
    val.to_string()
}

fn decompress_gzip(data: &[u8]) -> Result<Vec<u8>, io::Error> {
    use flate2::read::GzDecoder;
    let mut decoder = GzDecoder::new(data);
    let mut out = Vec::new();
    decoder.read_to_end(&mut out)?;
    Ok(out)
}

// ── CLI entry point ──────────────────────────────────────────────────────────

pub fn run_cli() {
    let input = if let Some(path) = std::env::args().nth(1) {
        std::fs::read_to_string(&path).unwrap_or_else(|e| {
            eprintln!("Failed to read {path}: {e}");
            std::process::exit(1);
        })
    } else {
        let mut buf = String::new();
        io::stdin()
            .lock()
            .read_to_string(&mut buf)
            .expect("Failed to read stdin");
        buf
    };

    let snap = Snapshot::parse(&input);
    print!("{}", snap.render());
}