#!/usr/bin/env bash

#
# SSH Connection Utility (SCU)
# - ssh connection utility with pre-defined hosts list
# https://github.com/dualmi/scu
#
# This project is licensed under GPL 3 License - see [https://www.gnu.org/licenses/gpl-3.0.en.html]
#
set -euo pipefail

if (( BASH_VERSINFO[0] < 4 || (BASH_VERSINFO[0] == 4 && BASH_VERSINFO[1] < 3) )); then
  echo "SCU requires bash >= 4.3" >&2
  exit 1
fi

declare -A sshhostlist;
: "${ansiblepath:=}"
: "${netbox_hostname:=}"
: "${curl_connect_timeout:=3}"
: "${netbox_listing_limit:=1000}"
: "${netbox_cache_refresh_time:=300}"
hosts_file="$HOME/.scurc";
host_connect="${1:-}"

format_epoch() {
  local ts="$1"

  [[ -n "$ts" && "$ts" != 0 ]] || {
    echo "unknown"
    return
  }

  if date -r 0 >/dev/null 2>&1; then
    date -r "$ts" "+%Y-%m-%d %H:%M:%S"
  else
    date --date="@$ts" "+%Y-%m-%d %H:%M:%S"
  fi
}

echo -e "\n\tSSH Connection Utility with predefined hosts list\n";

# Hosts file check part.
# If file doesn't exists it can be generated when 'init' param is set.
if [[ -f $hosts_file ]]; then
  # shellcheck source=/dev/null
  source "$hosts_file"
else
  if [[ ${1:-} = "init" ]]; then
    # At next step example config file will be generated
cat > "$hosts_file" <<EOL
# Array entry format:
# [host-useful-name]="username* server.address.or.ip* ssh-port-if-not-default other-ssh-arguments"
# * - is required

# If you'll set the path to ansible host file(s), ansible hosts will be loaded firstly.
# Hosts with the same names will be overwritten (from the last loaded inventory file)

# yq - https://github.com/mikefarah/yq - must be installed if you using an ansible inventory

# You could point several space separated inventory files
ansiblepath="
$HOME/.scurc.yaml
/wrong/file/name
"

# You can provide your Netbox schema://address here.
# Also you should set a NETBOX_TOKEN variable, usually via "export NETBOX_TOKEN=..."
#netbox_hostname="https://netbox.contoso.com/"

sshhostlist=(
  # here is 4 examples to show how it will be shown in terminal
  # while you just run 'scu' w/o parameters or
  # type only part of host name like 'scu host'
  [example-host-one]="sshuser host-one.domain.local"
  [example-host-two]="root host-two.domain.local 2222"
  [example-line-one]="games server.domain.local 22"
  [example-ext-host]="admin h24.contoso.com 22022 -i ~/.ssh/id_rsa_for_contoso"
);
EOL

cat > "$hosts_file.yaml" <<EOL
# Here is the inventory example which could be used by Ansible and scu
all:
  children:
    my_company:
      vars:
        ansible_user: root
        ansible_port: 22
      children:
        firstgroup:
          hosts:
            example-host-yaml:
              ansible_host: 192.168.1.1
            example-host-yaml-args:
              ansible_user: admin
              ansible_host: 192.168.2.2
              ansible_port: 3333
              ansible_ssh_common_args: '-o ProxyCommand="ssh -W %h:%p -q ssh-tunnel@10.10.2.2 -p 2222"'
EOL

    echo "Example config files $hosts_file and $hosts_file.yaml are generated.";
    echo -e "Edit them as you wish.\n";
    exit 0;
  fi
  echo "File with list of ssh hosts ($hosts_file) not found. No way to live.";
  echo "(!) You can initialize example file - run '$0 init' to do that.";
  echo "";
  exit 1;
fi

