#!/bin/bash
# Tor Browser Container Manager
# Usage: tor-browser.sh [create|start|stop|launch|status|destroy|shell]

set -e

CONTAINER="tor-browser"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SPEC_FILE="$SCRIPT_DIR/../specs/tor-browser.yaml"
DOWNLOADS_DIR="$HOME/Downloads/tor-browser"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

log() { echo -e "${GREEN}[tor-browser]${NC} $1"; }
warn() { echo -e "${YELLOW}[tor-browser]${NC} $1"; }
error() { echo -e "${RED}[tor-browser]${NC} $1"; }

container_exists() {
    sudo machinectl list-images 2>/dev/null | grep -q "^$CONTAINER "
}

container_running() {
    sudo machinectl list 2>/dev/null | grep -q "^$CONTAINER "
}

wait_for_boot() {
    log "Waiting for container to boot..."
    for i in {1..30}; do
        if sudo machinectl shell "$CONTAINER" /bin/true 2>/dev/null; then
            return 0
        fi
        sleep 1
    done
    error "Container failed to boot"
    return 1
}

setup_network() {
    log "Configuring network..."
    sudo machinectl shell "$CONTAINER" /bin/bash -c '
        # Create network config if not exists
        if [ ! -f /etc/systemd/network/80-container-host0.network ]; then
            cat > /etc/systemd/network/80-container-host0.network << EOF
[Match]
Name=host0

[Network]
DHCP=yes
EOF
        fi
        systemctl enable --now systemd-networkd 2>/dev/null || true
    ' 2>/dev/null

    # Wait for IP
    for i in {1..30}; do
        IP=$(sudo machinectl show "$CONTAINER" -p Addresses --value 2>/dev/null | grep -oE '192\.[0-9]+\.[0-9]+\.[0-9]+' | head -1)
        if [ -n "$IP" ]; then
            log "Container IP: $IP"
            return 0
        fi
        sleep 1
    done
    warn "Could not get container IP"
}

setup_tor() {
    log "Setting up Tor inside container..."
    sudo machinectl shell "$CONTAINER" /bin/bash -c '
        # Install Tor if not present
        if ! command -v tor &>/dev/null; then
            apt-get update
            apt-get install -y tor iptables curl
        fi

        # Configure Tor
        cat > /etc/tor/torrc << EOF
SocksPort 9050
DNSPort 5353
TransPort 9040 IsolateClientAddr IsolateClientProtocol IsolateDestAddr IsolateDestPort
AutomapHostsOnResolve 1
VirtualAddrNetworkIPv4 10.192.0.0/10
ExitNodes {SE},{DE}
StrictNodes 1
EOF

        # Configure transparent proxy iptables rules
        TOR_UID=$(id -u debian-tor)

        # NAT rules for transparent proxy
        iptables -t nat -F
        iptables -t nat -A OUTPUT -p udp --dport 53 -j REDIRECT --to-ports 5353
        iptables -t nat -A OUTPUT -m owner --uid-owner $TOR_UID -j RETURN
        iptables -t nat -A OUTPUT -p tcp --syn -j REDIRECT --to-ports 9040

        # Kill switch and logging rules (filter table)
        iptables -N LEAK_CHECK 2>/dev/null || iptables -F LEAK_CHECK

        # Allow loopback
        iptables -A LEAK_CHECK -o lo -j ACCEPT

        # Allow Tor daemon traffic (it needs to reach the network)
        iptables -A LEAK_CHECK -m owner --uid-owner $TOR_UID -j ACCEPT

        # Allow established connections (for Tor circuits)
        iptables -A LEAK_CHECK -m state --state ESTABLISHED,RELATED -j ACCEPT

        # Allow traffic to localhost (where Tor proxy listens)
        iptables -A LEAK_CHECK -d 127.0.0.0/8 -j ACCEPT

        # Log potential leaks before dropping
        iptables -A LEAK_CHECK -j LOG --log-prefix "TOR-LEAK: " --log-level 4

        # Drop everything else (kill switch)
        iptables -A LEAK_CHECK -j DROP

        # Insert the chain into OUTPUT
        iptables -C OUTPUT -j LEAK_CHECK 2>/dev/null || iptables -I OUTPUT 1 -j LEAK_CHECK

        # Save iptables rules (both nat and filter)
        mkdir -p /etc/iptables
        iptables-save > /etc/iptables/rules.v4

        # Create restore service
        cat > /etc/systemd/system/iptables-restore.service << EOF
[Unit]
Description=Restore iptables rules
Before=network-pre.target
Wants=network-pre.target

[Service]
Type=oneshot
ExecStart=/sbin/iptables-restore /etc/iptables/rules.v4
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
EOF

        systemctl daemon-reload
        systemctl enable iptables-restore
        systemctl enable tor
        systemctl restart tor@default
    ' 2>/dev/null

    log "Waiting for Tor to bootstrap..."
    for i in {1..60}; do
        if sudo machinectl shell "$CONTAINER" /bin/bash -c 'curl -s --connect-timeout 5 https://check.torproject.org/api/ip 2>/dev/null | grep -q IsTor' 2>/dev/null; then
            log "Tor is ready!"
            return 0
        fi
        sleep 2
    done
    warn "Tor bootstrap timeout - may still be connecting"
}

