#!/usr/bin/env bash
# mitm-redirect.sh — UID-scoped iptables redirect for MITM interception
#
# Creates a dedicated system user for the standalone LS and adds an iptables
# rule that ONLY redirects traffic from that user's UID. No /etc/hosts
# modification, no system-wide changes.
#
# Flow:
#   1. Standalone LS runs as 'antigravity-ls' user (via sudo -u)
#   2. iptables catches :443 traffic from that UID only → REDIRECT to MITM port
#   3. MITM terminates TLS (Go client trusts our CA via SSL_CERT_FILE)
#   4. MITM forwards upstream, captures usage
#
# What this does NOT affect:
#   - Your real Antigravity session (different UID)
#   - Any other software on your PC (different UID)
#   - DNS resolution (no /etc/hosts changes)
#
# Usage:
#   sudo ./scripts/mitm-redirect.sh install [mitm_port]
#   sudo ./scripts/mitm-redirect.sh uninstall [mitm_port]
#   sudo ./scripts/mitm-redirect.sh status

set -euo pipefail

MITM_PORT="${2:-8742}"
LS_USER="antigravity-ls"
DATA_DIR="/tmp/antigravity-standalone"
LS_BINARY="/usr/share/antigravity/resources/app/extensions/antigravity/bin/language_server_linux_x64"
SUDOERS_FILE="/etc/sudoers.d/antigravity-ls"

install() {
    if [[ $EUID -ne 0 ]]; then
        echo "Error: must run as root (sudo)"
        exit 1
    fi

    echo "[mitm-redirect] Installing UID-scoped iptables redirect → :$MITM_PORT"
    echo

    # ── 1. Create system user ───────────────────────────────────────────
    if id "$LS_USER" &>/dev/null; then
        echo "  ✓ user '$LS_USER' already exists (uid=$(id -u "$LS_USER"))"
    else
        useradd -r -s /usr/sbin/nologin -d "$DATA_DIR" "$LS_USER"
        echo "  + created user '$LS_USER' (uid=$(id -u "$LS_USER"))"
    fi
    local LS_UID
    LS_UID=$(id -u "$LS_USER")

    # ── 2. Create data directory (writable by both users) ────────────────
    mkdir -p "$DATA_DIR/.gemini"
    chmod 1777 "$DATA_DIR" "$DATA_DIR/.gemini"
    echo "  + data dir: $DATA_DIR (mode 1777, writable by all)"

    # ── 3. Sudoers entry ────────────────────────────────────────────────
    # Allow the invoking user (SUDO_USER) to run ANY command as antigravity-ls.
    # This is needed for the proxy to spawn the LS binary.
    local REAL_USER="${SUDO_USER:-$(logname 2>/dev/null || whoami)}"
    cat > "$SUDOERS_FILE" <<EOF
# Allow $REAL_USER to run commands as $LS_USER (for antigravity proxy)
$REAL_USER ALL=($LS_USER) NOPASSWD: ALL
# Allow $REAL_USER to kill $LS_USER's processes (for clean shutdown)
$REAL_USER ALL=(root) NOPASSWD: /usr/bin/kill -TERM *, /usr/bin/kill -KILL *, /usr/bin/pkill -TERM -u $LS_USER *, /usr/bin/pkill -KILL -u $LS_USER *
EOF
    chmod 440 "$SUDOERS_FILE"
    echo "  + sudoers: $REAL_USER can run as $LS_USER + kill $LS_USER processes"

    # ── 4. iptables REDIRECT (scoped to UID) ────────────────────────────
    # Remove existing rule first (idempotent)
    iptables -t nat -D OUTPUT -m owner --uid-owner "$LS_UID" \
        -p tcp --dport 443 -j REDIRECT --to-port "$MITM_PORT" 2>/dev/null || true

    iptables -t nat -A OUTPUT -m owner --uid-owner "$LS_UID" \
        -p tcp --dport 443 -j REDIRECT --to-port "$MITM_PORT"
    echo "  + iptables: uid=$LS_UID :443 → :$MITM_PORT"

    echo
    echo "[mitm-redirect] ✓ Installed (only affects uid=$LS_UID)"
    echo "  Restart the proxy to take effect:"
    echo "    RUST_LOG=info ./target/release/antigravity-proxy --standalone"
}

