#!/usr/bin/env bash
# ╔═══════════════════════════════════════════════════════════════════════════╗
# ║  Antigravity MITM — DNS-based redirect for targeted interception         ║
# ║                                                                          ║
# ║  Instead of redirecting ALL port 443 traffic (which breaks everything),  ║
# ║  this uses /etc/hosts to redirect ONLY the LLM API domain to localhost,  ║
# ║  then iptables redirects only localhost:443 → MITM port.                 ║
# ║                                                                          ║
# ║  Also adds the MITM CA to the system trust store so Go trusts it.        ║
# ╚═══════════════════════════════════════════════════════════════════════════╝
set -euo pipefail

MITM_PORT="${ANTIGRAVITY_MITM_PORT:-8742}"
MITM_CA="${HOME}/.config/antigravity-proxy/mitm-ca.pem"
# If run with sudo, use SUDO_USER's home
if [[ -n "${SUDO_USER:-}" ]]; then
    MITM_CA="$(eval echo "~${SUDO_USER}")/.config/antigravity-proxy/mitm-ca.pem"
fi

HOSTS_MARKER="# antigravity-mitm"
API_DOMAINS=(
    "daily-cloudcode-pa.googleapis.com"
)

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

cmd_install() {
    echo -e "${BOLD}${CYAN}Antigravity MITM DNS Redirect Setup${NC}"
    echo -e "────────────────────────────────────"
    echo ""

    # 1. Add MITM CA to system trust store
    if [[ ! -f "$MITM_CA" ]]; then
        echo -e "  ${RED}✗${NC} MITM CA not found: ${MITM_CA}"
        echo -e "    Start the proxy once first to generate it."
        exit 1
    fi

    local sys_cert="/usr/local/share/ca-certificates/antigravity-mitm.crt"
    cp "$MITM_CA" "$sys_cert"
    update-ca-certificates >/dev/null 2>&1
    echo -e "  ${GREEN}✓${NC} MITM CA added to system trust store"

    # 2. Add /etc/hosts entries for API domains → 127.0.0.1
    # First, cache the real IPs for the MITM to use later
    local real_ips_file="/tmp/antigravity-mitm-real-ips"
    > "$real_ips_file"

    for domain in "${API_DOMAINS[@]}"; do
        # Remove old entries
        sed -i "/${domain}.*${HOSTS_MARKER}/d" /etc/hosts

        # Resolve and cache the real IPs BEFORE redirecting
        local real_ip
        real_ip=$(dig +short "$domain" 2>/dev/null | grep -E '^[0-9]+\.' | head -1)
        if [[ -n "$real_ip" ]]; then
            echo "${domain}=${real_ip}" >> "$real_ips_file"
        fi

        # Add the /etc/hosts redirect
        echo "127.0.0.1 ${domain} ${HOSTS_MARKER}" >> /etc/hosts
        echo -e "  ${GREEN}✓${NC} /etc/hosts: ${domain} → 127.0.0.1 (real: ${real_ip:-unknown})"
    done

    # 3. iptables: redirect ONLY 127.0.0.1:443 → MITM port
    # This catches only the /etc/hosts redirected domains, nothing else!
    iptables -t nat -D OUTPUT -d 127.0.0.1 -p tcp --dport 443 \
        -j REDIRECT --to-port "$MITM_PORT" 2>/dev/null || true
    iptables -t nat -A OUTPUT -d 127.0.0.1 -p tcp --dport 443 \
        -j REDIRECT --to-port "$MITM_PORT"
    echo -e "  ${GREEN}✓${NC} iptables: 127.0.0.1:443 → localhost:${MITM_PORT}"

    echo ""
    echo -e "  ${GREEN}Done!${NC}"
    echo ""
    echo -e "  ${BOLD}How it works:${NC}"
    echo -e "    1. LS resolves ${API_DOMAINS[0]} → 127.0.0.1 (via /etc/hosts)"
    echo -e "    2. LS connects to 127.0.0.1:443"
    echo -e "    3. iptables redirects to MITM proxy on :${MITM_PORT}"
    echo -e "    4. MITM intercepts, decrypts (CA is trusted), proxies to real Google"
    echo ""
    echo -e "  Real upstream IPs cached in: ${real_ips_file}"
    echo -e "  Restart Antigravity for changes to take effect."
    echo -e "  Undo: sudo $0 uninstall"
    echo ""
}

cmd_uninstall() {
    echo -e "${BOLD}${CYAN}Removing MITM DNS Redirect${NC}"
    echo ""

    # Remove /etc/hosts entries
    sed -i "/${HOSTS_MARKER}/d" /etc/hosts
    echo -e "  ${GREEN}✓${NC} Removed /etc/hosts entries"

    # Remove iptables rule
    iptables -t nat -D OUTPUT -d 127.0.0.1 -p tcp --dport 443 \
        -j REDIRECT --to-port "$MITM_PORT" 2>/dev/null || true
    echo -e "  ${GREEN}✓${NC} Removed iptables rule"

    # Remove system CA (optional)
    rm -f /usr/local/share/ca-certificates/antigravity-mitm.crt
    update-ca-certificates >/dev/null 2>&1
    echo -e "  ${GREEN}✓${NC} Removed MITM CA from system trust store"

    echo ""
}

cmd_status() {
    echo -e "${BOLD}${CYAN}MITM DNS Redirect Status${NC}"
    echo ""

    # Check /etc/hosts
    local hosts_count
    hosts_count=$(grep -c "$HOSTS_MARKER" /etc/hosts 2>/dev/null || echo 0)
    if [[ "$hosts_count" -gt 0 ]]; then
        echo -e "  ${GREEN}✓${NC} /etc/hosts: ${hosts_count} domain(s) redirected"
        grep "$HOSTS_MARKER" /etc/hosts | sed 's/^/    /'
    else
        echo -e "  ${YELLOW}○${NC} /etc/hosts: no redirects"
    fi

    echo ""
    # Check iptables
    if iptables -t nat -L OUTPUT -n 2>/dev/null | grep -q "127.0.0.1.*REDIRECT.*${MITM_PORT}"; then
        echo -e "  ${GREEN}✓${NC} iptables: 127.0.0.1:443 → :${MITM_PORT}"
    else
        echo -e "  ${YELLOW}○${NC} iptables: no redirect"
    fi

    echo ""
    # Check system CA
    if [[ -f /usr/local/share/ca-certificates/antigravity-mitm.crt ]]; then
        echo -e "  ${GREEN}✓${NC} System CA: installed"
    else
        echo -e "  ${YELLOW}○${NC} System CA: not installed"
    fi
    echo ""
}

case "${1:-}" in
    install)
        cmd_install
        ;;
    uninstall)
        cmd_uninstall
        ;;
    status)
        cmd_status
        ;;
    *)
        echo "Usage: sudo $0 {install|uninstall|status}"
        echo ""
        echo "Redirects LLM API domain to localhost via /etc/hosts + iptables."
        echo "Only intercepts API traffic, everything else is untouched."
        exit 1
        ;;
esac