#!/usr/bin/env bash
# ╔═══════════════════════════════════════════════════════════════════════════╗
# ║  Antigravity MITM — iptables redirect for transparent interception       ║
# ║                                                                          ║
# ║  Redirects outbound port 443 traffic to the MITM proxy.                  ║
# ║  Uses a dedicated GID to exclude the proxy's own upstream traffic,       ║
# ║  preventing redirect loops.                                              ║
# ║                                                                          ║
# ║  Usage: sudo ./iptables-redirect.sh install                              ║
# ║         sudo ./iptables-redirect.sh uninstall                            ║
# ║         sudo ./iptables-redirect.sh status                               ║
# ╚═══════════════════════════════════════════════════════════════════════════╝
set -euo pipefail

MITM_PORT="${ANTIGRAVITY_MITM_PORT:-8742}"
CHAIN="ANTIGRAVITY_MITM"
BYPASS_GROUP="mitm-bypass"

# Resolve target user (the one whose traffic we redirect)
TARGET_USER="${SUDO_USER:-$(whoami)}"
TARGET_UID=$(id -u "$TARGET_USER" 2>/dev/null || echo "")

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 iptables Setup${NC}"
    echo -e "────────────────────────────────"
    echo ""

    if [[ -z "$TARGET_UID" ]]; then
        echo -e "  ${RED}✗${NC} Cannot resolve UID for user '${TARGET_USER}'"
        exit 1
    fi

    # Create bypass group (proxy runs with this GID to avoid redirect loop)
    if ! getent group "$BYPASS_GROUP" >/dev/null 2>&1; then
        groupadd "$BYPASS_GROUP"
        echo -e "  ${GREEN}✓${NC} Created group: ${BYPASS_GROUP}"
    else
        echo -e "  ${GREEN}✓${NC} Group exists: ${BYPASS_GROUP}"
    fi

    # Add user to bypass group (so they can use 'sg' to run proxy)
    if ! id -nG "$TARGET_USER" 2>/dev/null | grep -qw "$BYPASS_GROUP"; then
        usermod -aG "$BYPASS_GROUP" "$TARGET_USER"
        echo -e "  ${GREEN}✓${NC} Added ${TARGET_USER} to ${BYPASS_GROUP}"
    fi

    local bypass_gid
    bypass_gid=$(getent group "$BYPASS_GROUP" | cut -d: -f3)

    # Check MITM proxy is running
    if ! ss -tlnp 2>/dev/null | grep -q ":${MITM_PORT}"; then
        echo -e "  ${YELLOW}!${NC} MITM proxy not running on :${MITM_PORT} (will work once started)"
    else
        echo -e "  ${GREEN}✓${NC} MITM proxy listening on :${MITM_PORT}"
    fi

    # Create our chain
    iptables -t nat -N "$CHAIN" 2>/dev/null || true
    iptables -t nat -F "$CHAIN"

    # THE KEY RULE: redirect port 443 traffic UNLESS it's from the bypass group.
    # This prevents redirect loops — the proxy runs with GID=mitm-bypass,
    # so its upstream connections to Google are NOT redirected back to itself.
    iptables -t nat -A "$CHAIN" \
        -m owner ! --gid-owner "$bypass_gid" \
        -p tcp --dport 443 \
        -j REDIRECT --to-port "$MITM_PORT"

    echo -e "  ${GREEN}✓${NC} Redirect rule: tcp/443 → :${MITM_PORT} (skip GID ${bypass_gid})"

    # Hook into OUTPUT for target user only
    iptables -t nat -D OUTPUT -m owner --uid-owner "$TARGET_UID" \
        -p tcp --dport 443 -j "$CHAIN" 2>/dev/null || true
    iptables -t nat -A OUTPUT -m owner --uid-owner "$TARGET_UID" \
        -p tcp --dport 443 -j "$CHAIN"

    echo -e "  ${GREEN}✓${NC} OUTPUT hook: UID ${TARGET_UID} (${TARGET_USER})"

    echo ""
    echo -e "  ${GREEN}Done!${NC}"
    echo ""
    echo -e "  ${BOLD}IMPORTANT:${NC} Run the proxy with the bypass group to avoid loops:"
    echo -e "    ${CYAN}sg ${BYPASS_GROUP} -c 'RUST_LOG=info ./target/release/antigravity-proxy'${NC}"
    echo ""
    echo -e "  Then restart Antigravity to re-establish connections."
    echo -e "  Undo: sudo $0 uninstall"
    echo ""
}

cmd_uninstall() {
    echo -e "${BOLD}${CYAN}Removing iptables MITM redirect${NC}"
    echo ""

    local target_uid
    target_uid=$(id -u "$TARGET_USER" 2>/dev/null || echo "1000")

    # Remove jump from OUTPUT
    iptables -t nat -D OUTPUT -m owner --uid-owner "$target_uid" \
        -p tcp --dport 443 -j "$CHAIN" 2>/dev/null || true
    echo -e "  ${GREEN}✓${NC} Removed OUTPUT jump"

    # Flush and delete our chain
    iptables -t nat -F "$CHAIN" 2>/dev/null || true
    iptables -t nat -X "$CHAIN" 2>/dev/null || true
    echo -e "  ${GREEN}✓${NC} Removed ${CHAIN} chain"

    echo ""
    echo -e "  ${YELLOW}Note:${NC} Group '${BYPASS_GROUP}' left intact (harmless)."
    echo -e "        Remove with: sudo groupdel ${BYPASS_GROUP}"
    echo ""
}

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

    if iptables -t nat -L "$CHAIN" -n 2>/dev/null | grep -q "REDIRECT"; then
        echo -e "  ${GREEN}✓${NC} Chain ${CHAIN}: active"
        iptables -t nat -L "$CHAIN" -nv --line-numbers 2>/dev/null | \
            sed 's/^/    /'
    else
        echo -e "  ${YELLOW}○${NC} Chain ${CHAIN}: not installed"
    fi

    echo ""
    if iptables -t nat -L OUTPUT -n 2>/dev/null | grep -q "$CHAIN"; then
        echo -e "  ${GREEN}✓${NC} OUTPUT hook: installed"
        iptables -t nat -L OUTPUT -n 2>/dev/null | grep "$CHAIN" | sed 's/^/    /'
    else
        echo -e "  ${YELLOW}○${NC} OUTPUT hook: not installed"
    fi

    echo ""
    if getent group "$BYPASS_GROUP" >/dev/null 2>&1; then
        local gid
        gid=$(getent group "$BYPASS_GROUP" | cut -d: -f3)
        echo -e "  ${GREEN}✓${NC} Bypass group: ${BYPASS_GROUP} (GID ${gid})"
    else
        echo -e "  ${YELLOW}○${NC} Bypass group: not created"
    fi
    echo ""
}

case "${1:-}" in
    install)
        cmd_install
        ;;
    uninstall)
        cmd_uninstall
        ;;
    status)
        cmd_status
        ;;
    *)
        echo "Usage: sudo $0 {install|uninstall|status}"
        echo ""
        echo "Redirects outbound port 443 traffic to the MITM proxy."
        echo "The proxy must be run with 'sg mitm-bypass' to avoid redirect loops."
        exit 1
        ;;
esac