uninstall() {
    if [[ $EUID -ne 0 ]]; then
        echo "Error: must run as root (sudo)"
        exit 1
    fi

    echo "[mitm-redirect] Removing UID-scoped iptables redirect"
    echo

    # Remove iptables rule
    if id "$LS_USER" &>/dev/null; then
        local LS_UID
        LS_UID=$(id -u "$LS_USER")
        iptables -t nat -D OUTPUT -m owner --uid-owner "$LS_UID" \
            -p tcp --dport 443 -j REDIRECT --to-port "$MITM_PORT" 2>/dev/null || true
        echo "  - iptables: removed REDIRECT rule for uid=$LS_UID"
    fi

    # Remove sudoers entry
    rm -f "$SUDOERS_FILE"
    echo "  - sudoers: removed $SUDOERS_FILE"

    # Clean data dir
    rm -rf "$DATA_DIR"
    echo "  - data dir: removed $DATA_DIR"

    # Optionally remove user (commented out — user might want to keep it)
    # userdel "$LS_USER" 2>/dev/null || true
    echo "  ℹ user '$LS_USER' kept (run 'sudo userdel $LS_USER' to remove)"

    echo
    echo "[mitm-redirect] ✓ Uninstalled."
}

status() {
    echo "[mitm-redirect] Status"
    echo

    # Check user
    if id "$LS_USER" &>/dev/null; then
        local LS_UID
        LS_UID=$(id -u "$LS_USER")
        echo "  user: $LS_USER (uid=$LS_UID) ✓"
    else
        echo "  user: $LS_USER (not found) ✗"
        echo
        echo "  Run: sudo $0 install"
        return
    fi

    # Check sudoers
    if [[ -f "$SUDOERS_FILE" ]]; then
        echo "  sudoers: $SUDOERS_FILE ✓"
    else
        echo "  sudoers: $SUDOERS_FILE (not found) ✗"
    fi

    # Check iptables
    echo "  iptables:"
    if iptables -t nat -L OUTPUT -n 2>/dev/null | grep -q "owner UID match.*$LS_UID"; then
        iptables -t nat -L OUTPUT -n -v 2>/dev/null | grep "owner UID" | sed 's/^/    /'
    else
        echo "    (no rules for uid=$LS_UID)"
    fi

    # Check data dir
    echo "  data dir: $(ls -ld "$DATA_DIR" 2>/dev/null || echo '(not found)')"

    # Test sudo
    echo
    echo "  sudo test:"
    if sudo -n -u "$LS_USER" true 2>/dev/null; then
        echo "    ✓ can run as $LS_USER without password"
    else
        echo "    ✗ cannot run as $LS_USER (check sudoers)"
    fi
}

case "${1:-help}" in
    install)   install ;;
    uninstall) uninstall ;;
    status)    status ;;
    *)
        echo "Usage: sudo $0 {install|uninstall|status} [mitm_port]"
        echo
        echo "Redirects ONLY the standalone LS's outgoing :443 traffic through"
        echo "the MITM proxy using UID-scoped iptables rules."
        echo
        echo "This does NOT affect:"
        echo "  - Your real Antigravity coding session"
        echo "  - Any other software on your PC"
        echo "  - DNS resolution (/etc/hosts is untouched)"
        echo
        echo "  install [port]    Create user + iptables REDIRECT for that UID"
        echo "  uninstall [port]  Remove iptables rule + sudoers"
        echo "  status            Show current state"
        echo
        echo "Default MITM port: 8742"
        ;;
esac