#!/usr/bin/env bash
# ============================================================================
# RevSocks v4 — Proxy Helper Utility
#
# Manages proxychains configuration, tests SOCKS5 connectivity,
# and lists currently connected agent targets.
#
# Usage:
#   ./proxy_helper.sh list                  — Show connected targets
#   ./proxy_helper.sh test <target_id>      — Test SOCKS5 connectivity
#   ./proxy_helper.sh config <target_id>    — Generate proxychains4.conf
#   ./proxy_helper.sh status                — Show server status
#   ./proxy_helper.sh clean                 — Remove generated configs
#
# ============================================================================
set -euo pipefail

# --- Configuration ---
REVSOCKS_DIR="$(cd "$(dirname "$0")/.." && pwd)"
SERVER_HOST="${REVSOCKS_HOST:-127.0.0.1}"
SERVER_PORT="${REVSOCKS_PORT:-443}"
SOCKS_BASE_PORT="${REVSOCKS_SOCKS_BASE:-1080}"
CONFIG_DIR="${REVSOCKS_DIR}/output/proxychains"
CONTROL_SOCKET="${REVSOCKS_DIR}/output/server_control.sock"

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

# --- Helpers ---
die() { echo -e "${RED}[!] $*${NC}" >&2; exit 1; }
info() { echo -e "${CYAN}[*]${NC} $*"; }
ok() { echo -e "${GREEN}[+]${NC} $*"; }
warn() { echo -e "${YELLOW}[!]${NC} $*"; }

usage() {
    cat <<EOF
RevSocks v4 — Proxy Helper

Usage: $(basename "$0") <command> [args]

Commands:
  list                   List all connected agent targets
  test   <target_id>     Test SOCKS5 connectivity for a target
  config <target_id>     Generate proxychains4.conf for a target
  status                 Show server and listener status
  clean                  Remove all generated proxychains configs

Environment:
  REVSOCKS_HOST          Server host (default: 127.0.0.1)
  REVSOCKS_PORT          Server control port (default: 443)
  REVSOCKS_SOCKS_BASE    Base SOCKS5 port (default: 1080)

EOF
    exit 0
}

# --- Check dependencies ---
check_deps() {
    for cmd in curl nc ss; do
        if ! command -v "$cmd" &>/dev/null; then
            warn "Optional dependency '$cmd' not found"
        fi
    done
}

# --- List connected targets ---
cmd_list() {
    info "Connected targets:"
    echo ""

    # Try to read from server's status endpoint or state file
    local state_file="${REVSOCKS_DIR}/output/state.json"
    local status_url="http://${SERVER_HOST}:${SERVER_PORT}/api/targets"

    # Method 1: Read state file directly
    if [ -f "$state_file" ]; then
        if command -v python3 &>/dev/null; then
            python3 -c "
import json, sys
try:
    with open('$state_file') as f:
        state = json.load(f)
    targets = state.get('targets', state.get('agents', []))
    if not targets:
        print('  (no active targets)')
        sys.exit(0)
    print(f'  {\"ID\":<6} {\"Address\":<22} {\"SOCKS Port\":<12} {\"Status\":<10} {\"Last Seen\":<20}')
    print(f'  {\"-\"*6} {\"-\"*22} {\"-\"*12} {\"-\"*10} {\"-\"*20}')
    for i, t in enumerate(targets):
        tid = t.get('id', i)
        addr = t.get('address', t.get('ip', 'unknown'))
        sport = t.get('socks_port', $SOCKS_BASE_PORT + i)
        status = t.get('status', 'active')
        last = t.get('last_seen', 'now')
        print(f'  {tid:<6} {addr:<22} {sport:<12} {status:<10} {last:<20}')
except Exception as e:
    print(f'  Error reading state: {e}')
"
        else
            echo "  (python3 required for JSON parsing)"
            echo "  Raw state: $(head -c 500 "$state_file")"
        fi
        return
    fi

    # Method 2: Check listening SOCKS ports
    if command -v ss &>/dev/null; then
        echo "  Active SOCKS listeners:"
        ss -tlnp 2>/dev/null | grep -E ":(108[0-9]|109[0-9]|10[0-9]{2})" | while read -r line; do
            local port
            port=$(echo "$line" | grep -oP ':\K[0-9]+' | head -1)
            local pid_info
            pid_info=$(echo "$line" | grep -oP 'pid=\K[0-9]+' || echo "unknown")
            echo "    Port $port (PID: $pid_info)"
        done
    fi

    # Method 3: Scan common SOCKS ports
    info "Scanning SOCKS ports ${SOCKS_BASE_PORT}-$((SOCKS_BASE_PORT+19))..."
    local found=0
    for port in $(seq "$SOCKS_BASE_PORT" $((SOCKS_BASE_PORT + 19))); do
        if (echo >/dev/tcp/"$SERVER_HOST"/"$port") 2>/dev/null; then
            ok "  Target $((port - SOCKS_BASE_PORT)): SOCKS5 on ${SERVER_HOST}:${port}"
            found=$((found + 1))
        fi
    done

    if [ "$found" -eq 0 ]; then
        warn "  No active SOCKS listeners found"
    else
        echo ""
        info "$found active target(s)"
    fi
}

