#!/bin/bash

#===============================================================================
# OpenSSH Server Setup Script for Ubuntu 22.04
# Version: 1.0.0
# Hosted at: https://scripts.mavenmusic.network/openssh/openssh-server.sh
#===============================================================================
# This script automates the installation and configuration of OpenSSH Server
# on Ubuntu 22.04, including firewall configuration and root login setup.
#===============================================================================

set -euo pipefail
IFS=$'\n\t'

# Script metadata
SCRIPT_NAME="OpenSSH Server Setup"
SCRIPT_VERSION="1.0.0"
SCRIPT_URL="https://scripts.mavenmusic.network/openssh/openssh-server.sh"
GITHUB_REPO="https://github.com/mavenmusic/scripts"  # Update with your actual repo

# Color codes for output
readonly RED='\033[0;31m'
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[1;33m'
readonly BLUE='\033[0;34m'
readonly MAGENTA='\033[0;35m'
readonly CYAN='\033[0;36m'
readonly WHITE='\033[1;37m'
readonly NC='\033[0m' # No Color
readonly BOLD='\033[1m'

# Logging configuration
readonly LOG_DIR="/var/log/openssh-setup"
readonly LOG_FILE="${LOG_DIR}/setup-$(date +%Y%m%d-%H%M%S).log"
readonly SSH_CONFIG="/etc/ssh/sshd_config"
readonly SSH_CONFIG_DIR="/etc/ssh/sshd_config.d"

#===============================================================================
# Utility Functions
#===============================================================================

log() {
    local level="$1"
    shift
    local message="$*"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[${timestamp}] [${level}] ${message}" | tee -a "$LOG_FILE"
}

log_info() { log "INFO" "$@"; }
log_warn() { log "WARN" "$@" >&2; }
log_error() { log "ERROR" "$@" >&2; }
log_success() { log "SUCCESS" "$@"; }

print_header() {
    local message="$1"
    local length=${#message}
    local padding=$(( (60 - length) / 2 ))
    
    echo
    printf "${BLUE}${BOLD}%${padding}s" | tr ' ' '='
    printf " %s " "$message"
    printf '%*s${NC}\n' $(( 60 - length - padding )) | tr ' ' '='
    echo
}

print_step() {
    echo -e "${GREEN}[✓]${NC} ${BOLD}$1${NC}"
}

print_warning() {
    echo -e "${YELLOW}[!]${NC} ${BOLD}$1${NC}"
}

print_error() {
    echo -e "${RED}[✗]${NC} ${BOLD}$1${NC}"
}

print_info() {
    echo -e "${BLUE}[i]${NC} $1"
}

show_banner() {
    clear
    echo -e "${BLUE}${BOLD}"
    echo "╔══════════════════════════════════════════════════════╗"
    echo "║         OpenSSH Server Setup for Ubuntu 22.04        ║"
    echo "║                  Version ${SCRIPT_VERSION}                         ║"
    echo "╚══════════════════════════════════════════════════════╝"
    echo -e "${NC}"
    echo -e "${CYAN}  Maven Music Network - System Administration Scripts${NC}"
    echo -e "${CYAN}  ${SCRIPT_URL}${NC}"
    echo
}

check_root() {
    if [[ $EUID -ne 0 ]]; then
        print_error "This script must be run as root or with sudo"
        echo
        echo "Usage: sudo bash <(curl -s ${SCRIPT_URL})"
        echo "   or: sudo bash openssh-server.sh"
        exit 1
    fi
}

setup_logging() {
    mkdir -p "$LOG_DIR"
    touch "$LOG_FILE"
    chmod 640 "$LOG_FILE"
    log_info "Starting ${SCRIPT_NAME} v${SCRIPT_VERSION}"
    log_info "Log file: ${LOG_FILE}"
}

backup_file() {
    local file="$1"
    if [[ -f "$file" ]]; then
        local backup="${file}.backup.$(date +%Y%m%d_%H%M%S)"
        cp "$file" "$backup"
        log_info "Backup created: $backup"
        echo "$backup"
    fi
}

validate_ssh_config() {
    if sshd -t 2>/dev/null; then
        log_info "SSH configuration validation passed"
        return 0
    else
        log_error "SSH configuration validation failed"
        sshd -t 2>&1 | tee -a "$LOG_FILE"
        return 1
    fi
}

#===============================================================================
# System Check Functions
#===============================================================================

check_ubuntu_version() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        if [[ "${ID}" != "ubuntu" ]] || [[ "${VERSION_ID}" != "22.04" ]]; then
            print_warning "This script is designed for Ubuntu 22.04"
            print_info "Detected: ${PRETTY_NAME}"
            print_info "The script may still work, but it's not tested on your version"
            echo
            
            read -p "Do you want to continue anyway? (y/N): " -n 1 -r
            echo
            if [[ ! $REPLY =~ ^[Yy]$ ]]; then
                log_info "User chose to exit - unsupported OS version"
                exit 0
            fi
        fi
    fi
}

check_internet_connection() {
    log_info "Checking internet connectivity..."
    if ping -c 2 8.8.8.8 &>/dev/null || ping -c 2 1.1.1.1 &>/dev/null; then
        log_success "Internet connection available"
        return 0
    else
        log_warn "No internet connection detected. Package installation may fail."
        return 1
    fi
}

#===============================================================================
# Installation Functions
#===============================================================================

install_openssh_server() {
    print_header "Installing OpenSSH Server"
    
    # Update package list
    log_info "Updating package lists..."
    if ! apt update -y &>> "$LOG_FILE"; then
        log_error "Failed to update package lists"
        print_error "Failed to update package lists. Check $LOG_FILE for details."
        return 1
    fi
    print_step "Package lists updated"
    
    # Check if already installed
    if dpkg -l | grep -q "^ii.*openssh-server"; then
        local version=$(dpkg -l | grep openssh-server | awk '{print $3}')
        print_info "OpenSSH Server is already installed (version: ${version})"
        log_info "OpenSSH Server already installed: $version"
        return 0
    fi
    
    # Install OpenSSH Server
    log_info "Installing OpenSSH Server..."
    print_info "Installing OpenSSH Server package..."
    
    if apt install -y openssh-server &>> "$LOG_FILE"; then
        local version=$(dpkg -l | grep openssh-server | awk '{print $3}')
        print_step "OpenSSH Server installed successfully (version: ${version})"
        log_success "OpenSSH Server installed: $version"
        return 0
    else
        log_error "Failed to install OpenSSH Server"
        print_error "Failed to install OpenSSH Server"
        return 1
    fi
}

configure_ssh_service() {
    print_header "Configuring SSH Service"
    
    # Start SSH service
    log_info "Starting SSH service..."
    if systemctl start ssh &>> "$LOG_FILE"; then
        print_step "SSH service started"
    else
        log_error "Failed to start SSH service"
        print_error "Failed to start SSH service"
        return 1
    fi
    
    # Enable SSH on boot
    log_info "Enabling SSH service on boot..."
    if systemctl enable ssh &>> "$LOG_FILE"; then
        print_step "SSH service enabled on boot"
    else
        log_warn "Failed to enable SSH service on boot"
        print_warning "SSH service will not start automatically on boot"
    fi
    
    # Check service status
    if systemctl is-active --quiet ssh; then
        print_step "SSH service is running"
        log_success "SSH service is active and running"
    else
        log_error "SSH service is not running"
        print_error "SSH service failed to start"
        systemctl status ssh --no-pager -l &>> "$LOG_FILE"
        return 1
    fi
    
    return 0
}

#===============================================================================
# Firewall Configuration
#===============================================================================

configure_firewall() {
    print_header "Configuring Firewall"
    
    local ssh_port=22
    
    # Try UFW first (preferred for Ubuntu)
    if command -v ufw &> /dev/null; then
        configure_ufw "$ssh_port"
    else
        log_warn "UFW not found, using iptables"
        configure_iptables "$ssh_port"
    fi
}

configure_ufw() {
    local port="$1"
    
    log_info "Configuring UFW firewall..."
    
    # Check if UFW is active
    if ! ufw status | grep -q "Status: active"; then
        print_warning "UFW is not active. Enabling UFW..."
        log_info "Enabling UFW..."
        
        # Allow SSH before enabling to prevent lockout
        ufw allow "${port}/tcp" &>> "$LOG_FILE"
        
        # Enable UFW with force to skip confirmation
        if ufw --force enable &>> "$LOG_FILE"; then
            print_step "UFW enabled successfully"
        else
            log_error "Failed to enable UFW"
            print_error "Failed to enable UFW"
            return 1
        fi
    fi
    
    # Add SSH rule
    if ufw status | grep -q "${port}/tcp.*ALLOW"; then
        print_info "SSH port ${port} is already allowed in UFW"
    else
        log_info "Adding SSH rule to UFW..."
        if ufw allow "${port}/tcp" &>> "$LOG_FILE"; then
            print_step "SSH port ${port} allowed in UFW"
        else
            log_error "Failed to add SSH rule to UFW"
            return 1
        fi
    fi
    
    # Reload UFW
    ufw reload &>> "$LOG_FILE"
    
    # Display rules
    print_info "Current UFW rules for SSH:"
    ufw status | grep -E "SSH|${port}/tcp" | while read line; do
        echo "  ${line}"
    done
    
    log_success "UFW configured successfully"
    return 0
}

configure_iptables() {
    local port="$1"
    
    log_info "Configuring iptables firewall..."
    print_info "Adding iptables rule for SSH port ${port}..."
    
    # Check if rule already exists
    if iptables -C INPUT -p tcp --dport "${port}" -j ACCEPT &>/dev/null; then
        print_info "SSH port ${port} rule already exists in iptables"
    else
        if iptables -A INPUT -p tcp --dport "${port}" -j ACCEPT &>> "$LOG_FILE"; then
            print_step "SSH port ${port} allowed in iptables"
        else
            log_error "Failed to add iptables rule"
            return 1
        fi
        
        # Save iptables rules
        if command -v netfilter-persistent &> /dev/null; then
            netfilter-persistent save &>> "$LOG_FILE"
        elif [ -f /etc/iptables/rules.v4 ]; then
            iptables-save > /etc/iptables/rules.v4
        fi
    fi
    
    log_success "iptables configured successfully"
    return 0
}

#===============================================================================
# Root Login Configuration
#===============================================================================

enable_root_login() {
    print_header "Configuring Root Login"
    
    print_warning "Enabling root login via SSH is not recommended for production"
    print_warning "Consider using key-based authentication or a non-root user"
    echo
    
    read -p "Do you want to enable root login? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        print_info "Root login will NOT be enabled"
        log_info "User chose not to enable root login"
        return 0
    fi
    
    # Backup SSH config
    local backup_file=$(backup_file "$SSH_CONFIG")
    
    # Enable root login
    log_info "Enabling root login in SSH configuration..."
    if grep -q "^PermitRootLogin" "$SSH_CONFIG"; then
        sed -i 's/^PermitRootLogin.*/PermitRootLogin yes/' "$SSH_CONFIG"
    else
        echo "PermitRootLogin yes" >> "$SSH_CONFIG"
    fi
    
    # Ensure password authentication is enabled
    if grep -q "^PasswordAuthentication" "$SSH_CONFIG"; then
        sed -i 's/^PasswordAuthentication.*/PasswordAuthentication yes/' "$SSH_CONFIG"
    else
        echo "PasswordAuthentication yes" >> "$SSH_CONFIG"
    fi
    
    # Check for override files
    if [[ -d "$SSH_CONFIG_DIR" ]]; then
        print_warning "Found SSH config directory: $SSH_CONFIG_DIR"
        print_info "Checking for conflicting configurations..."
        
        for conf_file in "$SSH_CONFIG_DIR"/*.conf; do
            if [[ -f "$conf_file" ]]; then
                if grep -q "PermitRootLogin" "$conf_file" 2>/dev/null; then
                    print_warning "Found PermitRootLogin in $conf_file"
                    log_warn "Override file may conflict: $conf_file"
                fi
            fi
        done
    fi
    
    # Validate SSH config
    if validate_ssh_config; then
        print_step "SSH configuration validated"
    else
        print_error "SSH configuration is invalid. Rolling back..."
        if [[ -n "$backup_file" ]]; then
            cp "$backup_file" "$SSH_CONFIG"
        fi
        return 1
    fi
    
    # Set root password if needed
    check_root_password
    
    # Restart SSH service
    log_info "Restarting SSH service to apply changes..."
    if systemctl restart ssh &>> "$LOG_FILE"; then
        print_step "SSH service restarted successfully"
        log_success "Root login enabled successfully"
    else
        log_error "Failed to restart SSH service. Rolling back..."
        if [[ -n "$backup_file" ]]; then
            cp "$backup_file" "$SSH_CONFIG"
            systemctl restart ssh
        fi
        return 1
    fi
    
    return 0
}

check_root_password() {
    log_info "Checking root password status..."
    
    local passwd_status=$(passwd -S root 2>/dev/null | awk '{print $2}')
    
    case "$passwd_status" in
        P)
            print_info "Root password is set"
            ;;
        NP|L)
            print_warning "Root password is not set or account is locked"
            read -p "Do you want to set the root password now? (y/N): " -n 1 -r
            echo
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                print_info "Please enter the new root password:"
                if passwd root; then
                    print_step "Root password set successfully"
                    log_success "Root password updated"
                else
                    print_error "Failed to set root password"
                    log_error "Failed to set root password"
                fi
            else
                print_warning "Root password not set. You won't be able to login as root without setting a password first"
            fi
            ;;
        *)
            print_warning "Unknown root password status: $passwd_status"
            ;;
    esac
}

#===============================================================================
# Security Recommendations
#===============================================================================

display_security_tips() {
    print_header "Security Recommendations"
    
    cat << EOF
${YELLOW}${BOLD}⚠️  IMPORTANT SECURITY RECOMMENDATIONS:${NC}

${BOLD}1. Use SSH Key Authentication${NC}
   Generate and use SSH keys instead of passwords:
   ${CYAN}ssh-keygen -t ed25519 -C "your_email@example.com"${NC}
   ${CYAN}ssh-copy-id user@your_server${NC}

${BOLD}2. Change Default SSH Port${NC}
   Edit ${CYAN}${SSH_CONFIG}${NC} and change:
   ${CYAN}Port 2222  # Instead of default 22${NC}

${BOLD}3. Install Fail2Ban${NC}
   ${CYAN}sudo apt install fail2ban -y${NC}

${BOLD}4. Disable Root Login After Setup${NC}
   Set ${CYAN}PermitRootLogin no${NC} in ${CYAN}${SSH_CONFIG}${NC}

${BOLD}5. Limit User Access${NC}
   Add ${CYAN}AllowUsers your_username${NC} to ${CYAN}${SSH_CONFIG}${NC}

${BOLD}6. Use 2FA (Two-Factor Authentication)${NC}
   ${CYAN}sudo apt install libpam-google-authenticator -y${NC}

${BOLD}7. Regular Updates${NC}
   Keep system updated: ${CYAN}sudo apt update && sudo apt upgrade -y${NC}

${BLUE}${BOLD}Configuration Backup:${NC} ${LOG_DIR}/
${BLUE}${BOLD}Log File:${NC} ${LOG_FILE}
${BLUE}${BOLD}SSH Config:${NC} ${SSH_CONFIG}

EOF
}

#===============================================================================
# Testing Functions
#===============================================================================

test_ssh_connection() {
    print_header "Testing SSH Connection"
    
    read -p "Do you want to test the SSH connection to localhost? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        return 0
    fi
    
    print_info "Testing SSH connection..."
    
    local test_user=""
    if grep -q "^PermitRootLogin yes" "$SSH_CONFIG" 2>/dev/null; then
        read -p "Test connection as root? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            test_user="root"
        fi
    fi
    
    if [[ -z "$test_user" ]]; then
        test_user="$SUDO_USER"
        if [[ -z "$test_user" ]]; then
            print_warning "No user to test with"
            return 0
        fi
    fi
    
    print_info "Testing SSH connection as $test_user..."
    if ssh -o StrictHostKeyChecking=no -o ConnectTimeout=5 -o BatchMode=no \
        "${test_user}@localhost" "echo 'Connection successful'" 2>&1 | grep -q "Connection successful"; then
        print_step "SSH connection test successful!"
        log_success "SSH connection test passed for user: $test_user"
    else
        print_warning "Could not verify SSH connection (this may be normal if password is required)"
        log_warn "SSH connection test could not be verified"
    fi
}

#===============================================================================
# Summary
#===============================================================================

display_summary() {
    print_header "Installation Summary"
    
    local ssh_version=$(dpkg -l | grep openssh-server | awk '{print $3}')
    local ssh_status=$(systemctl is-active ssh)
    local ssh_enabled=$(systemctl is-enabled ssh 2>/dev/null || echo "unknown")
    local root_login=$(grep "^PermitRootLogin" "$SSH_CONFIG" 2>/dev/null | awk '{print $2}' || echo "not configured")
    local ip_address=$(ip -4 addr show | grep -oP '(?<=inet\s)\d+(\.\d+){3}' | grep -v '127.0.0.1' | head -n 1)
    
    cat << EOF
${BOLD}OpenSSH Server Information:${NC}
  • Version: ${GREEN}${ssh_version}${NC}
  • Service Status: ${GREEN}${ssh_status}${NC}
  • Auto-start on Boot: ${GREEN}${ssh_enabled}${NC}
  • Root Login: ${YELLOW}${root_login}${NC}

${BOLD}Connection Details:${NC}
  • Server IP: ${GREEN}${ip_address}${NC}
  • SSH Port: ${GREEN}22${NC}
  • Command: ${CYAN}ssh user@${ip_address}${NC}

${BOLD}Firewall Status:${NC}
$(if command -v ufw &>/dev/null && ufw status | grep -q "active"; then
    echo "  • UFW: Active"
    ufw status | grep "22/tcp" | while read line; do echo "  • Rule: $line"; done
else
    echo "  • iptables: Configured"
    iptables -L INPUT -v -n 2>/dev/null | grep "dpt:22" | while read line; do echo "  • Rule: $line"; done
fi)

${BOLD}Logs & Backups:${NC}
  • Log File: ${LOG_FILE}
  • Backup Directory: ${LOG_DIR}/

${BOLD}Script Information:${NC}
  • Version: ${SCRIPT_VERSION}
  • URL: ${SCRIPT_URL}

EOF

    display_security_tips
    
    echo -e "${GREEN}${BOLD}✓ Setup completed successfully!${NC}"
    echo
}

#===============================================================================
# Main Function
#===============================================================================

main() {
    # Show banner
    show_banner
    
    # Initial checks
    check_root
    setup_logging
    check_ubuntu_version
    check_internet_connection
    
    # Trap for cleanup on exit
    trap 'log_info "Script execution completed (exit code: $?)"' EXIT
    
    # Execute installation steps
    install_openssh_server || {
        print_error "Failed to install OpenSSH Server"
        exit 1
    }
    
    configure_ssh_service || {
        print_error "Failed to configure SSH service"
        exit 1
    }
    
    configure_firewall || {
        print_warning "Firewall configuration had issues, but SSH may still work"
    }
    
    enable_root_login || {
        print_warning "Root login configuration had issues"
    }
    
    test_ssh_connection
    
    display_summary
    
    log_success "Script completed successfully"
}

#===============================================================================
# Script Entry Point
#===============================================================================

# Allow script to be sourced without executing
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi