#!/bin/bash

# This script can activate or deactivate the SSH service on the system.
# It doesn't matter which state the SSH service is currently in, 
# the script will always try to set the desired state.
#
# Note: We are using the ssh.service instead of the ssh.socket, mainly because it's the default configuration in case of Debian.
#       The difference between the two approaches is that when using the ssh.service sshd is started by systemd as a daemon 
#       which is then listening to the port 22.
#       When using the ssh.socket systemd itself is listening to the port 22 and whenever there is a new connection it will spawn
#       a new instance of an ssh@<...>.service which then finally creates an sshd process.
#       The ssh.socket is not started as long as the ssh.service is active. But that means also if the ssh.socket is not properly 
#       deactivated it can be that disabling the ssh.service and rebooting the systems leads to enabling the ssh.socket 
#       and therefore enabling SSH.
#       Because of that we always assure before activating or deactivating the ssh.service that the ssh.socket is disabled and masked.
# 
# Description:
#   - The script will does the following steps
#           |--> Check if the script is executed as root.
#           |--> Parsing the input parameter.
#           |--> Validate the input parameter.
#           |--> Check if ssh.socket is masked and disable it first if it is not.
#           |--> Unmask, enable and start or stop, disable and unmask the ssh.service 
#                depending on the input parameter and the current state of the ssh.service.
#

# ==========================
#      Functions
# ==========================

usage() {

	echo -e "ssh-service-manager:\n"

    echo -e "\tScript which manage the SSH server on the system.\n"

    echo -e "\tUsage: $0 [enable|disable]\n"

    echo -e "\tversion: ${SCRIPT_VERSION}\n"
}

disable_ssh_socket(){
    # Check if ssh.socket is masked and disable it first if it is not
    #   This is necessary because in case the ssh.socket is active the enabling of the 
    #   ssh.service will simply stop the socket but not disabling it. 
    #   So when you then disable the ssh.service later on again then the socket will be 
    #   started again after a reboot, so there is a potentially security risk of not disabling it properly.
    #   Because of the same reason it makes also sense to do that when disabling the ssh.service.

    # Get the state of the ssh.socket unit (enabled, disabled or masked)
    UNIT_FILE_STATE=$(systemctl show --property UnitFileState --value ssh.socket)

    if [ "${UNIT_FILE_STATE}" != "masked" ]
    then
        echo "ssh-service-manager: SSH socket is not masked and will be disabled and masked"

        # Stop and disable the ssh.socket
        #   In case the socket is disabled already then the 
        #   command will don't create any output and exit with 0
        systemctl disable --now ssh.socket

        # Mask the ssh.socket
        systemctl mask ssh.socket

        # Update state of the ssh.socket unit
        UNIT_FILE_STATE=$(systemctl show --property UnitFileState --value ssh.socket)

        # Get the state of the ssh.service
        if [ "${UNIT_FILE_STATE}" != "masked" ]
        then
            echo "ssh-service-manager: The SSH socket could not be masked. Please check the logs for more information."
            exit 1
        fi
    fi
}

# ==========================
#      Initialization
# ==========================

if [ "$(whoami)" != root ]
then
    # Inform the user that he has to run the script as root
    echo "ssh-service-manager: To execute this script you need root privileges."
    exit 1
fi

SCRIPT_VERSION="1.0.1"

# Check the amount of input parameters
if [ "$#" -ne 1 ]
then
    echo -e "ssh-service-manager: Illegal amount of input parameters.\n"
    usage
    exit 1
fi

# Get the input parameter
COMMAND="${1}"


# Validate the input parameter
if [ -z "${COMMAND}" ]
then
    echo -e "ssh-service-manager: No action was provided. Please provide the action as input parameter [enable|disable].\n"
    usage
    exit 1
fi

if [ "${COMMAND}" != "enable" ] && [ "${COMMAND}" != "disable" ]
then
    echo -e "ssh-service-manager: The provided action (${COMMAND}) is not valid. Please provide a valid command [enable|disable].\n"
    usage
    exit 1
fi

# ==========================
#      Main part
# ==========================

# Check if ssh.socket is masked and disable it first if it is not
disable_ssh_socket


if [ "${COMMAND}" == "enable" ]
then 

    # Get the state of the ssh.service
    STATE=$(systemctl show --property ActiveState --value ssh.service)
    SUB_STATE=$(systemctl show --property SubState --value ssh.service)

    # Get the state of the ssh.service unit (enabled, disabled or masked)
    UNIT_FILE_STATE=$(systemctl show --property UnitFileState --value ssh.service)

    # Check if the ssh.service is already active and running and if yes return early
    if [ "${STATE}" == "active" ] && [ "${SUB_STATE}" == "running" ] && [ "${UNIT_FILE_STATE}" == "enabled" ]
    then
        echo "ssh-service-manager: The SSH service is already running."
        exit 0
    fi

    if [ "${UNIT_FILE_STATE}" == "masked" ]
    then
        echo "ssh-service-manager: SSH service is masked and will be unmasked"

        # Unmask the ssh.service
        systemctl unmask ssh.service

        # Update state of the ssh.service unit
        UNIT_FILE_STATE=$(systemctl show --property UnitFileState --value ssh.service)
        
        # Validate if unmasking the ssh.service was successful
        if [ "${UNIT_FILE_STATE}" != "disabled" ]
        then
            echo "ssh-service-manager: The SSH service could not be unmasked. Please check the logs for more information."
            exit 1
        fi
    fi

    if [ "${UNIT_FILE_STATE}" == "disabled" ]
    then
        echo "ssh-service-manager: SSH service is disabled and will be enabled"

        # Enable the ssh.service
        systemctl enable ssh.service

        # Update state of the ssh.service unit
        UNIT_FILE_STATE=$(systemctl show --property UnitFileState --value ssh.service)

        # Validate if enabling the ssh.service was successful
        if [ "${UNIT_FILE_STATE}" != "enabled" ]
        then
            echo "ssh-service-manager: The SSH service could not be enabled. Please check the logs for more information."
            exit 1
        fi
    fi

    if [ "${UNIT_FILE_STATE}" == "enabled" ]
    then

        echo "ssh-service-manager: SSH service is enabled and will be started"

        # Start the ssh.service
        systemctl start ssh.service

        # Update the state of the ssh.service
        STATE=$(systemctl show --property ActiveState --value ssh.service)
        SUB_STATE=$(systemctl show --property SubState --value ssh.service)

        # Validate if starting the ssh.service was successful
        if [ "${STATE}" != "active" ] || [ "${SUB_STATE}" != "running" ]
        then
            echo "ssh-service-manager: The SSH service could not be started. Please check the logs for more information."
            exit 1
        fi
        echo "ssh-service-manager: The SSH service was started successfully."
    fi

    exit 0

else

    # Get the state of the ssh.service unit (enabled, disabled or masked)
    UNIT_FILE_STATE=$(systemctl show --property UnitFileState --value ssh.service)

    # Check if it is already masked and if yes return early
    if [ "${UNIT_FILE_STATE}" == "masked" ]
    then
        echo "ssh-service-manager: The SSH service is already disabled and masked."
        exit 0
    fi

    if [ "${UNIT_FILE_STATE}" == "enabled" ]
    then
        echo "ssh-service-manager: SSH service is enabled and will be disabled"

        # Disable the ssh.service
        systemctl disable --now ssh.service

        # Update the state of the ssh.service
        UNIT_FILE_STATE=$(systemctl show --property UnitFileState --value ssh.service)

        # Validate if disabling the ssh.service was successful
        if [ "${UNIT_FILE_STATE}" != "disabled" ]
        then
            echo "ssh-service-manager: The SSH service could not be disabled. Please check the logs for more information."
            exit 1
        fi
    fi

    if [ "${UNIT_FILE_STATE}" == "disabled" ]
    then

        echo "ssh-service-manager: SSH service is disabled and will be masked"

        # Mask the ssh.service
        systemctl mask ssh.service

        # Update the state of the ssh.service unit (enabled, disabled or masked)
        UNIT_FILE_STATE=$(systemctl show --property UnitFileState --value ssh.service)

        # Validate if masking the ssh.service was successful
        if [ "${UNIT_FILE_STATE}" != "masked" ]
        then
            echo "ssh-service-manager: The SSH service could not be masked. Please check the logs for more information."
            exit 1
        fi
        echo "ssh-service-manager: The SSH service was masked successfully."
    fi

    # Check for active SSH connections and terminate them
    #   After disabling and masking the ssh.service als also assuring that the ssh.socket is masked
    #   we can be sure that no new connections can be established. 
    #   But already active connections are still running and we have to kill them manually.

    ACTIVE_SSHD_PROCESSES=$(pgrep -c sshd)

    if [ "${ACTIVE_SSHD_PROCESSES}" -gt 0 ]
    then
        echo "ssh-service-manager: There are still active SSH connections which will be now terminated."

        # Try to terminated all active sshd processes 
        pkill --signal HUP sshd

        # Wait for 2 second to give the sshd processes the chance to terminate
        sleep 2

        # Check again if there are still active sshd processes
        ACTIVE_SSHD_PROCESSES=$(pgrep -c sshd)

        if [ "${ACTIVE_SSHD_PROCESSES}" -gt 0 ]
        then
            echo "ssh-service-manager: There are still active SSH connections which could not be terminated. They will be now killed."

            # Kill all active sshd processes
            pkill --signal KILL sshd

            # Wait for 1 second to assure that state was already refreshed
            sleep 1

            # Check again if there are still active sshd processes
            ACTIVE_SSHD_PROCESSES=$(pgrep -c sshd)
        fi
    fi

    if [ "${ACTIVE_SSHD_PROCESSES}" -eq 0 ]
    then
        echo "ssh-service-manager: There are no active SSH connections."
        exit 0
    else
        echo "ssh-service-manager: There are still active SSH connections which could not be killed. Please check the logs for more information."
        exit 1
    fi

fi