# --- Test SOCKS5 connectivity ---
cmd_test() {
    local target_id="${1:-}"
    [ -z "$target_id" ] && die "Usage: $(basename "$0") test <target_id>"

    local socks_port=$((SOCKS_BASE_PORT + target_id))
    info "Testing SOCKS5 proxy on ${SERVER_HOST}:${socks_port}..."

    # Test 1: Port open?
    if ! (echo >/dev/tcp/"$SERVER_HOST"/"$socks_port") 2>/dev/null; then
        die "Port $socks_port is not listening"
    fi
    ok "Port $socks_port is open"

    # Test 2: SOCKS5 handshake
    if command -v python3 &>/dev/null; then
        python3 -c "
import socket, sys
try:
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(5)
    s.connect(('$SERVER_HOST', $socks_port))
    # SOCKS5 greeting: version 5, 1 auth method (no auth)
    s.send(b'\\x05\\x01\\x00')
    resp = s.recv(2)
    if resp == b'\\x05\\x00':
        print('[+] SOCKS5 handshake OK (no auth)')
    elif resp == b'\\x05\\x02':
        print('[+] SOCKS5 handshake OK (user/pass auth required)')
    elif len(resp) >= 2 and resp[0] == 5:
        print(f'[+] SOCKS5 responded (method: {resp[1]})')
    else:
        print(f'[-] Unexpected response: {resp.hex()}')
        sys.exit(1)
    s.close()
except Exception as e:
    print(f'[-] SOCKS5 test failed: {e}')
    sys.exit(1)
" || die "SOCKS5 handshake failed"
    else
        # Fallback: just verify port is open with nc
        if command -v nc &>/dev/null; then
            echo -ne '\x05\x01\x00' | nc -w3 "$SERVER_HOST" "$socks_port" | head -c2 | xxd -p
        fi
    fi

    # Test 3: Try a connection through the proxy
    if command -v curl &>/dev/null; then
        info "Testing HTTP through SOCKS5..."
        if curl -s --max-time 10 --socks5 "${SERVER_HOST}:${socks_port}" \
                "http://ifconfig.me" -o /dev/null -w "%{http_code}" 2>/dev/null | grep -q "200"; then
            ok "HTTP through SOCKS5 works"
            local ext_ip
            ext_ip=$(curl -s --max-time 10 --socks5 "${SERVER_HOST}:${socks_port}" "http://ifconfig.me" 2>/dev/null)
            ok "Exit IP: $ext_ip"
        else
            warn "HTTP test failed (target may have restricted internet)"
        fi
    fi

    echo ""
    ok "Target $target_id (port $socks_port) is functional"
}

