#!/bin/bash

# OpenConnect-SSO VPN Connection Script
# Usage: ./openconnect-vpn [-c|--connect] [-d|--disconnect] [-s|--status] [--setup-keyring] [--help]

# Credentials from environment variables (set by runtipi)
VPN_EMAIL="${VPN_EMAIL:-}"
VPN_PASSWORD="${VPN_PASSWORD:-}"
VPN_TOTP_SECRET="${VPN_TOTP_SECRET:-}"
VPN_HOST="${VPN_HOST:-}"
TARGET_IP="${TARGET_IP:-}"
VPN_INTERFACE="${VPN_INTERFACE:-tun0}"


# Log directory
LOG_DIR="/var/log/openconnect-vpn"
LOG_RETENTION_DAYS=7
mkdir -p "$LOG_DIR" 2>/dev/null

get_log_file() {
    echo "$LOG_DIR/$(date '+%Y-%m-%d').log"
}

cleanup_old_logs() {
    find "$LOG_DIR" -name "*.log" -type f -mtime +$LOG_RETENTION_DAYS -delete 2>/dev/null
}

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

print_banner() {
    echo -e "${CYAN}========================================${NC}"
    echo -e "${CYAN}    OpenConnect-SSO VPN Connection      ${NC}"
    echo -e "${CYAN}========================================${NC}"
    echo ""
}

# Flags
DO_CONNECT=false
DO_DISCONNECT=false
DO_SETUP_KEYRING=false

# Logging
log() {
    local level="$1"
    local msg="$2"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    local timestamp_short=$(date '+%H:%M:%S')
    local log_file=$(get_log_file)

    local msg_plain=$(echo -e "$msg" | sed 's/\x1b\[[0-9;]*m//g')
    echo "[$timestamp] [$level] $msg_plain" >> "$log_file"

    case $level in
        INFO)  echo -e "${GRAY}[$timestamp_short]${NC} ${GREEN}[INFO]${NC}  $msg" ;;
        WARN)  echo -e "${GRAY}[$timestamp_short]${NC} ${YELLOW}[WARN]${NC}  $msg" ;;
        ERROR) echo -e "${GRAY}[$timestamp_short]${NC} ${RED}[ERROR]${NC} $msg" ;;
        DEBUG) echo -e "${GRAY}[$timestamp_short]${NC} ${CYAN}[DEBUG]${NC} $msg" ;;
        *)     echo -e "${GRAY}[$timestamp_short]${NC} $msg" ;;
    esac
}

run_cmd() {
    local desc="$1"
    shift
    log DEBUG "$desc: $*"
    output=$("$@" 2>&1)
    local rc=$?
    if [ -n "$output" ]; then
        echo "$output" | while IFS= read -r line; do
            echo -e "           ${GRAY}│${NC} $line"
        done
    fi
    return $rc
}

# Fetch server certificate fingerprint from VPN host
get_server_cert() {
    local host="$1"
    # Strip protocol and path to get just hostname:port
    local server=$(echo "$host" | sed -E 's|^https?://||; s|/.*$||')
    # Add default port if not specified
    [[ "$server" != *:* ]] && server="${server}:443"

    log DEBUG "Fetching server certificate from $server..."

    local cert_pin=$(echo | openssl s_client -connect "$server" 2>/dev/null \
        | openssl x509 -pubkey -noout 2>/dev/null \
        | openssl pkey -pubin -outform der 2>/dev/null \
        | openssl dgst -sha256 -binary \
        | openssl enc -base64)

    if [[ -n "$cert_pin" ]]; then
        echo "pin-sha256:${cert_pin}"
    else
        log WARN "Could not fetch server certificate"
        echo ""
    fi
}

# Setup keyring with credentials for openconnect-sso
setup_keyring() {
    log INFO "Setting up keyring credentials..."

    if [[ -z "$VPN_EMAIL" ]]; then
        log ERROR "VPN_EMAIL is not set"
        return 1
    fi

    python3 << PYTHON
import keyring
from keyrings.alt.file import PlaintextKeyring

# Use plaintext keyring (works without desktop environment)
keyring.set_keyring(PlaintextKeyring())

email = "$VPN_EMAIL"
password = "$VPN_PASSWORD"
totp_secret = "$VPN_TOTP_SECRET"

# Store password
if password:
    keyring.set_password('openconnect-sso', email, password)
    print(f"Password stored in keyring for {email}")

# Store TOTP secret
if totp_secret:
    keyring.set_password('openconnect-sso', f'totp/{email}', totp_secret.upper())
    print(f"TOTP secret stored in keyring for {email}")

print("Keyring setup complete")
PYTHON

    if [ $? -eq 0 ]; then
        log INFO "Keyring credentials configured"
    else
        log ERROR "Failed to setup keyring"
        return 1
    fi
}

