#!/bin/bash
# =============================================================================
# SSH hardening module
# =============================================================================

# Set script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Source common functions and variables
source "./common.sh"
SOURCE_KEY="./custom/custom_key"

# Function to create SSH keys for root user
create_root_ssh_keys() {
  local SSH_DIR="/root/.ssh"
  local SYMBOLIQUE_DIR="/etc/ssh/authorized_keys/"
  local FILE_AUTHORIZED_KEYS="$SSH_DIR/authorized_keys"
  local SSH_KEY_DOC="/root/root-ssh-keys-documentation.txt"

  # Create .ssh directory if it doesn't exist
  mkdir -p "$SSH_DIR"
  chmod 700 "$SSH_DIR"

  # Create or backup authorized_keys file
  if [ -f "$FILE_AUTHORIZED_KEYS" ]; then
    backup_file "$FILE_AUTHORIZED_KEYS"
  else
    touch "$FILE_AUTHORIZED_KEYS"
     cat "$SOURCE_KEY" > "$FILE_AUTHORIZED_KEYS"
  fi

  chmod 600 "$FILE_AUTHORIZED_KEYS"

  # 
  mkdir -p "$SYMBOLIQUE_DIR"
  cd "$SYMBOLIQUE_DIR"
  ln -s "$FILE_AUTHORIZED_KEYS" root

  # Create documentation
  log_message "INFO" "Creating SSH key documentation for root user"

  cat > "$SSH_KEY_DOC" << EOF
# Root SSH Keys Documentation
# 
# To add SSH public keys for root user, add them to the authorized_keys file:
# $FILE_AUTHORIZED_KEYS
#
# Format:
# ssh-rsa AAAAB3NzaC1yc2EA... comment
#
# For security:
# - Key-based authentication is more secure than password-based authentication
# - Use strong, unique keys for each user or service
# - Regularly rotate SSH keys
# - Remove keys that are no longer needed
#
# Remember to maintain proper permissions:
# chmod 700 $SSH_DIR
# chmod 600 $FILE_AUTHORIZED_KEYS
EOF

  log_message "SUCCESS" "SSH key documentation for root user created at $SSH_KEY_DOC"
}

# Function to create SSH keys for non-root users
create_non_root_ssh_keys() {
  local SSH_KEY_DOC="/etc/skel/.ssh-documentation.txt"

  # Create /etc/skel/.ssh directory
  mkdir -p "/etc/skel/.ssh"
  chmod 700 "/etc/skel/.ssh"
  touch "/etc/skel/.ssh/authorized_keys"
  chmod 600 "/etc/skel/.ssh/authorized_keys"

  # Create documentation
  log_message "INFO" "Creating SSH key documentation for non-root users"

  cat > "$SSH_KEY_DOC" << EOF
# User SSH Keys Documentation
#
# To add SSH public keys for this user, add them to the authorized_keys file:
# ~/.ssh/authorized_keys
#
# Format:
# ssh-rsa AAAAB3NzaC1yc2EA... comment
#
# For security:
# - Key-based authentication is more secure than password-based authentication
# - Use strong, unique keys for each user or service
# - Regularly rotate SSH keys
# - Remove keys that are no longer needed
#
# Remember to maintain proper permissions:
# chmod 700 ~/.ssh
# chmod 600 ~/.ssh/authorized_keys
EOF

  log_message "SUCCESS" "SSH key documentation for non-root users created at $SSH_KEY_DOC"
}

# Function to harden SSH configuration
harden_ssh_config() {
  local SSH_CONFIG="/etc/ssh/sshd_config"
  local SSH_CONFIG_BACKUP="$BACKUP_DIR/etc/ssh/sshd_config"

  # Backup current SSH configuration
  backup_file "$SSH_CONFIG"

  log_message "INFO" "Hardening SSH configuration"

  # Create a new SSH configuration with hardened settings
  cat > "$SSH_CONFIG" << EOF
# SSH Server Configuration
# Hardened configuration from security script

# Basic SSH server settings
Port 22
#AddressFamily inet
#ListenAddress 0.0.0.0
#ListenAddress ::

# Authentication settings
#AuthorizedKeysFile /root/.ssh/authorized_keys/%u
#PermitRootLogin prohibit-password
#PubkeyAuthentication yes
#PasswordAuthentication no
#PermitEmptyPasswords no
#ChallengeResponseAuthentication no
UsePAM yes

# Security settings
X11Forwarding no
PrintMotd no
AcceptEnv LANG LC_*
Subsystem sftp /usr/lib/openssh/sftp-server

# Hardening options
LoginGraceTime 30
MaxAuthTries 5
MaxSessions 3
ClientAliveInterval 300
ClientAliveCountMax 3

#
Banner /etc/banner

# Disable weak crypto
#KexAlgorithms curve25519-sha256@libssh.org,diffie-hellman-group-exchange-sha256
#Ciphers chacha20-poly1305@openssh.com,aes256-gcm@openssh.com,aes128-gcm@openssh.com,aes256-ctr,aes192-ctr,aes128-ctr
#MACs hmac-sha2-512-etm@openssh.com,hmac-sha2-256-etm@openssh.com,hmac-sha2-512,hmac-sha2-256

# Logging
SyslogFacility AUTH
LogLevel VERBOSE
EOF

  log_message "INFO" "Configured SSH to use port 2222 and disabled root password login"
  log_message "SUCCESS" "SSH configuration hardened"

  # Restart SSH service to apply changes
  log_message "INFO" "Restarting SSH service"
  service restart sshd

  if [ $? -eq 0 ]; then
    log_message "SUCCESS" "SSH service restarted successfully"
  else
    log_message "ERROR" "Failed to restart SSH service"
    
    # Revert to backup
    cp "$SSH_CONFIG_BACKUP" "$SSH_CONFIG"
    service restart sshd
    
    log_message "WARNING" "Reverted to original SSH configuration"
  fi
}

#
log_message "SCRIPT" "ssh_hardening.sh"

# Main execution for SSH hardening
create_root_ssh_keys
create_non_root_ssh_keys
harden_ssh_config

log_message "SUCCESS" "SSH hardening completed"