#!/usr/bin/env bash
# SSH Server Setup and Hardening Script

set -euo pipefail

# Load environment variables
source /tmp/env-desktop

log() {
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] SSH: $1"
}

error() {
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] SSH ERROR: $1" >&2
    exit 1
}

# Detect OS
detect_os() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS=$ID
        VERSION=$VERSION_ID
    else
        error "Cannot detect OS"
    fi
    log "Detected OS: $OS $VERSION"
}

# Install SSH server
install_ssh_server() {
    log "Installing SSH server..."
    
    case $OS in
        ubuntu|debian)
            apt-get update
            apt-get install -y openssh-server openssh-client
            ;;
        centos|rhel|fedora)
            if command -v dnf >/dev/null 2>&1; then
                dnf install -y openssh-server openssh-clients
            else
                yum install -y openssh-server openssh-clients
            fi
            ;;
        *)
            error "Unsupported OS for SSH installation: $OS"
            ;;
    esac
    
    log "SSH server installed"
}

# Configure SSH server
configure_ssh_server() {
    log "Configuring SSH server..."
    
    local ssh_port="${SSH_PORT:-22}"
    local password_auth="${SSH_PASSWORD_AUTH:-yes}"
    local key_auth="${SSH_KEY_AUTH:-yes}"
    local root_login="${SSH_ROOT_LOGIN:-prohibit-password}"
    local max_auth_tries="${SSH_MAX_AUTH_TRIES:-3}"
    local client_alive_interval="${SSH_CLIENT_ALIVE_INTERVAL:-300}"
    local client_alive_count_max="${SSH_CLIENT_ALIVE_COUNT_MAX:-2}"
    
    # Backup original config
    cp /etc/ssh/sshd_config /etc/ssh/sshd_config.backup.$(date +%Y%m%d_%H%M%S)
    
    # Create new SSH configuration
    cat > /etc/ssh/sshd_config << EOF
# SSH Configuration for Desktop Environment
# Generated by provisioning system

# Connection settings
Port $ssh_port
AddressFamily any
ListenAddress 0.0.0.0
ListenAddress ::

# Host keys
HostKey /etc/ssh/ssh_host_rsa_key
HostKey /etc/ssh/ssh_host_ecdsa_key
HostKey /etc/ssh/ssh_host_ed25519_key

# Ciphers and keying
RekeyLimit default none

# Logging
SyslogFacility AUTH
LogLevel INFO

# Authentication
LoginGraceTime 2m
PermitRootLogin $root_login
StrictModes yes
MaxAuthTries $max_auth_tries
MaxSessions 10

PubkeyAuthentication $key_auth
AuthorizedKeysFile .ssh/authorized_keys .ssh/authorized_keys2

# Password authentication
PasswordAuthentication $password_auth
PermitEmptyPasswords no
ChallengeResponseAuthentication no

# Kerberos and GSSAPI (disabled for security)
KerberosAuthentication no
GSSAPIAuthentication no

# Connection timeouts
ClientAliveInterval $client_alive_interval
ClientAliveCountMax $client_alive_count_max
TCPKeepAlive yes

# Compression
Compression delayed

# Environment
AcceptEnv LANG LC_*
AcceptEnv XMODIFIERS

# X11 forwarding (enabled for desktop environment)
X11Forwarding yes
X11DisplayOffset 10
X11UseLocalhost yes

# Agent forwarding (be careful with security)
AllowAgentForwarding yes

# TCP forwarding
AllowTcpForwarding yes
GatewayPorts no

# Tunnel device forwarding
PermitTunnel no

# chroot directory
ChrootDirectory none

# Banner
Banner none

# Subsystem
Subsystem sftp /usr/lib/openssh/sftp-server

# User/Group restrictions
EOF

    # Add user restrictions if specified
    if [[ -n "${SSH_ALLOWED_USERS:-}" ]]; then
        echo "AllowUsers $SSH_ALLOWED_USERS" >> /etc/ssh/sshd_config
        log "SSH access restricted to users: $SSH_ALLOWED_USERS"
    fi
    
    if [[ -n "${SSH_DENIED_USERS:-}" ]]; then
        echo "DenyUsers $SSH_DENIED_USERS" >> /etc/ssh/sshd_config
        log "SSH access denied for users: $SSH_DENIED_USERS"
    fi
    
    # Fix sftp-server path for different distributions
    case $OS in
        ubuntu|debian)
            sed -i 's|/usr/lib/openssh/sftp-server|/usr/lib/openssh/sftp-server|' /etc/ssh/sshd_config
            ;;
        centos|rhel|fedora)
            sed -i 's|/usr/lib/openssh/sftp-server|/usr/libexec/openssh/sftp-server|' /etc/ssh/sshd_config
            ;;
    esac
    
    # Test SSH configuration
    sshd -t || error "SSH configuration is invalid"
    
    log "SSH server configured"
}

