#!/bin/bash

if [[ $# -lt 2 ]]; then
  echo -e "\e[31m[Error]: Not enough arguments!\e[0m"
  echo -e "\e[31m[!] Position Matters!\e[0m"
  echo -e "\e[32m[+] Usage: ssh-config.sh [-i identity_file] [-p port] user@host alias\e[0m"
  exit 1
fi

while getopts "i:p:" opt; do
  case ${opt} in
    i )
      PRIVATE_KEY=${OPTARG}
      ;;
    p )
      CONF_PORT=${OPTARG}
      ;;
    \? )
      echo -e "\e[31m[Error]: Invalid option!\e[0m"
      echo -e "\e[31mUsage: SSHCONFIG [-i identity_file] [-P port] user@host alias\e[0m"
      echo -e "\e[31m[!] Position Matters!\e[0m"
      exit 1
      ;;
  esac
done
shift $((OPTIND -1))

SSH_ARGS=$1
HOST=$(echo $1 | cut -d '@' -f 2)
ALIAS=$2
PORT="${CONF_PORT:-22}"

if [[ -n "$PRIVATE_KEY" ]]; then
  cat <<EOF >> "/home/${USER}/.ssh/config"

Host $ALIAS
  HostName ${HOST} 
  User ${SSH_ARGS%@*}
  Port ${PORT}
  IdentityFile $(pwd)/$PRIVATE_KEY
EOF
  echo -e "\e[32m[Success]: SSH configuration added for $ALIAS\e[0m"
else
  cat <<EOF >> "/home/${USER}/.ssh/config"

Host $ALIAS
  HostName ${HOST} 
  User ${SSH_ARGS%@*}
  Port ${PORT}
EOF
  echo -e "\e[32m[Success]: SSH configuration added for $ALIAS\e[0m"
fi