#!/bin/bash

# Function to generate SSH keys
generate_ssh_keys() {
    dialog --yesno "Do you want to generate SSH keys?" 10 60

    # Check if user canceled
    if [ $? -ne 0 ]; then
        return 1
    fi

    SSH_DIR=~/.ssh
    mkdir -p $SSH_DIR
    SSH_KEY_FILE=$SSH_DIR/id_rsa
    ssh-keygen -t rsa -b 4096 -f $SSH_KEY_FILE -N ""

    dialog --msgbox "SSH keys generated in $SSH_KEY_FILE." 10 60
}

# Function to add SSH key to authorized_keys
add_ssh_key() {
    KEY_FILE=$(dialog --inputbox "Enter the path to the SSH public key file:" 10 60 3>&1 1>&2 2>&3 3>&-)

    # Check if user canceled
    if [ $? -ne 0 ]; then
        return 1
    fi

    if [ -z "$KEY_FILE" ]; then
        dialog --msgbox "Please enter a valid path to the SSH public key file." 10 60
    else
        mkdir -p ~/.ssh
        cat "$KEY_FILE" >> ~/.ssh/authorized_keys
        dialog --msgbox "SSH key added to authorized_keys." 10 60
    fi
}

# Function to configure advanced SSH options
configure_advanced_ssh() {
    while true; do
        ADV_CHOICE=$(dialog --clear --backtitle "SSH Advanced Configuration" \
                            --title "SSH Advanced Menu" \
                            --menu "Choose an option:" 20 60 10 \
                            1 "Edit SSH Config File" \
                            2 "Restart SSH Service" \
                            3 "Check SSH Service Status" \
                            4 "Install OpenSSH Server" \
                            5 "Uninstall OpenSSH Server" \
                            6 "View SSH Logs" \
                            7 "Backup SSH Config" \
                            8 "Restore SSH Config" \
                            9 "Change SSH Port" \
                            10 "Enable/Disable Password Authentication" \
                            11 "List SSH Connections" \
                            12 "Back to SSH Menu" \
                            3>&1 1>&2 2>&3 3>&-)

        # Check if user canceled
        if [ $? -ne 0 ]; then
            break
        fi

        case $ADV_CHOICE in
            1) ${EDITOR:-nano} /etc/ssh/sshd_config ;;
            2) restart_ssh_service ;;
            3) check_ssh_service_status ;;
            4) install_openssh_server ;;
            5) uninstall_openssh_server ;;
            6) view_ssh_logs ;;
            7) backup_ssh_config ;;
            8) restore_ssh_config ;;
            9) change_ssh_port ;;
            10) toggle_password_authentication ;;
            11) list_ssh_connections ;;
            12) break ;;
            *) dialog --msgbox "Invalid option. Please choose a valid option." 10 60 ;;
        esac
    done
}

# Function to install OpenSSH Server
install_openssh_server() {
    dialog --yesno "Do you want to install OpenSSH Server?" 10 60

    # Check if user canceled
    if [ $? -ne 0 ]; then
        return 1
    fi

    case $INIT_SYSTEM in
        "sysvinit")
            apt update
            apt install -y openssh-server
            service ssh start
            update-rc.d ssh defaults
            ;;
        "openrc")
            apk update
            apk add openssh
            rc-update add sshd
            rc-service sshd start
            ;;
        "runit")
            apk update
            apk add openssh
            ln -s /etc/runit/sv/sshd /run/runit/service
            ;;
        *)
            dialog --msgbox "Unsupported init system: $INIT_SYSTEM" 10 60
            return 1
            ;;
    esac

    dialog --msgbox "OpenSSH Server installed and enabled." 10 60
}

# Function to uninstall OpenSSH Server
uninstall_openssh_server() {
    dialog --yesno "Do you want to uninstall OpenSSH Server?" 10 60

    # Check if user canceled
    if [ $? -ne 0 ]; then
        return 1
    fi

    case $INIT_SYSTEM in
        "sysvinit")
            service ssh stop
            update-rc.d -f ssh remove
            apt remove -y openssh-server
            ;;
        "openrc")
            rc-service sshd stop
            rc-update del sshd
            apk del openssh
            ;;
        "runit")
            rm /run/runit/service/sshd
            apk del openssh
            ;;
        *)
            dialog --msgbox "Unsupported init system: $INIT_SYSTEM" 10 60
            return 1
            ;;
    esac

    dialog --msgbox "OpenSSH Server uninstalled." 10 60
}

# Function to view SSH logs
view_ssh_logs() {
    if command -v journalctl >/dev/null 2>&1; then
        dialog --msgbox "$(journalctl -u ssh --no-pager)" 30 80
    else
        dialog --msgbox "Journalctl not found. Cannot view SSH logs." 10 60
    fi
}

# Function to backup SSH config
backup_ssh_config() {
    BACKUP_DIR="/var/backups"
    BACKUP_FILE="sshd_config_$(date +"%Y%m%d_%H%M%S").bak"
    cp /etc/ssh/sshd_config $BACKUP_DIR/$BACKUP_FILE
    dialog --msgbox "SSH config backed up to $BACKUP_DIR/$BACKUP_FILE." 10 60
}

