#!/bin/bash

# SSH Tunneling Management Script
# This script sets up bi-directional SSH tunnels for remote access through NAT

# Configuration Variables
SSHD_ADDRESS="user@example.com"      # Public SSH server address
REMOTE_ADDRESS="192.168.1.100"       # Remote computer's internal IP
LOCAL_RDP_PORT="3389"                # Local RDP port
REMOTE_RDP_PORT="3389"               # Remote RDP port
TUNNEL_PORT="2222"                   # Tunnel port on public server
SOCKS_PORT="8765"                    # Local SOCKS proxy port
LOG_FILE="/var/log/ssh-tunnel.log"   # Log file location
KEEP_ALIVE="60"                      # SSH keepalive interval in seconds
MAX_RETRIES=3                        # Maximum number of connection retries

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Logging function
log() {
    local level=$1
    local message=$2
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo -e "${timestamp} [${level}] ${message}" >> "$LOG_FILE"
    
    case $level in
        "INFO")
            echo -e "${GREEN}${message}${NC}"
            ;;
        "WARNING")
            echo -e "${YELLOW}${message}${NC}"
            ;;
        "ERROR")
            echo -e "${RED}${message}${NC}"
            ;;
    esac
}

# Check if required commands exist
check_requirements() {
    local requirements=("ssh" "nc")
    for cmd in "${requirements[@]}"; do
        if ! command -v "$cmd" &> /dev/null; then
            log "ERROR" "Required command '$cmd' not found. Please install it."
            exit 1
        fi
    done
}

# Test SSH connection
test_ssh_connection() {
    log "INFO" "Testing SSH connection to $SSHD_ADDRESS..."
    if ! ssh -q -o BatchMode=yes -o ConnectTimeout=5 "$SSHD_ADDRESS" exit; then
        log "ERROR" "Cannot connect to SSH server. Check your credentials and connection."
        return 1
    fi
    return 0
}

# Check if port is available
check_port() {
    local port=$1
    if nc -z localhost "$port" 2>/dev/null; then
        log "WARNING" "Port $port is already in use"
        return 1
    fi
    return 0
}

# Setup reverse tunnel and SOCKS proxy
setup_reverse_tunnel() {
    log "INFO" "Setting up reverse tunnel and SOCKS proxy..."
    ssh -R "0.0.0.0:${TUNNEL_PORT}:127.0.0.1:${LOCAL_RDP_PORT}" \
        -D "$SOCKS_PORT" \
        -o "ServerAliveInterval=${KEEP_ALIVE}" \
        -o "ServerAliveCountMax=3" \
        -o "ExitOnForwardFailure=yes" \
        -N "$SSHD_ADDRESS" &
    REVERSE_PID=$!
    
    # Check if tunnel is established
    sleep 2
    if kill -0 "$REVERSE_PID" 2>/dev/null; then
        log "INFO" "Reverse tunnel established (PID: $REVERSE_PID)"
    else
        log "ERROR" "Failed to establish reverse tunnel"
        return 1
    fi
}

# Setup local tunnel
setup_local_tunnel() {
    log "INFO" "Setting up local tunnel..."
    ssh -L "${TUNNEL_PORT}:${REMOTE_ADDRESS}:${REMOTE_RDP_PORT}" \
        -o "ServerAliveInterval=${KEEP_ALIVE}" \
        -o "ServerAliveCountMax=3" \
        -o "ExitOnForwardFailure=yes" \
        -N "$SSHD_ADDRESS" &
    LOCAL_PID=$!
    
    # Check if tunnel is established
    sleep 2
    if kill -0 "$LOCAL_PID" 2>/dev/null; then
        log "INFO" "Local tunnel established (PID: $LOCAL_PID)"
    else
        log "ERROR" "Failed to establish local tunnel"
        return 1
    fi
}

# Cleanup function
cleanup() {
    log "INFO" "Cleaning up tunnels..."
    [[ -n $REVERSE_PID ]] && kill "$REVERSE_PID" 2>/dev/null
    [[ -n $LOCAL_PID ]] && kill "$LOCAL_PID" 2>/dev/null
    exit 0
}

# Handle signals
trap cleanup SIGINT SIGTERM

# Main function
main() {
    local tunnel_type=$1
    local retry_count=0
    
    # Create log directory if it doesn't exist
    mkdir -p "$(dirname "$LOG_FILE")"
    
    # Check requirements
    check_requirements
    
    # Test SSH connection
    test_ssh_connection || exit 1
    
    # Check ports
    check_port "$TUNNEL_PORT" || exit 1
    [[ "$tunnel_type" == "reverse" ]] && check_port "$SOCKS_PORT" || exit 1
    
    # Main loop with retry logic
    while true; do
        case $tunnel_type in
            "reverse")
                setup_reverse_tunnel
                ;;
            "local")
                setup_local_tunnel
                ;;
            *)
                log "ERROR" "Invalid tunnel type. Use 'reverse' or 'local'"
                exit 1
                ;;
        esac
        
        if [[ $? -ne 0 ]]; then
            ((retry_count++))
            if [[ $retry_count -ge $MAX_RETRIES ]]; then
                log "ERROR" "Maximum retry attempts reached. Exiting."
                exit 1
            fi
            log "WARNING" "Tunnel failed. Retrying in 10 seconds... (Attempt $retry_count/$MAX_RETRIES)"
            sleep 10
            continue
        fi
        
        retry_count=0
        wait
    done
}

# Display usage information
usage() {
    echo "Usage: $0 [reverse|local]"
    echo "  reverse - Setup reverse tunnel and SOCKS proxy"
    echo "  local   - Setup local tunnel for remote access"
    exit 1
}

# Check command line arguments
[[ $# -ne 1 ]] && usage

# Start the script
main "$1"