#!/bin/bash
# NOTE All keys are referring to root users
# Source the sourcer.sh file with specific arguments
source /vagrant/sourcer.sh variables.sh

# DESCRIPTION #
log $DEBUG_LEVEL $VERBOSE_INFO " Setting up the ssh key of th VM ..."

# Function to generate SSH keys if they don't exist
# NOTE root keys of VM to use for stfp server : ssh connection from VM to host
generate_ssh_keys() {
  log $DEBUG_LEVEL $VERBOSE_INFO "Generating SSH keys..."
  if [ ! -f /root/.ssh/$KEY_TYPE_SSH_VM ]; then
    log $DEBUG_LEVEL $VERBOSE_DEBUG "SSH keys do not exist. Generating new keys..."
    ssh-keygen -t rsa -b 2048 -f /root/.ssh/$KEY_TYPE_SSH_VM -C "$HOSTNAME" -N ""

  else
    log $DEBUG_LEVEL $VERBOSE_DEBUG "SSH keys already exist."
  fi
}

# Function to collect the public key and store it as <hostname>.pub
collect_public_key() {
  local hostname=$1
  log $DEBUG_LEVEL $VERBOSE_INFO "Collecting public key for hostname: $hostname"

  # Define the target directory and file path
  local target_dir="$PATH_TO_SSH_KEYS_FOLDER"
  local target_key_file="$target_dir/${hostname}.pub"

  # Check if the target directory exists; if not, log a warning and create it
  if [ ! -d "$target_dir" ]; then
    log $DEBUG_LEVEL $VERBOSE_WARNING "Target directory $target_dir does not exist. Creating it now."
    mkdir -p "$target_dir"
  fi

  # Copy the existing public key file to the target location with the desired filename
  cp /root/.ssh/$KEY_TYPE_SSH_VM.pub "$target_key_file"

  # Verify the copy and log the results
  if [ -s "$target_key_file" ]; then
    log $DEBUG_LEVEL $VERBOSE_INFO "Public key collected and stored in $target_key_file"
  else
    log $DEBUG_LEVEL $VERBOSE_ERROR "Failed to store public key in $target_key_file"
  fi
}


# Function to distribute the host public key to the specified VM
# NOTE ssh connection from home host to root VM
distribute_host_public_key() {
  local host_vm_name=$1
  local keys_file="$PATH_TO_SSH_HOST_KEY"

  # Check if the host_vm_name is provided
  if [ -z "$host_vm_name" ]; then
    log $DEBUG_LEVEL $VERBOSE_ERROR "Error: Host VM name is not provided."
    return 1
  fi

  log $DEBUG_LEVEL $VERBOSE_INFO "Distributing host public key to VM: $host_vm_name"

  # Create the .ssh directory if it doesn't exist
  log $DEBUG_LEVEL $VERBOSE_DEBUG "Creating .ssh directory if it doesn't exist..."
  mkdir -p /root/.ssh
  chmod 700 /root/.ssh

  # Create the authorized_keys file if it doesn't exist
  if [ ! -f /root/.ssh/authorized_keys ]; then
    log $DEBUG_LEVEL $VERBOSE_DEBUG "Creating the authorized_keys file..."
    touch /root/.ssh/authorized_keys
    chmod 600 /root/.ssh/authorized_keys
  fi
    # Add the public key of the host machine to the authorized_keys file
  log $DEBUG_LEVEL $VERBOSE_DEBUG "Adding the public key of the host machine to the authorized_keys file..."
  # Copy the entire keys_file content to /root/.ssh/authorized_keys
  cat "$keys_file" | sudo tee -a /root/.ssh/authorized_keys > /dev/null

  log $DEBUG_LEVEL $VERBOSE_INFO "Host key from $keys_file has been copied to /root/.ssh/authorized_keys of VM: $host_vm_name"

  # Print the contents of the authorized_keys file for debugging
  log $DEBUG_LEVEL $VERBOSE_DEBUG "Contents of /root/.ssh/authorized_keys: $(cat /root/.ssh/authorized_keys)"
 

}

# Main script
if [ "$1" == "generate" ]; then
  log $DEBUG_LEVEL $VERBOSE_INFO "Running generate_ssh_keys and collect_public_key..."
  generate_ssh_keys
  collect_public_key "$2"
elif [ "$1" == "distribute" ]; then
  log $DEBUG_LEVEL $VERBOSE_INFO "Running distribute_host_public_key..."
  distribute_host_public_key "$2"
else
  log $DEBUG_LEVEL $VERBOSE_ERROR "Usage: $0 {generate|distribute} <hostname>"
  exit 1
fi
