#!/bin/bash
#
# check_connection.sh - Verify server is running and reachable
#
# Usage: ./check_connection.sh <SERVER_IP> [TUNNEL_PORT] [PAYLOAD_PORT] [SOCKS_PORT]
#

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

SERVER_IP="${1}"
TUNNEL_PORT="${2:-443}"
PAYLOAD_PORT="${3:-8443}"
SOCKS_PORT="${4:-1080}"

if [ -z "$SERVER_IP" ]; then
    echo -e "${RED}[!] Usage: $0 <SERVER_IP> [TUNNEL_PORT] [PAYLOAD_PORT] [SOCKS_PORT]${NC}"
    exit 1
fi

echo -e "${CYAN}"
echo "╔══════════════════════════════════════╗"
echo "║     CONNECTION CHECKER               ║"
echo "╠══════════════════════════════════════╣"
echo "║  Server : ${SERVER_IP}"
echo "║  Tunnel : ${TUNNEL_PORT}"
echo "║  Payload: ${PAYLOAD_PORT}"
echo "║  SOCKS5 : ${SOCKS_PORT}"
echo "╚══════════════════════════════════════╝"
echo -e "${NC}"

PASS=0
FAIL=0

check() {
    local desc="$1"
    local result="$2"
    if [ "$result" -eq 0 ]; then
        echo -e "  ${GREEN}[PASS]${NC} $desc"
        ((PASS++))
    else
        echo -e "  ${RED}[FAIL]${NC} $desc"
        ((FAIL++))
    fi
}

# ============ BASIC CONNECTIVITY ============
echo -e "${YELLOW}[1/5] Basic connectivity${NC}"

ping -c 1 -W 3 "$SERVER_IP" &>/dev/null
check "ICMP ping to ${SERVER_IP}" $?

# ============ TUNNEL PORT ============
echo -e "${YELLOW}[2/5] Tunnel port (${TUNNEL_PORT}/tcp)${NC}"

timeout 5 bash -c "echo >/dev/tcp/${SERVER_IP}/${TUNNEL_PORT}" 2>/dev/null
check "TCP connect to ${SERVER_IP}:${TUNNEL_PORT}" $?

# Also check with nc if available
if command -v nc &>/dev/null; then
    nc -z -w3 "$SERVER_IP" "$TUNNEL_PORT" 2>/dev/null
    check "NC probe ${SERVER_IP}:${TUNNEL_PORT}" $?
fi

# ============ PAYLOAD PORT ============
echo -e "${YELLOW}[3/5] Payload server (${PAYLOAD_PORT}/tcp)${NC}"

timeout 5 bash -c "echo >/dev/tcp/${SERVER_IP}/${PAYLOAD_PORT}" 2>/dev/null
check "TCP connect to ${SERVER_IP}:${PAYLOAD_PORT}" $?

# Try HTTP GET on payload
if command -v curl &>/dev/null; then
    HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" --max-time 5 "http://${SERVER_IP}:${PAYLOAD_PORT}/payload" 2>/dev/null)
    if [ "$HTTP_CODE" = "200" ]; then
        check "HTTP GET payload (status: ${HTTP_CODE})" 0
        
        # Check payload size
        PAYLOAD_SIZE=$(curl -s --max-time 10 "http://${SERVER_IP}:${PAYLOAD_PORT}/payload" 2>/dev/null | wc -c)
        if [ "$PAYLOAD_SIZE" -gt 0 ]; then
            echo -e "  ${GREEN}[INFO]${NC} Payload size: ${PAYLOAD_SIZE} bytes"
        fi
    else
        check "HTTP GET payload (status: ${HTTP_CODE})" 1
    fi
elif command -v wget &>/dev/null; then
    wget -q --timeout=5 -O /dev/null "http://${SERVER_IP}:${PAYLOAD_PORT}/payload" 2>/dev/null
    check "HTTP GET payload (wget)" $?
fi

# ============ SOCKS5 PORT (local only) ============
echo -e "${YELLOW}[4/5] SOCKS5 proxy (127.0.0.1:${SOCKS_PORT})${NC}"

if [ "$SERVER_IP" = "127.0.0.1" ] || [ "$SERVER_IP" = "localhost" ]; then
    timeout 3 bash -c "echo >/dev/tcp/127.0.0.1/${SOCKS_PORT}" 2>/dev/null
    check "SOCKS5 listening on 127.0.0.1:${SOCKS_PORT}" $?
    
    # Test SOCKS5 handshake
    if command -v curl &>/dev/null; then
        curl -s --max-time 5 --socks5 "127.0.0.1:${SOCKS_PORT}" "http://ifconfig.me" &>/dev/null
        if [ $? -eq 0 ]; then
            check "SOCKS5 proxy functional" 0
        else
            echo -e "  ${YELLOW}[WARN]${NC} SOCKS5 port open but no agent connected (expected if no client yet)"
        fi
    fi
else
    echo -e "  ${YELLOW}[SKIP]${NC} SOCKS5 is local-only (127.0.0.1), run this check on the server"
fi

# ============ SERVER PROCESS ============
echo -e "${YELLOW}[5/5] Server process${NC}"

if [ "$SERVER_IP" = "127.0.0.1" ] || [ "$SERVER_IP" = "localhost" ]; then
    if pgrep -f "server.py" &>/dev/null; then
        check "server.py process running" 0
        PID=$(pgrep -f "server.py" | head -1)
        echo -e "  ${GREEN}[INFO]${NC} PID: ${PID}"
    else
        check "server.py process running" 1
        echo -e "  ${YELLOW}[TIP]${NC}  Start with: python3 server.py &"
    fi
    
    # Check listening ports
    echo -e "\n  ${CYAN}Listening ports:${NC}"
    ss -tlnp 2>/dev/null | grep -E "(${TUNNEL_PORT}|${PAYLOAD_PORT}|${SOCKS_PORT})" | while read line; do
        echo -e "  ${GREEN}  $line${NC}"
    done
else
    # Remote check via SSH if possible
    echo -e "  ${YELLOW}[SKIP]${NC} Process check requires local access or SSH"
fi

# ============ FIREWALL CHECK ============
echo ""
echo -e "${YELLOW}[*] Firewall hints:${NC}"
echo -e "  ${CYAN}Ensure these ports are open on server:${NC}"
echo -e "    - ${TUNNEL_PORT}/tcp  (agent reverse connection)"
echo -e "    - ${PAYLOAD_PORT}/tcp (payload download)"
echo -e "  ${CYAN}Commands:${NC}"
echo -e "    ufw allow ${TUNNEL_PORT}/tcp"
echo -e "    ufw allow ${PAYLOAD_PORT}/tcp"
echo -e "    # OR"
echo -e "    iptables -A INPUT -p tcp --dport ${TUNNEL_PORT} -j ACCEPT"
echo -e "    iptables -A INPUT -p tcp --dport ${PAYLOAD_PORT} -j ACCEPT"

# ============ AGENT STATUS ============
echo ""
echo -e "${YELLOW}[*] Agent connection test:${NC}"
echo -e "  ${CYAN}To verify agent callback, watch server logs for:${NC}"
echo -e "    [+] Agent connected: <IP>:<PORT>"
echo ""

# ============ SUMMARY ============
echo ""
TOTAL=$((PASS + FAIL))
echo -e "${CYAN}═══════════════════════════════════════${NC}"
echo -e "  Results: ${GREEN}${PASS} passed${NC} / ${RED}${FAIL} failed${NC} / ${TOTAL} total"
echo -e "${CYAN}═══════════════════════════════════════${NC}"

if [ $FAIL -eq 0 ]; then
    echo -e "\n  ${GREEN}All checks passed. Server is operational.${NC}"
    echo -e "  ${GREEN}Waiting for agent connection on :${TUNNEL_PORT}${NC}"
else
    echo -e "\n  ${RED}Some checks failed. Troubleshoot:${NC}"
    echo -e "  ${YELLOW}1. Is server.py running?     python3 server.py &${NC}"
    echo -e "  ${YELLOW}2. Firewall blocking ports?   ufw status${NC}"
    echo -e "  ${YELLOW}3. Correct IP?                ip addr show${NC}"
    echo -e "  ${YELLOW}4. payload.enc exists?        ls -la client/payload.enc${NC}"
fi

exit $FAIL