setup_firefox() {
    log "Installing Firefox and graphics libraries..."
    sudo machinectl shell "$CONTAINER" /bin/bash -c '
        if ! command -v firefox-esr &>/dev/null; then
            apt-get update
            apt-get install -y firefox-esr libegl1 libgl1 libgbm1
        fi

        # Create Firefox home directory and profile
        mkdir -p /tmp/firefox-home

        # Create default profile directory
        PROFILE_DIR="/tmp/firefox-home/.mozilla/firefox/tor.default"
        mkdir -p "$PROFILE_DIR/chrome"

        # Create profiles.ini
        cat > /tmp/firefox-home/.mozilla/firefox/profiles.ini << PROFILES
[Profile0]
Name=default
IsRelative=1
Path=tor.default
Default=1

[General]
StartWithLastProfile=1
PROFILES

        # Create userChrome.css for purple/onion themed title bar
        cat > "$PROFILE_DIR/chrome/userChrome.css" << CHROME
/* Tor Browser - Purple/Onion Theme Title Bar */
@namespace url("http://www.mozilla.org/keymaster/gatekeeper/there.is.only.xul");

/* Main window title bar - Tor purple */
#navigator-toolbox {
    background: linear-gradient(to bottom, #7D4698 0%, #5E3370 100%) !important;
}

/* Tab bar background */
#titlebar, #TabsToolbar {
    background: #5E3370 !important;
}

/* Navigation bar */
#nav-bar {
    background: linear-gradient(to bottom, #4A2858 0%, #3D1F4A 100%) !important;
    border-bottom: 2px solid #9B59B6 !important;
}

/* URL bar */
#urlbar-background {
    background-color: #2C1A33 !important;
    border: 1px solid #9B59B6 !important;
}

#urlbar-input {
    color: #E8D5F0 !important;
}

/* Tab styling - unselected tabs */
.tabbrowser-tab {
    background: #3D1F4A !important;
}

.tabbrowser-tab .tab-background {
    background: #3D1F4A !important;
}

/* Selected tab - lighter purple with dark text */
.tabbrowser-tab[selected="true"] {
    background: #9B59B6 !important;
}

.tabbrowser-tab[selected="true"] .tab-background {
    background: #9B59B6 !important;
}

/* Tab text - light for unselected, dark for selected */
.tabbrowser-tab .tab-label {
    color: #D4B8E0 !important;
}

.tabbrowser-tab[selected="true"] .tab-label {
    color: #1A0D20 !important;
    font-weight: bold !important;
}

/* Tab close button */
.tabbrowser-tab .tab-close-button {
    color: #D4B8E0 !important;
}

.tabbrowser-tab[selected="true"] .tab-close-button {
    color: #1A0D20 !important;
}