# --- Generate proxychains config ---
cmd_config() {
    local target_id="${1:-}"
    [ -z "$target_id" ] && die "Usage: $(basename "$0") config <target_id>"

    local socks_port=$((SOCKS_BASE_PORT + target_id))

    # Verify port is listening
    if ! (echo >/dev/tcp/"$SERVER_HOST"/"$socks_port") 2>/dev/null; then
        warn "Port $socks_port not listening — generating config anyway"
    fi

    mkdir -p "$CONFIG_DIR"
    local conf_file="${CONFIG_DIR}/proxychains_target${target_id}.conf"

    cat > "$conf_file" <<PCEOF
# proxychains4 configuration — RevSocks v4
# Target: $target_id
# SOCKS5: ${SERVER_HOST}:${socks_port}
# Generated: $(date -u '+%Y-%m-%d %H:%M:%S UTC')
#
# Usage:
#   proxychains4 -f $conf_file <command>
#   proxychains4 -f $conf_file nmap -sT -Pn <target>
#   proxychains4 -f $conf_file curl http://internal-host/

# Quiet mode — less output noise
quiet_mode

# Proxy DNS through the chain
proxy_dns

# Timeouts
tcp_read_time_out 15000
tcp_connect_time_out 10000

# Chain type
# strict_chain = all proxies must succeed in order
# dynamic_chain = skip dead proxies
dynamic_chain

[ProxyList]
socks5 ${SERVER_HOST} ${socks_port}
PCEOF

    ok "Config written: $conf_file"
    echo ""
    info "Usage:"
    echo "  proxychains4 -f $conf_file nmap -sT -Pn <target>"
    echo "  proxychains4 -f $conf_file ssh user@internal-host"
    echo "  proxychains4 -f $conf_file curl http://internal:8080/"
    echo ""

    # Also set up env shortcut
    echo ""
    info "Quick alias (paste in terminal):"
    echo "  alias pc${target_id}='proxychains4 -f $conf_file'"
}

# --- Server status ---
cmd_status() {
    info "RevSocks v4 Server Status"
    echo ""

    # Check server process
    if pgrep -f "server\.py" &>/dev/null; then
        local pid
        pid=$(pgrep -f "server\.py" | head -1)
        ok "Server process running (PID: $pid)"
    else
        warn "Server process not found"
    fi

    # Check CF bridge
    if pgrep -f "cf_bridge\.py" &>/dev/null; then
        local pid
        pid=$(pgrep -f "cf_bridge\.py" | head -1)
        ok "CF Bridge running (PID: $pid)"
    else
        warn "CF Bridge not running"
    fi

    # Check port 443
    if (echo >/dev/tcp/"$SERVER_HOST"/443) 2>/dev/null; then
        ok "Port 443 listening"
    else
        warn "Port 443 not listening"
    fi

    # Check port 4443 (CF bridge)
    if (echo >/dev/tcp/"$SERVER_HOST"/4443) 2>/dev/null; then
        ok "Port 4443 (CF bridge) listening"
    else
        warn "Port 4443 (CF bridge) not listening"
    fi

    # Memory usage
    if pgrep -f "server\.py" &>/dev/null; then
        local pid
        pid=$(pgrep -f "server\.py" | head -1)
        local mem
        mem=$(ps -o rss= -p "$pid" 2>/dev/null | tr -d ' ')
        if [ -n "$mem" ]; then
            local mem_mb=$((mem / 1024))
            info "Server memory: ${mem_mb} MB"
        fi
    fi

    echo ""
    cmd_list
}

# --- Clean generated configs ---
cmd_clean() {
    if [ -d "$CONFIG_DIR" ]; then
        local count
        count=$(find "$CONFIG_DIR" -name "proxychains_target*.conf" 2>/dev/null | wc -l)
        rm -f "$CONFIG_DIR"/proxychains_target*.conf
        ok "Removed $count config file(s)"
    else
        info "No configs to clean"
    fi
}

# --- Main ---
check_deps

case "${1:-}" in
    list)       cmd_list ;;
    test)       cmd_test "${2:-}" ;;
    config)     cmd_config "${2:-}" ;;
    status)     cmd_status ;;
    clean)      cmd_clean ;;
    -h|--help|help|"")  usage ;;
    *)          die "Unknown command: $1 (use --help)" ;;
esac