# # If path to ansible host file (yaml format) is set, then using it and updating original list from .scurc
if [[ $ansiblepath ]]; then
  if [[ $(yq --version | awk '{ print $NF }' | sed -e 's,^[^0-9]\([0-9]*\).*,\1,g') -lt 4 ]]; then
    echo -e "Your yq version not suitable for this script\n please install the latest version from https://github.com/mikefarah/yq"
  fi
  for inventory in $ansiblepath; do
    if [[ -r $inventory ]]; then
      OLDIFS=$IFS
      IFS=$'\n'
      for line in $(
        # shellcheck disable=SC2016
        yq -r '.all.children.* as $company
          | $company.children.*.*
          | to_entries
          | select(.[].value | type=="!!map").[]
          | [{.key: (.value.ansible_user // $company.vars.ansible_user // "root")
           + " " + .value.ansible_host + " " + (.value.ansible_port // $company.vars.ansible_port // 22)
           + " " + (.value.ansible_ssh_common_args // $company.vars.ansible_ssh_common_args // "")}]
          | .[]' "$inventory" | grep -E -v '^#' | awk '{ gsub(/^"|^'\''/, "", $2); gsub(/"$|'\''$/, "", $NF); print $0; }'
        ); do
          sshhostlist["${line%%:*}"]=${line#*:}
      done
      IFS=$OLDIFS
    else
      echo -e "\n! Path $inventory is set in ansiblepath variable, but file is not readable!\n! Either clear this or specify the correct path\n"
    fi
  done
fi

# Netbox inventory implementation
if [[ $netbox_hostname ]]; then
  : "${NETBOX_TOKEN:?Netbox Token must be provided or \"netbox_hostname\" variable should be removed from ${hosts_file} file. You can use \"export NETBOX_TOKEN=...\"}"

  cache_dir="${XDG_CACHE_HOME:-$HOME/.cache}/scu"
  cache_file="$cache_dir/netbox.cache.txt"
  ts_file="$cache_dir/netbox.cache.ts"
  lock_dir="$cache_dir/netbox.cache.lock"
  mkdir -p "$cache_dir"

  cache_load() {
    [[ -s "$cache_file" ]] || return 1
    echo "Loading Netbox cache from $(format_epoch "$(cache_refresh_time)")..."
    while IFS= read -r line
    do
      sshhostlist["${line%%:*}"]=${line#*:}
    done < "$cache_file"
  }

  cache_refresh_time() {
    local last=0
    [[ -r "$ts_file" ]] && read -r last <"$ts_file" || true
    echo -n "$last"
  }

  cache_should_refresh() {
    local now last
    now=$(date +%s)
    last=$(cache_refresh_time)
    [[ $((now - last)) -ge "${netbox_cache_refresh_time}" ]]
  }

  cache_mark_refresh_attempt() {
    date +%s >"$ts_file";
  }

  curl_netbox_request() {
    local type="${1:-}"
    local location=""
    local outfile="${2:?Outfile must be provided during the function call}"

    if [[ "$type" == "virtual_machines" ]]; then
      location="/api/virtualization/virtual-machines/?limit=${netbox_listing_limit}&include=primary_ip"
    elif [[ "$type" == "devices" ]]; then
      location="/api/dcim/devices/?limit=${netbox_listing_limit}&include=primary_ip"
    fi
    # shellcheck disable=SC2086,SC2090
    curl --connect-timeout ${curl_connect_timeout} -fsS -H "Authorization: Token ${NETBOX_TOKEN}" -H 'Accept: application/json' "${netbox_hostname}${location}" \
      | jq -r '.results[]
        | select(.primary_ip4.address != null)
        | "\(.name):\(.config_context.ansible_user // "root" ) \(.primary_ip4.address | split("/")[0]) \(.config_context.ansible_port // 22 ) \(.config_context.ansible_ssh_common_args // "")"' | sed 's/[[:space:]]*$//' >> "$outfile" || return 1
  }

  cache_refresh_job() {
    mkdir "$lock_dir" 2>/dev/null || return 0
    trap 'rmdir "$lock_dir" 2>/dev/null || true' EXIT

    local tmp
    tmp="$(mktemp "$cache_dir/netbox.cache.XXXXXX")"

    if curl_netbox_request virtual_machines "$tmp" && curl_netbox_request devices "$tmp"; then
      if awk '
          {
            if ($1 ~ /^[a-zA-Z0-9._-]+:[a-zA-Z0-9._-]+$/ &&
                $2 ~ /^([0-9]+\.){3}[0-9]+$/ &&
                $3 ~ /^[0-9]+$/ ) {
              ok=1
            }
          }
          END {
            exit(ok ? 0 : 1)
          }
        ' "$tmp"; then
        mv -f "$tmp" "$cache_file"
        return 0
      fi
    fi

    rm -f "$tmp"
    return 1
  }

  cache_refresh_background() {
    cache_mark_refresh_attempt
    (
      exec </dev/null >/dev/null 2>&1
      cache_refresh_job
    ) &
    disown 2>/dev/null || true
  }

  if cache_load; then
    : # sshhostlist loaded
  else
    cache_mark_refresh_attempt
    if cache_refresh_job; then
      cache_load || true
    fi
  fi

  if [[ -s "$cache_file" ]] && cache_should_refresh; then
    cache_refresh_background
  fi
fi

# Indexing hosts list. After that connection by host number will be available.
declare -A array_index;
array_i=1;
for value in "${!sshhostlist[@]}"; do
  array_index[$value]=$array_i;
  array_i=$(( array_i+1 ))
done

# connect to selected host
function sshconnect() {
  local host_to_connect=$1
  for sshhost in "${!sshhostlist[@]}"; do
    if [[ $sshhost = "$host_to_connect" ]] || [[ ${array_index[$sshhost]} = "$host_to_connect" ]]; then
      read -r USERNAME ADDRESS PORT SSHARGS <<<"${sshhostlist[$sshhost]}"
      PORT=${PORT:-22}
      SSHARGS=${SSHARGS//\\\"/\"}
      echo -n "Connecting to $sshhost with command: ";
      bash -x -c "ssh -p $(printf %q "$PORT") $SSHARGS $(printf %q "$USERNAME@$ADDRESS")"
      exit 0;
    fi
  done
}

# Connecting to host if exists
sshconnect "$host_connect"

#
# Usage information and hosts list
# which will be displayed if no valid host was choosen
#
echo -e "\t\tUsage: $0 <host number|host name>\n";
echo "Available hosts:";
counter=1;
counter2=1;
for value in "${!sshhostlist[@]}"; do
  name="$value";
  value="${array_index[$value]}: $value";
  if [[ ${#value} -gt 7 ]]; then
    if [[ ${#value} -gt 15 ]]; then
      RES="$value\t";
    else
      RES="$value\t\t\t";
    fi
  else
    RES="$value\t\t\t\t";
  fi
  if [[ -n $host_connect ]]; then
    [[ $RES =~ $host_connect ]] || continue;
    MATCHVALUE="$name";
    echo -ne "$RES";
  else
    echo -ne "$RES";
  fi
  counter=$(( counter+1 ))
  counter2=$(( counter2+1 ))
  if [[ $counter -gt 3 ]]; then
    echo "";
    counter=1;
  fi
done
echo "";
if [ $counter2 -eq 2 ] && [ -n "$MATCHVALUE" ]; then
  echo "Connecting to only one matched host $MATCHVALUE...";
  sshconnect "$MATCHVALUE";
  exit 0;
fi
if [ -n "$host_connect" ] && [ $counter2 -eq 1 ]; then
  echo "There is no host to connect with name you wrote...";
fi
echo "";
if [ -n "$host_connect" ] && [ $counter2 -gt 1 ]; then
  echo -n "Enter host number / host name to connect: "
  read -r input
  sshconnect "$input"
  exit 0;
fi