/* Toolbar buttons */
#nav-bar toolbarbutton {
    color: #E8D5F0 !important;
}

/* Bookmarks bar */
#PersonalToolbar {
    background: #3D1F4A !important;
}

/* Menu bar */
#toolbar-menubar {
    background: #5E3370 !important;
    color: white !important;
}
CHROME

        # Enable userChrome.css in Firefox preferences
        cat > "$PROFILE_DIR/user.js" << USERJS
// Enable userChrome.css
user_pref("toolkit.legacyUserProfileCustomizations.stylesheets", true);

// Privacy settings for Tor
user_pref("privacy.resistFingerprinting", true);
user_pref("privacy.trackingprotection.enabled", true);
user_pref("media.peerconnection.enabled", false);
user_pref("geo.enabled", false);
user_pref("dom.battery.enabled", false);

// Disable telemetry
user_pref("toolkit.telemetry.enabled", false);
user_pref("datareporting.healthreport.uploadEnabled", false);

// Homepage
user_pref("browser.startup.homepage", "https://check.torproject.org");

// Show title bar (for the colored theme)
user_pref("browser.tabs.inTitlebar", 0);

// Disable DNS over HTTPS - Tor handles DNS
user_pref("network.trr.mode", 5);
user_pref("network.trr.uri", "");
user_pref("doh-rollout.disable-heuristics", true);
USERJS

        echo "Firefox profile configured with Tor theme"
    ' 2>/dev/null
}

create_nspawn_config() {
    log "Creating nspawn configuration..."
    sudo tee /etc/systemd/nspawn/$CONTAINER.nspawn > /dev/null << EOF
[Exec]
Boot=yes
PrivateUsers=no
Environment=DISPLAY=:0
Environment=WAYLAND_DISPLAY=wayland-0
Environment=XDG_RUNTIME_DIR=/tmp
Environment=MOZ_ENABLE_WAYLAND=1
Environment=MOZ_WEBRENDER=1
Environment=MOZ_DISABLE_RTC=1

[Network]
VirtualEthernet=yes
Bridge=br0

[Files]
Bind=/tmp/.X11-unix:/tmp/.X11-unix
Bind=/run/user/1000/wayland-0:/tmp/wayland-0
Bind=$DOWNLOADS_DIR:/home/browser/Downloads
EOF
}

cmd_create() {
    if container_exists; then
        error "Container already exists. Use 'destroy' first or 'start' to run it."
        exit 1
    fi

    # Create downloads directory
    mkdir -p "$DOWNLOADS_DIR"

    log "Creating container from spec..."

    # Find gaol binary
    GAOL_BIN=$(which gaol)
    if [ -z "$GAOL_BIN" ]; then
        error "gaol not found in PATH"
        exit 1
    fi

    sudo "$GAOL_BIN" apply "$SPEC_FILE"

    # Update nspawn config with correct settings
    create_nspawn_config

    log "Starting container..."
    sudo machinectl start "$CONTAINER"
    wait_for_boot

    setup_network
    setup_tor
    setup_firefox

    log "Container created and configured!"
    log "Run '$0 launch' to start Firefox"
}

cmd_start() {
    if ! container_exists; then
        error "Container does not exist. Run 'create' first."
        exit 1
    fi

    if container_running; then
        warn "Container already running"
        return 0
    fi

    log "Starting container..."
    sudo machinectl start "$CONTAINER"
    wait_for_boot
    setup_network

    # Restore iptables and restart Tor
    sudo machinectl shell "$CONTAINER" /bin/bash -c '
        /sbin/iptables-restore /etc/iptables/rules.v4 2>/dev/null || true
        systemctl restart tor@default
    ' 2>/dev/null

    log "Waiting for Tor..."
    for i in {1..30}; do
        if sudo machinectl shell "$CONTAINER" /bin/bash -c 'curl -s --connect-timeout 5 https://check.torproject.org/api/ip 2>/dev/null | grep -q IsTor' 2>/dev/null; then
            log "Container started and Tor is ready!"
            return 0
        fi
        sleep 2
    done
    warn "Tor may still be bootstrapping"
}

