#!/usr/bin/env bash
set -euo pipefail

# =========================
# Config
# =========================
NEW_USER="${1:-dev}"
SSH_PORT="${2:-22}"
DISABLE_PASSWORD_LOGIN="${3:-false}"
PUBKEY="${4:-}"   # optional public key

# =========================
# UI
# =========================

log(){ echo -e "\033[1;32m[INFO]\033[0m $1"; }
warn(){ echo -e "\033[1;33m[WARN]\033[0m $1"; }
err(){ echo -e "\033[1;31m[ERROR]\033[0m $1"; }

confirm(){
    read -rp "$1 [y/N]: " ans
    [[ "$ans" == "y" || "$ans" == "Y" ]]
}

pause(){
    read -rp "Press Enter to continue..."
}

# =========================
# System Detect
# =========================

detect_os(){
    if [[ -f /etc/arch-release ]]; then
        OS="arch"
    elif grep -qi "debian\|ubuntu" /etc/os-release; then
        OS="debian"
    else
        err "Unsupported OS"
        exit 1
    fi
    log "Detected OS: $OS"
}

pkg_install(){
    case "$OS" in
        arch) pacman -Sy --noconfirm "$@" ;;
        debian) apt-get update && apt-get install -y "$@" ;;
    esac
}

# =========================
# User
# =========================

create_user(){
    if id "$NEW_USER" &>/dev/null; then
        warn "User exists"
    else
        useradd -m -s /bin/bash "$NEW_USER"
        passwd "$NEW_USER"
    fi
    usermod -aG sudo "$NEW_USER" 2>/dev/null || usermod -aG wheel "$NEW_USER"
}

setup_ssh_key(){
    mkdir -p /home/$NEW_USER/.ssh
    chmod 700 /home/$NEW_USER/.ssh

    if [[ -n "$PUBKEY" ]]; then
        log "Adding provided public key"
        echo "$PUBKEY" > /home/$NEW_USER/.ssh/authorized_keys
    else
        warn "No public key provided, please paste one"
        read -rp "Paste your SSH public key: " key
        echo "$key" > /home/$NEW_USER/.ssh/authorized_keys
    fi

    chmod 600 /home/$NEW_USER/.ssh/authorized_keys
    chown -R $NEW_USER:$NEW_USER /home/$NEW_USER/.ssh
}

# =========================
# SSH Hardening
# =========================

set_ssh(){
    local key="$1"
    local val="$2"

    if grep -q "^#\?$key" /etc/ssh/sshd_config; then
        sed -i "s/^#\?$key.*/$key $val/" /etc/ssh/sshd_config
    else
        echo "$key $val" >> /etc/ssh/sshd_config
    fi
}

config_ssh(){
    cp /etc/ssh/sshd_config /etc/ssh/sshd_config.bak.$(date +%s)

    set_ssh "PermitRootLogin" "no"
    set_ssh "PubkeyAuthentication" "yes"
    set_ssh "Port" "$SSH_PORT"

    if [[ "$DISABLE_PASSWORD_LOGIN" == "true" ]]; then
        warn "Disabling password login"
        confirm "Are you sure?" && set_ssh "PasswordAuthentication" "no"
    else
        set_ssh "PasswordAuthentication" "yes"
    fi
}

restart_ssh(){
    systemctl restart ssh 2>/dev/null || systemctl restart sshd
}

# =========================
# Firewall (UFW)
# =========================

setup_ufw(){
    log "Setting up UFW firewall"

    pkg_install ufw

    ufw allow "$SSH_PORT"/tcp
    ufw allow 80/tcp
    ufw allow 443/tcp

    ufw --force enable
}

# =========================
# Fail2ban
# =========================

setup_fail2ban(){
    log "Installing fail2ban"

    pkg_install fail2ban

    cat > /etc/fail2ban/jail.local <<EOF
[sshd]
enabled = true
port = $SSH_PORT
maxretry = 5
bantime = 1h
EOF

    systemctl enable fail2ban
    systemctl restart fail2ban
}

# =========================
# Main
# =========================

main(){
    if [[ $EUID -ne 0 ]]; then
        err "Run as root"
        exit 1
    fi

    detect_os

    echo "================================="
    echo "User: $NEW_USER"
    echo "Port: $SSH_PORT"
    echo "Disable Password: $DISABLE_PASSWORD_LOGIN"
    echo "================================="

    confirm "Start?" || exit 0

    pause
    log "Create user"
    create_user
    setup_ssh_key

    pause
    log "Configure SSH"
    config_ssh

    pause
    log "Setup Firewall"
    setup_ufw

    pause
    log "Setup Fail2ban"
    setup_fail2ban

    pause
    log "Restart SSH"
    restart_ssh

    echo "================================="
    log "DONE"
    echo "Test SSH BEFORE exit!"
}

main "$@"