show_help() {
    echo -e "${CYAN}OpenConnect-SSO VPN Connection Script${NC}"
    echo ""
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  -c, --connect       Connect to VPN"
    echo "  -d, --disconnect    Disconnect from VPN"
    echo "  -s, --status        Show VPN status"
    echo "  -r, --routes        Show routing table"
    echo "  --setup-keyring     Setup keyring credentials only"
    echo "  --help              Show this help message"
    echo ""
    echo "Environment Variables:"
    echo "  VPN_EMAIL           Email/username for SSO"
    echo "  VPN_PASSWORD        Password for SSO"
    echo "  VPN_TOTP_SECRET     TOTP secret for 2FA"
    echo "  VPN_HOST            VPN server hostname"
    echo "  TARGET_IP           Target IP for connectivity test"
    echo "  VPN_INTERFACE       TUN interface name (default: tun0)"
}

get_vpn_interface() {
    local iface=$(ip link show | grep -oP 'tun\d+(?=:.*UP)' | head -1)
    if [ -z "$iface" ]; then
        iface=$(ip link show | grep -oP 'tun\d+' | head -1)
    fi
    echo "$iface"
}

get_container_ip() {
    ip addr show eth0 2>/dev/null | grep -oP 'inet \K[\d.]+' | head -1
}

get_vpn_ip() {
    local iface=$(get_vpn_interface)
    if [ -n "$iface" ]; then
        ip addr show "$iface" 2>/dev/null | grep -oP 'inet \K[\d.]+' | head -1
    fi
}

check_vpn_status() {
    local vpn_iface=$(get_vpn_interface)
    if [ -n "$vpn_iface" ]; then
        local vpn_ip=$(get_vpn_ip)
        log INFO "VPN is ${GREEN}CONNECTED${NC}"
        log DEBUG "  Interface: $vpn_iface"
        log DEBUG "  VPN IP: $vpn_ip"
        return 0
    else
        log WARN "VPN is ${RED}NOT CONNECTED${NC}"
        return 1
    fi
}

show_routes() {
    echo -e "${CYAN}========================================${NC}"
    echo -e "${CYAN}         Current Routing Table          ${NC}"
    echo -e "${CYAN}========================================${NC}"
    echo ""
    echo -e "${GREEN}IPv4 Routes:${NC}"
    ip -4 route show | while IFS= read -r line; do
        if echo "$line" | grep -qE "tun[0-9]"; then
            echo -e "  ${YELLOW}$line${NC}"
        elif echo "$line" | grep -qE "10\.|172\.(1[6-9]|2[0-9]|3[01])\.|192\.168\."; then
            echo -e "  ${CYAN}$line${NC}"
        else
            echo "  $line"
        fi
    done
    echo ""
    echo -e "${GREEN}VPN Interface:${NC}"
    local vpn_iface=$(get_vpn_interface)
    if [ -n "$vpn_iface" ]; then
        ip addr show "$vpn_iface" 2>/dev/null | grep -E "inet|link" | while IFS= read -r line; do
            echo "  $line"
        done
    else
        echo -e "  ${RED}No VPN interface found${NC}"
    fi
}

show_network_status() {
    log INFO "Current network status:"
    echo ""
    log DEBUG "Container Network Interfaces:"
    ip -4 addr show | grep -E "inet |^[0-9]+:" | while IFS= read -r line; do
        echo -e "           ${GRAY}│${NC} $line"
    done

    echo ""
    local vpn_iface=$(get_vpn_interface)
    if [ -n "$vpn_iface" ]; then
        local vpn_ip=$(get_vpn_ip)
        log INFO "VPN Status: ${GREEN}CONNECTED${NC}"
        log DEBUG "  Interface: $vpn_iface"
        log DEBUG "  VPN IP: $vpn_ip"
    else
        log WARN "VPN Status: ${RED}NOT CONNECTED${NC}"
    fi

    local container_ip=$(get_container_ip)
    if [ -n "$container_ip" ]; then
        log DEBUG "Container IP: $container_ip"
    fi

    echo ""
    log DEBUG "Default gateway:"
    ip route show default | while IFS= read -r line; do
        echo -e "           ${GRAY}│${NC} $line"
    done
    echo ""
}

kill_vpn_processes() {
    log INFO "Killing VPN processes..."
    local killed=0

    for pid in $(pgrep -x "openconnect" 2>/dev/null); do
        log DEBUG "Killing openconnect (PID $pid)"
        kill -9 "$pid" 2>/dev/null && ((killed++))
    done

    for pid in $(pgrep -f "openconnect-sso" 2>/dev/null); do
        log DEBUG "Killing openconnect-sso (PID $pid)"
        kill -9 "$pid" 2>/dev/null && ((killed++))
    done

    if [ $killed -eq 0 ]; then
        log INFO "No VPN processes were running"
    else
        log INFO "Killed $killed process(es)"
        sleep 2
    fi
}

disconnect_vpn() {
    log INFO "Disconnecting VPN..."
    kill_vpn_processes

    if check_vpn_status; then
        log WARN "VPN still appears connected"
        return 1
    fi

    log INFO "VPN disconnected"
    return 0
}

setup_forwarding() {
    log INFO "Setting up IP forwarding rules..."

    local vpn_iface=$(get_vpn_interface)
    if [ -z "$vpn_iface" ]; then
        log ERROR "No VPN interface found! Is VPN connected?"
        return 1
    fi

    local vpn_ip=$(get_vpn_ip)
    local container_ip=$(get_container_ip)

    log DEBUG "VPN interface: $vpn_iface"
    log DEBUG "VPN IP: $vpn_ip"
    log DEBUG "Container IP: $container_ip"

    run_cmd "Enabling IP forwarding" sysctl -w net.ipv4.ip_forward=1

    # NAT masquerade for container network
    if ! iptables -t nat -C POSTROUTING -o "$vpn_iface" -j MASQUERADE 2>/dev/null; then
        run_cmd "Adding NAT masquerade" iptables -t nat -A POSTROUTING -o "$vpn_iface" -j MASQUERADE
    fi

    log INFO "Forwarding rules configured"

    # Trigger host routing service restart if available
    if [ -d /runtime ]; then
        touch /runtime/restart-routing 2>/dev/null || true
    fi
}

connect_vpn() {
    log INFO "=== Starting OpenConnect-SSO VPN ==="
    echo ""

    # Kill any existing VPN processes
    kill_vpn_processes

    # Validate required variables
    if [[ -z "$VPN_HOST" ]]; then
        log ERROR "VPN_HOST is not set"
        return 1
    fi

    # Setup keyring credentials
    setup_keyring

    # Build openconnect-sso command
    local sso_args=()
    sso_args+=("-s" "$VPN_HOST")

    if [[ -n "$VPN_EMAIL" ]]; then
        sso_args+=("-u" "$VPN_EMAIL")
    fi

    sso_args+=("--browser-display-mode" "hidden")

    # Fetch server certificate
    log INFO "Fetching server certificate..."
    local servercert=$(get_server_cert "$VPN_HOST")

    # Build openconnect args
    local oc_args=()
    oc_args+=("--protocol=anyconnect")
    oc_args+=("--interface" "$VPN_INTERFACE")
    oc_args+=("--script" "/usr/share/vpnc-scripts/vpnc-script")

    if [[ -n "$servercert" ]]; then
        oc_args+=("--servercert" "$servercert")
        log DEBUG "Server cert: $servercert"
    fi

    log INFO "Connecting to: $VPN_HOST"
    log DEBUG "Interface: $VPN_INTERFACE"

    # Launch openconnect-sso
    log INFO "Launching openconnect-sso..."
    openconnect-sso "${sso_args[@]}" -- /usr/sbin/openconnect "${oc_args[@]}" &
    OC_PID=$!
    disown $OC_PID
    log DEBUG "openconnect-sso started with PID $OC_PID"

    # Wait for VPN to connect
    log INFO "Waiting for VPN connection..."
    local wait_count=0
    local max_wait=300

    while [ -z "$(get_vpn_interface)" ]; do
        sleep 2
        ((wait_count+=2))
        if [ $((wait_count % 10)) -eq 0 ]; then
            log DEBUG "Still waiting for VPN... (${wait_count}s)"
        fi
        if [ $wait_count -ge $max_wait ]; then
            log ERROR "Timeout waiting for VPN connection after ${max_wait}s"
            return 1
        fi
    done

    log INFO "VPN connected!"
    local vpn_iface=$(get_vpn_interface)
    local vpn_ip=$(get_vpn_ip)
    log DEBUG "  Interface: $vpn_iface"
    log DEBUG "  VPN IP: $vpn_ip"

    sleep 3
    setup_forwarding

    # Test connection if TARGET_IP is set
    if [[ -n "$TARGET_IP" ]]; then
        log INFO "Testing connection to $TARGET_IP..."
        if ping -c 2 -W 3 "$TARGET_IP" &>/dev/null; then
            log INFO "Connection test: ${GREEN}SUCCESS${NC}"
        else
            log WARN "Connection test: ${RED}FAILED${NC}"
        fi
    fi

    log INFO "VPN setup complete"
    return 0
}

# Main menu
main_menu() {
    echo -e "${GREEN}Options:${NC}"
    echo -e "  ${CYAN}1${NC} - Connect VPN"
    echo -e "  ${CYAN}2${NC} - Disconnect VPN"
    echo -e "  ${CYAN}3${NC} - Show VPN status"
    echo -e "  ${CYAN}4${NC} - Setup IP forwarding"
    echo -e "  ${CYAN}5${NC} - Test connection"
    echo -e "  ${CYAN}6${NC} - Show network status"
    echo -e "  ${CYAN}7${NC} - Show routing table"
    echo -e "  ${CYAN}8${NC} - Setup keyring"
    echo -e "  ${CYAN}q${NC} - Quit"
    echo ""
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -c|--connect)
                DO_CONNECT=true
                shift
                ;;
            -d|--disconnect)
                DO_DISCONNECT=true
                shift
                ;;
            -s|--status)
                print_banner
                check_vpn_status
                echo ""
                show_network_status
                exit 0
                ;;
            -r|--routes)
                show_routes
                exit 0
                ;;
            --setup-keyring)
                DO_SETUP_KEYRING=true
                shift
                ;;
            --help)
                show_help
                exit 0
                ;;
            *)
                echo "Unknown option: $1"
                echo "Use --help for usage information"
                exit 1
                ;;
        esac
    done
}

# Main
parse_args "$@"
cleanup_old_logs

echo "" >> "$(get_log_file)"
echo "========================================" >> "$(get_log_file)"
log INFO "openconnect-vpn script started"
log DEBUG "VPN_EMAIL=$VPN_EMAIL"
log DEBUG "VPN_HOST=$VPN_HOST"
log DEBUG "TARGET_IP=$TARGET_IP"
log DEBUG "VPN_TOTP_SECRET is $([ -n "$VPN_TOTP_SECRET" ] && echo 'set' || echo 'NOT SET')"

print_banner

if [ "$DO_SETUP_KEYRING" = "true" ]; then
    setup_keyring
    exit $?
fi

if [ "$DO_DISCONNECT" = "true" ]; then
    disconnect_vpn
    exit $?
fi

if [ "$DO_CONNECT" = "true" ]; then
    connect_vpn
    exit $?
fi

# Interactive menu mode
while true; do
    echo ""
    main_menu
    echo -ne "${CYAN}Choice: ${NC}"
    read -r choice
    echo ""

    [[ -z "${choice// }" ]] && continue

    case $choice in
        1) connect_vpn ;;
        2) disconnect_vpn ;;
        3) check_vpn_status ;;
        4) setup_forwarding ;;
        5) if [[ -n "$TARGET_IP" ]]; then
               log INFO "Testing connection to $TARGET_IP..."
               ping -c 3 "$TARGET_IP" && log INFO "Connection test: ${GREEN}SUCCESS${NC}" || log ERROR "Connection test: ${RED}FAILED${NC}"
           else
               log WARN "TARGET_IP not set"
           fi ;;
        6) show_network_status ;;
        7) show_routes ;;
        8) setup_keyring ;;
        q|Q) log INFO "Goodbye!"; exit 0 ;;
        *) ;;
    esac
done