# Function to restore SSH config
restore_ssh_config() {
    BACKUP_FILE=$(dialog --inputbox "Enter the full path of the SSH config backup file:" 10 60 3>&1 1>&2 2>&3 3>&-)

    # Check if user canceled
    if [ $? -ne 0 ]; then
        return 1
    fi

    if [ -z "$BACKUP_FILE" ]; then
        dialog --msgbox "Please enter a valid path to the SSH config backup file." 10 60
    elif [ ! -f "$BACKUP_FILE" ]; then
        dialog --msgbox "Backup file not found." 10 60
    else
        cp "$BACKUP_FILE" /etc/ssh/sshd_config

        case $INIT_SYSTEM in
            "sysvinit")
                service ssh restart
                ;;
            "openrc")
                rc-service sshd restart
                ;;
            "runit")
                sv restart sshd
                ;;
            *)
                dialog --msgbox "Unsupported init system: $INIT_SYSTEM" 10 60
                return 1
                ;;
        esac

        dialog --msgbox "SSH config restored from $BACKUP_FILE." 10 60
    fi
}

# Function to change SSH port
change_ssh_port() {
    NEW_PORT=$(dialog --inputbox "Enter the new SSH port:" 10 60 3>&1 1>&2 2>&3 3>&-)

    # Check if user canceled
    if [ $? -ne 0 ]; then
        return 1
    fi

    if [[ $NEW_PORT =~ ^[0-9]+$ ]]; then
        sed -i "s/#Port 22/Port $NEW_PORT/" /etc/ssh/sshd_config

        case $INIT_SYSTEM in
            "sysvinit")
                service ssh restart
                ;;
            "openrc")
                rc-service sshd restart
                ;;
            "runit")
                sv restart sshd
                ;;
            *)
                dialog --msgbox "Unsupported init system: $INIT_SYSTEM" 10 60
                return 1
                ;;
        esac

        dialog --msgbox "SSH port changed to $NEW_PORT." 10 60
    else
        dialog --msgbox "Invalid port number." 10 60
    fi
}

# Function to enable/disable password authentication in SSH
toggle_password_authentication() {
    STATUS=$(dialog --menu "Choose an option:" 10 60 2 \
                    1 "Enable Password Authentication" \
                    2 "Disable Password Authentication" \
                    3>&1 1>&2 2>&3 3>&-)

    # Check if user canceled
    if [ $? -ne 0 ]; then
        return 1
    fi

    case $STATUS in
        1) sed -i 's/PasswordAuthentication no/PasswordAuthentication yes/' /etc/ssh/sshd_config ;;
        2) sed -i 's/PasswordAuthentication yes/PasswordAuthentication no/' /etc/ssh/sshd_config ;;
    esac

    case $INIT_SYSTEM in
        "sysvinit")
            service ssh restart
            ;;
        "openrc")
            rc-service sshd restart
            ;;
        "runit")
            sv restart sshd
            ;;
        *)
            dialog --msgbox "Unsupported init system: $INIT_SYSTEM" 10 60
            return 1
            ;;
    esac

    dialog --msgbox "Password authentication $(echo $STATUS | tr '[:lower:]' '[:upper:]')." 10 60
}

# Function to list SSH connections
list_ssh_connections() {
    netstat -tnpa | grep 'ESTABLISHED.*sshd'
    dialog --msgbox "List of SSH connections displayed." 10 60
}

# Function to check if OpenSSH is installed
check_ssh_installation() {
    case $INIT_SYSTEM in
        "sysvinit")
            dpkg -l openssh-server >/dev/null 2>&1
            ;;
        "openrc")
            apk info openssh >/dev/null 2>&1
            ;;
        "runit")
            apk info openssh >/dev/null 2>&1
            ;;
        *)
            dialog --msgbox "Unsupported init system: $INIT_SYSTEM" 10 60
            return 1
            ;;
    esac

    if [ $? -eq 0 ]; then
        dialog --msgbox "OpenSSH is installed." 10 60
    else
        dialog --msgbox "OpenSSH is not installed." 10 60
    fi
}

# Function to configure SSH options
configure_ssh() {
    while true; do
        SSH_CHOICE=$(dialog --clear --backtitle "SSH Configuration" \
                            --title "SSH Menu" \
                            --menu "Choose an option:" 20 60 10 \
                            1 "Generate SSH Keys" \
                            2 "Add SSH Key to authorized_keys" \
                            3 "Advanced Configuration" \
                            4 "Check SSH Installation" \
                            5 "Back to Main Menu" \
                            3>&1 1>&2 2>&3 3>&-)

        # Check if user canceled
        if [ $? -ne 0 ]; then
            break
        fi

        case $SSH_CHOICE in
            1) generate_ssh_keys ;;
            2) add_ssh_key ;;
            3) configure_advanced_ssh ;;
            4) check_ssh_installation ;;
            5) break ;;
            *) dialog --msgbox "Invalid option. Please choose a valid option." 10 60 ;;
        esac
    done
}

# Determine the init system and execute the main function to configure SSH
if command -v systemctl >/dev/null 2>&1; then
    INIT_SYSTEM="sysvinit"
elif command -v rc-service >/dev/null 2>&1; then
    INIT_SYSTEM="openrc"
elif command -v sv >/dev/null 2>&1; then
    INIT_SYSTEM="runit"
else
    dialog --msgbox "Unsupported init system." 10 60
    exit 1
fi

configure_ssh