# Setup SSH keys for desktop user
setup_ssh_keys() {
    local desktop_user="${DESKTOP_USER:-desktop}"
    local desktop_home="${DESKTOP_HOME:-/home/$desktop_user}"
    
    log "Setting up SSH keys for user $desktop_user"
    
    # Create user if not exists
    if ! id "$desktop_user" &>/dev/null; then
        useradd -m -s /bin/bash "$desktop_user"
        log "Created user $desktop_user"
    fi
    
    # Create .ssh directory
    sudo -u "$desktop_user" mkdir -p "$desktop_home/.ssh"
    chmod 700 "$desktop_home/.ssh"
    
    # Generate SSH key pair if not exists
    if [[ ! -f "$desktop_home/.ssh/id_rsa" ]]; then
        log "Generating SSH key pair for $desktop_user"
        sudo -u "$desktop_user" ssh-keygen -t rsa -b 4096 -f "$desktop_home/.ssh/id_rsa" -N "" -C "$desktop_user@$(hostname)"
        log "SSH key pair generated"
    fi
    
    # Create authorized_keys file
    sudo -u "$desktop_user" touch "$desktop_home/.ssh/authorized_keys"
    chmod 600 "$desktop_home/.ssh/authorized_keys"
    
    # Set proper ownership
    chown -R "$desktop_user:$desktop_user" "$desktop_home/.ssh"
    
    log "SSH keys setup completed for $desktop_user"
}

# Setup fail2ban for SSH protection
setup_fail2ban() {
    log "Setting up fail2ban for SSH protection..."
    
    case $OS in
        ubuntu|debian)
            apt-get install -y fail2ban
            ;;
        centos|rhel|fedora)
            if command -v dnf >/dev/null 2>&1; then
                dnf install -y fail2ban
            else
                yum install -y fail2ban
            fi
            ;;
    esac
    
    # Create fail2ban configuration for SSH
    cat > /etc/fail2ban/jail.local << EOF
[DEFAULT]
# Ban time in seconds (10 minutes)
bantime = 600

# Find time window (10 minutes)
findtime = 600

# Max retry attempts
maxretry = 3

[sshd]
enabled = true
port = ${SSH_PORT:-22}
filter = sshd
logpath = /var/log/auth.log
maxretry = ${SSH_MAX_AUTH_TRIES:-3}
bantime = 3600
EOF
    
    # Start and enable fail2ban
    systemctl enable fail2ban
    systemctl start fail2ban
    
    log "fail2ban configured and started"
}

# Setup firewall rules for SSH
setup_firewall() {
    log "Setting up firewall rules for SSH"
    
    local ssh_port="${SSH_PORT:-22}"
    
    # Try different firewall tools
    if command -v ufw >/dev/null 2>&1; then
        ufw allow "$ssh_port/tcp" comment "SSH"
        log "UFW rule added for SSH port $ssh_port"
    elif command -v firewall-cmd >/dev/null 2>&1; then
        if [[ "$ssh_port" != "22" ]]; then
            firewall-cmd --permanent --add-port="$ssh_port/tcp"
        else
            firewall-cmd --permanent --add-service=ssh
        fi
        firewall-cmd --reload
        log "FirewallD rule added for SSH port $ssh_port"
    else
        log "WARNING: No supported firewall tool found. Manual firewall configuration may be needed."
    fi
}

# Start and enable SSH service
start_ssh_service() {
    log "Starting SSH service..."
    
    # Enable and start SSH service
    systemctl enable ssh sshd 2>/dev/null || systemctl enable sshd
    systemctl restart ssh sshd 2>/dev/null || systemctl restart sshd
    
    # Check service status
    if systemctl is-active --quiet ssh || systemctl is-active --quiet sshd; then
        log "SSH service is running"
    else
        error "Failed to start SSH service"
    fi
    
    log "SSH service started and enabled"
}

# Display connection information
display_connection_info() {
    local desktop_user="${DESKTOP_USER:-desktop}"
    local ssh_port="${SSH_PORT:-22}"
    local server_ip=$(ip route get 1.1.1.1 | grep -oP 'src \K\S+' 2>/dev/null || echo "$(hostname -I | awk '{print $1}')")
    
    log "SSH setup completed!"
    log ""
    log "SSH Connection Information:"
    log "  Server IP: $server_ip"
    log "  SSH Port: $ssh_port"
    log "  Desktop User: $desktop_user"
    log ""
    log "Connect via SSH:"
    log "  ssh $desktop_user@$server_ip -p $ssh_port"
    log ""
    log "Public key location (for key-based auth):"
    log "  /home/$desktop_user/.ssh/id_rsa.pub"
    log ""
    log "To copy your public key to another machine:"
    log "  ssh-copy-id -i /home/$desktop_user/.ssh/id_rsa.pub user@remote-host"
}

# Main installation function
main() {
    if [[ "${SSH_ENABLED:-true}" != "true" ]]; then
        log "SSH is disabled, skipping installation"
        return 0
    fi
    
    log "Starting SSH server installation and configuration..."
    
    detect_os
    install_ssh_server
    configure_ssh_server
    setup_ssh_keys
    setup_fail2ban
    setup_firewall
    start_ssh_service
    display_connection_info
    
    log "SSH setup completed successfully!"
}

# Run main function if script is executed directly
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi