#!/bin/bash

# WebSocket SSH Client Watchdog Script
# Copyright (C) 2024 Stefy Lanza <stefy@nexlab.net> and SexHack.me
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# Configuration
DAEMON_NAME="wssshc"
DAEMON_PATH="/usr/bin/wssshc"
PID_FILE="/var/run/wssshc.pid"
WATCHDOG_PID_FILE="/var/run/wssshc-watchdog.pid"
LOG_FILE="/var/log/wssshc/watchdog.log"
CHECK_INTERVAL=30
MAX_RESTARTS=20
RESTART_WINDOW=60  # 1 minute

# Default configuration values (can be overridden by /etc/default/wssshc)
START=yes
DAEMON_ARGS=""

# Load configuration if available
if [ -f /etc/default/wssshc ]; then
    . /etc/default/wssshc
fi

# Function to log messages
log_message() {
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[$timestamp] $*" >> "$LOG_FILE"
    logger -t "$DAEMON_NAME-watchdog" "$*"
}

# Function to check if daemon is running (double-check process)
is_daemon_running() {
    if [ -f "$PID_FILE" ]; then
        local pid=$(cat "$PID_FILE")
        # First check if process exists
        if kill -0 "$pid" 2>/dev/null; then
            # Double-check: verify the process is actually our daemon
            if ps -p "$pid" -o comm= 2>/dev/null | grep -q "^wssshc$"; then
                return 0  # Running and correct process
            else
                log_message "PID file exists but process $pid is not wssshc"
                rm -f "$PID_FILE"
            fi
        else
            log_message "PID file exists but process $pid is not running"
            rm -f "$PID_FILE"
        fi
    fi
    return 1  # Not running
}

# Function to start daemon
start_daemon() {
    log_message "Starting $DAEMON_NAME daemon..."

    # Create necessary directories
    mkdir -p /var/log/wssshc 2>/dev/null || log_message "Warning: Could not create /var/log/wssshc"
    chown wssshc:wssshc /var/log/wssshc 2>/dev/null || log_message "Warning: Could not chown /var/log/wssshc"

    # Check if daemon binary exists
    if [ ! -x "$DAEMON_PATH" ]; then
        log_message "Error: Daemon binary $DAEMON_PATH not found or not executable"
        return 1
    fi

    # Try to start daemon as wssshc user, fallback to current user if that fails
    log_message "Attempting to start daemon with start-stop-daemon..."
    if [ -n "$DAEMON_ARGS" ]; then
        start-stop-daemon --start --quiet --pidfile "$PID_FILE" \
            --chuid wssshc:wssshc --background --make-pidfile \
            --exec "$DAEMON_PATH" -- $DAEMON_ARGS 2>/dev/null
        local result=$?
    else
        start-stop-daemon --start --quiet --pidfile "$PID_FILE" \
            --chuid wssshc:wssshc --background --make-pidfile \
            --exec "$DAEMON_PATH" 2>/dev/null
        local result=$?
    fi

    # If start-stop-daemon failed, try running directly
    if [ $result -ne 0 ]; then
        log_message "start-stop-daemon failed (exit code: $result), trying direct execution..."
        if [ -n "$DAEMON_ARGS" ]; then
            "$DAEMON_PATH" $DAEMON_ARGS &
            echo $! > "$PID_FILE"
        else
            "$DAEMON_PATH" &
            echo $! > "$PID_FILE"
        fi
        result=$?
    fi

    if [ $result -eq 0 ]; then
        log_message "$DAEMON_NAME started successfully"
        return 0
    else
        log_message "Failed to start $DAEMON_NAME (exit code: $result)"
        return 1
    fi
}

# Function to check restart limits
check_restart_limits() {
    local current_time=$(date +%s)
    local restart_count=0
    local window_start=$((current_time - RESTART_WINDOW))

    # Count successful starts within the restart window
    if [ -f "$LOG_FILE" ]; then
        # Count starts in the last RESTART_WINDOW seconds
        restart_count=$(awk -v window_start="$window_start" '
            BEGIN { count = 0 }
            {
                # Extract timestamp from log line [YYYY-MM-DD HH:MM:SS]
                if (match($0, /\[([0-9]{4}-[0-9]{2}-[0-9]{2} [0-9]{2}:[0-9]{2}:[0-9]{2})\]/, arr)) {
                    timestamp = arr[1]
                    # Convert to epoch time
                    cmd = "date -d \"" timestamp "\" +%s 2>/dev/null"
                    cmd | getline epoch_time
                    close(cmd)
                    if (epoch_time >= window_start && $0 ~ /started successfully/) {
                        count++
                    }
                }
            }
            END { print count }
        ' "$LOG_FILE" 2>/dev/null || echo "0")
    fi

    if [ "$restart_count" -ge "$MAX_RESTARTS" ]; then
        log_message "Too many restarts ($restart_count) in $RESTART_WINDOW seconds. Watchdog will exit."
        return 1
    fi

    return 0
}

# Function to stop daemon (double-check process)
stop_daemon() {
    log_message "Stopping $DAEMON_NAME daemon..."

    if [ -f "$PID_FILE" ]; then
        local pid=$(cat "$PID_FILE")
        # Double-check: verify the process is actually our daemon
        if ps -p "$pid" -o comm= 2>/dev/null | grep -q "^wssshc$"; then
            start-stop-daemon --stop --quiet --pidfile "$PID_FILE" --retry=TERM/30/KILL/5
            local result=$?
            rm -f "$PID_FILE"
            return $result
        else
            log_message "PID file exists but process $pid is not wssshc"
            rm -f "$PID_FILE"
            return 1
        fi
    else
        log_message "PID file not found, daemon may not be running"
        return 1
    fi
}


# Function to cleanup on exit
cleanup() {
    log_message "Watchdog shutting down..."
    if [ -f "$WATCHDOG_PID_FILE" ]; then
        rm -f "$WATCHDOG_PID_FILE"
    fi
    exit 0
}

# Trap signals
trap cleanup SIGTERM SIGINT

# Main watchdog function (continuous monitoring)
main() {
    # Debug: Log the START value
    log_message "START configuration value: '$START'"

    # Check if START is enabled (accept various forms: yes, YES, Y, 1, true, TRUE)
    START_LOWER=$(echo "$START" | tr '[:upper:]' '[:lower:]')
    log_message "START_LOWER: '$START_LOWER'"

    if [ "$START_LOWER" != "yes" ] && [ "$START_LOWER" != "y" ] && [ "$START_LOWER" != "1" ] && [ "$START_LOWER" != "true" ]; then
        log_message "START is not set to a valid enabled value in /etc/default/wssshc. Exiting."
        exit 0
    fi

    log_message "START validation passed, proceeding with watchdog initialization"

    log_message "Watchdog started for $DAEMON_NAME"
    log_message "Check interval: $CHECK_INTERVAL seconds"
    log_message "Max restarts: $MAX_RESTARTS per $RESTART_WINDOW seconds"

    log_message "Entering monitoring loop"
    local loop_count=0

    while true; do
        loop_count=$((loop_count + 1))
        log_message "Monitoring loop iteration $loop_count"

        if ! is_daemon_running; then
            log_message "$DAEMON_NAME is not running"

            # Check restart limits before attempting to start
            if ! check_restart_limits; then
                log_message "Restart limits exceeded, watchdog will exit"
                break
            fi

            # Attempt to start daemon
            if start_daemon; then
                log_message "$DAEMON_NAME restarted successfully"
                # Give daemon time to fully start before checking
                sleep 3
            else
                log_message "Failed to restart $DAEMON_NAME"
                # If daemon fails to start, exit watchdog
                log_message "Watchdog exiting due to daemon restart failure"
                break
            fi
        else
            log_message "$DAEMON_NAME is running, monitoring continues"
        fi

        log_message "Sleeping for $CHECK_INTERVAL seconds"
        sleep "$CHECK_INTERVAL"
    done

    log_message "Watchdog exiting"
    cleanup
}

# Handle command line arguments
case "$1" in
    start)
        echo "Starting watchdog..."
        if [ -f "$WATCHDOG_PID_FILE" ]; then
            echo "Watchdog is already running"
            exit 1
        fi
        echo "Calling main() function..."

        # Run main() function in background and capture its PID
        main &
        MAIN_PID=$!

        # Create PID file with the background process PID
        echo $MAIN_PID > "$WATCHDOG_PID_FILE"

        # Give it a moment to start
        sleep 1

        # Check if the process is still running
        if kill -0 $MAIN_PID 2>/dev/null; then
            echo "Watchdog started successfully (PID: $MAIN_PID)"
            exit 0  # Exit immediately to avoid killing background main process
        else
            echo "Watchdog failed to start"
            rm -f "$WATCHDOG_PID_FILE"
            exit 1
        fi
        ;;
    stop)
        # First, try to stop the process from PID file
        if [ -f "$WATCHDOG_PID_FILE" ]; then
            watchdog_pid=$(cat "$WATCHDOG_PID_FILE")
            # Double-check: verify the process is actually our watchdog
            if ps -p "$watchdog_pid" -o comm= 2>/dev/null | grep -q "wssshc-watchdog"; then
                kill "$watchdog_pid" 2>/dev/null
                rm -f "$WATCHDOG_PID_FILE"
                echo "Watchdog stopped (from PID file)"
            else
                echo "PID file exists but process is not wssshc-watchdog"
                rm -f "$WATCHDOG_PID_FILE"
            fi
        fi

        # Also kill any other watchdog processes that might be running
        # Find all bash processes running our watchdog script
        watchdog_processes=$(ps aux | grep "/usr/sbin/wssshc-watchdog start" | grep -v grep | awk '{print $2}')
        if [ -n "$watchdog_processes" ]; then
            echo "Found additional watchdog processes: $watchdog_processes"
            for pid in $watchdog_processes; do
                if [ "$pid" != "$$" ]; then  # Don't kill ourselves
                    kill "$pid" 2>/dev/null && echo "Killed watchdog process $pid"
                fi
            done
        fi

        # Clean up any remaining PID files
        rm -f "$WATCHDOG_PID_FILE" 2>/dev/null

        echo "Watchdog stop process completed"
        ;;
    status)
        if [ -f "$WATCHDOG_PID_FILE" ]; then
            watchdog_pid=$(cat "$WATCHDOG_PID_FILE")
            if kill -0 "$watchdog_pid" 2>/dev/null; then
                # Double-check: verify the process is actually our watchdog
                if ps -p "$watchdog_pid" -o comm= 2>/dev/null | grep -q "wssshc-watchdog"; then
                    echo "Watchdog is running (PID: $watchdog_pid)"
                else
                    echo "Watchdog PID file exists but process is not wssshc-watchdog"
                    rm -f "$WATCHDOG_PID_FILE"
                fi
            else
                echo "Watchdog PID file exists but process is not running"
                rm -f "$WATCHDOG_PID_FILE"
            fi
        else
            echo "Watchdog is not running"
        fi

        # Also check daemon status
        if is_daemon_running; then
            daemon_pid=$(cat "$PID_FILE")
            echo "Daemon is running (PID: $daemon_pid)"
        else
            echo "Daemon is not running"
        fi
        ;;
    restart)
        $0 stop
        sleep 2
        $0 start
        ;;
    *)
        echo "Usage: $0 {start|stop|status|restart}"
        exit 1
        ;;
esac

exit 0