cmd_stop() {
    if ! container_running; then
        warn "Container not running"
        return 0
    fi

    log "Stopping container..."
    sudo machinectl stop "$CONTAINER"
    log "Container stopped"
}

cmd_launch() {
    if ! container_running; then
        log "Container not running, starting..."
        cmd_start
    fi

    log "Launching Firefox with Tor theme..."
    sudo machinectl shell "$CONTAINER" /bin/bash -c '
        # Ensure browser user exists with UID 1000 (matches host user for Wayland access)
        if ! id -u browser 2>/dev/null; then
            useradd -u 1000 -m -s /bin/bash browser 2>/dev/null || true
        fi

        # Kill existing Firefox if any
        pkill -9 firefox-esr 2>/dev/null || true
        sleep 1

        # Ensure profile exists
        PROFILE_DIR="/home/browser/.mozilla/firefox/tor.default"
        if [ ! -f "$PROFILE_DIR/user.js" ]; then
            mkdir -p "$PROFILE_DIR/chrome"

            cat > /home/browser/.mozilla/firefox/profiles.ini << PROFILES
[Profile0]
Name=default
IsRelative=1
Path=tor.default
Default=1

[General]
StartWithLastProfile=1
PROFILES

            cat > "$PROFILE_DIR/chrome/userChrome.css" << CHROME
@namespace url("http://www.mozilla.org/keymaster/gatekeeper/there.is.only.xul");
#navigator-toolbox { background: linear-gradient(to bottom, #7D4698 0%, #5E3370 100%) !important; }
#titlebar, #TabsToolbar { background: #5E3370 !important; }
#nav-bar { background: linear-gradient(to bottom, #4A2858 0%, #3D1F4A 100%) !important; border-bottom: 2px solid #9B59B6 !important; }
/* URL bar - light background with dark text for readability */
#urlbar-background { background-color: #e8d5f0 !important; border: 2px solid #9B59B6 !important; }
#urlbar { background-color: #e8d5f0 !important; color: #1A0D20 !important; }
#urlbar * { color: #1A0D20 !important; }
#urlbar-input { color: #1A0D20 !important; }
#urlbar-input-container { color: #1A0D20 !important; }
#identity-box { color: #1A0D20 !important; }
.tabbrowser-tab { background: #3D1F4A !important; }
.tabbrowser-tab .tab-background { background: #3D1F4A !important; }
.tabbrowser-tab[selected="true"] { background: #9B59B6 !important; }
.tabbrowser-tab[selected="true"] .tab-background { background: #9B59B6 !important; }
.tabbrowser-tab .tab-label { color: #D4B8E0 !important; }
.tabbrowser-tab[selected="true"] .tab-label { color: #1A0D20 !important; font-weight: bold !important; }
#nav-bar toolbarbutton { color: #E8D5F0 !important; }
#PersonalToolbar { background: #3D1F4A !important; }
#toolbar-menubar { background: #5E3370 !important; color: white !important; }
CHROME

            cat > "$PROFILE_DIR/user.js" << USERJS
user_pref("toolkit.legacyUserProfileCustomizations.stylesheets", true);
user_pref("privacy.resistFingerprinting", true);
user_pref("privacy.trackingprotection.enabled", true);
user_pref("media.peerconnection.enabled", false);
user_pref("geo.enabled", false);
user_pref("dom.battery.enabled", false);
user_pref("toolkit.telemetry.enabled", false);
user_pref("datareporting.healthreport.uploadEnabled", false);
user_pref("browser.startup.homepage", "https://check.torproject.org");
user_pref("browser.tabs.inTitlebar", 0);

// Disable DNS over HTTPS - Tor handles DNS
user_pref("network.trr.mode", 5);
user_pref("network.trr.uri", "");
user_pref("doh-rollout.disable-heuristics", true);
USERJS
            chown -R browser:browser /home/browser/.mozilla
        fi

        # Fix dconf permissions
        mkdir -p /tmp/dconf
        chmod 777 /tmp/dconf

        # Launch Firefox as browser user (UID 1000 matches host for Wayland socket access)
        runuser -u browser -- bash -c "
            export WAYLAND_DISPLAY=wayland-0
            export XDG_RUNTIME_DIR=/tmp
            export XDG_CACHE_HOME=/home/browser/.cache
            export XDG_CONFIG_HOME=/home/browser/.config
            export MOZ_ENABLE_WAYLAND=1
            export NO_AT_BRIDGE=1
            export LIBGL_ALWAYS_SOFTWARE=1
            export HOME=/home/browser
            mkdir -p /home/browser/.cache /home/browser/.config
            setsid firefox-esr --profile /home/browser/.mozilla/firefox/tor.default https://check.torproject.org </dev/null >/dev/null 2>&1 &
        "
        echo "Firefox launched with Tor profile"
    ' 2>/dev/null

    log "Firefox launched with purple Tor theme - check your desktop"
}

cmd_status() {
    echo "=== Container Status ==="
    if container_exists; then
        echo "Image: EXISTS"
    else
        echo "Image: NOT FOUND"
        return
    fi

    if container_running; then
        echo "State: RUNNING"
        IP=$(sudo machinectl show "$CONTAINER" -p Addresses --value 2>/dev/null | head -1)
        echo "IP: $IP"

        # Check Tor
        if sudo machinectl shell "$CONTAINER" /bin/bash -c 'systemctl is-active tor@default' 2>/dev/null | grep -q active; then
            echo "Tor: RUNNING"
            TOR_IP=$(sudo machinectl shell "$CONTAINER" /bin/bash -c 'curl -s --connect-timeout 5 https://check.torproject.org/api/ip 2>/dev/null' 2>/dev/null | grep -oP '"IP":"[^"]+"' | cut -d'"' -f4)
            if [ -n "$TOR_IP" ]; then
                echo "Tor Exit IP: $TOR_IP"
            fi
        else
            echo "Tor: STOPPED"
        fi

        # Check Firefox
        if sudo machinectl shell "$CONTAINER" /bin/bash -c 'pgrep firefox-esr' 2>/dev/null | grep -q .; then
            echo "Firefox: RUNNING"
        else
            echo "Firefox: NOT RUNNING"
        fi
    else
        echo "State: STOPPED"
    fi
}

cmd_destroy() {
    if container_running; then
        log "Stopping container..."
        sudo machinectl stop "$CONTAINER" 2>/dev/null || true
        sleep 2
    fi

    if container_exists; then
        log "Removing container..."
        sudo machinectl remove "$CONTAINER" 2>/dev/null || true
        sudo rm -f /etc/systemd/nspawn/$CONTAINER.nspawn
        log "Container destroyed"
    else
        warn "Container does not exist"
    fi
}

cmd_shell() {
    if ! container_running; then
        error "Container not running. Run 'start' first."
        exit 1
    fi

    sudo machinectl shell "$CONTAINER"
}

cmd_help() {
    cat << EOF
Tor Browser Container Manager

Usage: $0 <command>

Commands:
  create   Create and configure the Tor browser container
  start    Start the container (if stopped)
  stop     Stop the container
  launch   Launch Firefox in the container (starts if needed)
  status   Show container and Tor status
  destroy  Remove the container completely
  shell    Open a shell in the container
  help     Show this help message

The container routes ALL traffic through Tor with exit nodes
restricted to Sweden (SE) and Germany (DE).

Downloads are saved to: $DOWNLOADS_DIR
EOF
}

# Main
case "${1:-help}" in
    create)  cmd_create ;;
    start)   cmd_start ;;
    stop)    cmd_stop ;;
    launch)  cmd_launch ;;
    status)  cmd_status ;;
    destroy) cmd_destroy ;;
    shell)   cmd_shell ;;
    help|--help|-h) cmd_help ;;
    *)
        error "Unknown command: $1"
        cmd_help
        exit 1
        ;;
esac
