#!/usr/bin/env bash
#
# Keymanage provides functionality to push, remove, and swap/refresh ssh keys
# on a list of systems for the specified users
#
# Copyright (C) 2014  Aaron Ball <nullspoon@iohq.net>
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
# 
manifest=''
key=''
action=''
id=''
user=''


#
# Backups by a specific ssh key to <date_modified>.<key_name>
#
# @param ssh_base string Path to where the ssh keys and configs are stored
# @param key string Name of the key to backup
#
# @return string The filename of the key backup
#
function backup_keys {
  local ssh_base=$1
  local key=$2
  moved=0;
  date=""
  priv_ls=$(ls -l --time-style='+%Y%m%d%H%M%S' ${ssh_base}/${key})
  date=$(echo "${priv_ls}" | tr -s ' ' | cut -d ' ' -f 6);
  # Rename the old key
  if [[ -e "${ssh_base}/${key}" ]]; then
    mv ${ssh_base}/${key} ${ssh_base}/${date}.${key}
    moved=1;
  fi
  # Rename the old public key
  if [[ -e "${ssh_base}/${key}.pub" ]]; then
    mv ${ssh_base}/${key}.pub ${ssh_base}/${date}.${key}.pub
    moved=1;
  fi
  if [[ ${moved} == 0 ]]; then
    echo ''
  else
    chmod 700 ${ssh_base}/*
    echo "${ssh_base}/${date}.${key}"
  fi
}


#
# Pushes specific public key to remote user's authorized_keys
#
# @param user string User owning the authorized_keys file to be modified
# @param server string Server the user's authorized_keys file is on
# @param old_key string The key to use for authentication
# @param new_key string The key, public or private, to be pushed
#
function push_key {
  local conn=$1
  local old_key=$2
  local new_key=$3
  if [[ ${#new_key} -lt '4' ]]; then
    echo "Key to be pushed is not a public key."
    exit
  fi

  ispub=$(key_is_public ${new_key})
  if [[ ${ispub} == 0 ]]; then
    # Append .pub because a public key wasn't specified
    new_key="${new_key}.pub"
  fi

  local cmd="if [[ ! -d ~/.ssh/ ]]; then mkdir ~/.ssh/; fi"
  cmd="${cmd} && echo '$(cat ${new_key})' >> ~/.ssh/authorized_keys"

  # Unset our identity file if it doesn't exist
  local id_file="-i ${old_key}"
  if [[ ${old_key} == '' ]]; then
    id_file=''
  fi
  contents=$(cat ${new_key})
  ssh -q ${id_file} ${conn} "${cmd}"
}


#
# Removes the specified public key from a remote user's authorized_keys file
#
# @param user string User owning the authorized_keys file to be modified
# @param server string Server the user's authorized_keys file is on
# @param key string The key to use for authentication which is to be removed
#
function remove_remote_key {
  local conn=$1
  local key=$2
  pub_key=''
  priv_key=''
  ispub=$(key_is_public ${key})
  if [[ ${ispub} == 0 ]]; then
    priv_key="${key}"
    pub_key="${key}.pub"
  else
    priv_key="${key:0:-4}"
    pub_key="${key}"
  fi
  contents=$(cat "${pub_key}")
  local cmd="if [[ ! -d ~/.ssh/ ]]; then mkdir ~/.ssh/; fi"
  cmd="${cmd} && cat ~/.ssh/authorized_keys | grep -v '${contents}' "
  cmd="${cmd} > ~/.ssh/auth_keys"
  cmd="${cmd} && mv ~/.ssh/auth_keys ~/.ssh/authorized_keys"
  if [[ ${id} != '' ]]; then
    ssh -q -i ${id} ${conn} "${cmd}"
  else
    ssh -q ${conn} "${cmd}"
  fi
}


#
# Determines if the specified key is public (or not which would be private).
#
# @param key string Path to the key to check
#
# @return int Whether or not the key is public
#
function key_is_public {
  key=$1
  if [[ ${#key} -lt '4' ]]; then
    echo 0;
  fi
  # Check the extension
  ext=${key:$((${#key}-4)):${#key}}
  if [[ ${ext} == '.pub' ]]; then
    echo 1;
  fi
  echo 0
}


#
# Generates a new ssh key of the length 4096
#
# @param filepath string Path to where the new ssh key will be written
# @param bits int Number of bits in the new key (eg: 2048, 4096, 8192, etc.)
#
function gen_key {
  local filepath=$1
  local bits=$2
  ssh-keygen -b ${bits} -f "${filepath}" -N ''
}


#
# Prints the help text
#
function get_help {
	echo "
Manages ssh keys en masse. Designed to perform pushes, removals, and swaps
of ssh keys on a llist of servers for the current or specified user.

Usage: keymanage.sh action --manifest systems.list

Actions: push, update, remove

Arguments:
 -m,--manifest   Text file containing a list of systems, delimited by new
                 lines.
 -k,--key        Path to a key to perform an action (push or remove) with.
 -i,--id         Key to use for automated logins. Not used when performing an
                 update.
 -u,--user       Username on remote systems to work on (defaults to root).
 
"
}


function parse_args {
  argv=(${@})
  # Parse the arguments
  for(( i=0; i<${#argv[*]}; i++ )); do
    if [[ ${argv[$i]} == '-h' || ${argv[$i]} == '--help' ]]; then
      echo "$(get_help)"
      exit
    elif [[ ${argv[$i]} == '-m' || ${argv[$i]} == '--manifest' ]]; then
      manifest=${argv[$i+1]}
      i=$(( ${i} + 1 ))
    elif [[ ${argv[$i]} == '-k' || ${argv[$i]} == '--key' ]]; then
      key=${argv[$i+1]}
      i=$(( ${i} + 1 ))
    elif [[ ${argv[$i]} == '-i' || ${argv[$i]} == '--id' ]]; then
      id=${argv[$i+1]}
      i=$(( ${i} + 1 ))
    elif [[ ${argv[$i]} == '-u' || ${argv[$i]} == '--user' ]]; then
      user=${argv[$i+1]}
      i=$(( ${i} + 1 ))
    else
      action=${argv[$i]}
    fi
  done

  # Enforce some business rules
  echo
  exit=0;
  if [[ ${action} == '' ]]; then
    echo "Please specify an action.";
    echo "  Available actions: push, remove, update."
    echo
    exit=1;
  fi
  if [[ ${manifest} == '' ]]; then
    echo "Please specify a manifest file."
    echo "  Example: keymanage.sh action [-m|--manifest] ./systems.txt"
    echo
    exit=1;
  fi
  if [[ ${exit} == 1 ]]; then
    exit
  fi
}


#
# Determines the path to the parent directory containing a file.
#
# @param filepath string Path to the file to get the parent directory for
#
# @return string Path to the file's parent directory
#
function get_file_path {
  filepath=$1
  filename=$(basename ${filepath})
  echo ${filepath} | sed "s/\(.*\)${filename}/\1/"
}


#
# Push main function. One param because the rest are global
#
function key_push {
  argv=( ${@} )
  if [[ ${id} == '' ]]; then
    echo "No identity file specified (-i). This will likely be painful."
  fi
  for (( i=0; i<${#argv[*]}; i++ )); do
    dest=${argv[$i]}
    if [[ ${id} == '' ]]; then
      push_key "${dest}" '' ${key} 
    else
      push_key "${dest}" ${id} ${key} 
    fi
    echo "Key ${key} added for ${dest}."
  done
}


#
# Update main function. One param because the rest are global
#
function key_update {
  argv=( ${@} )
  ssh_base=$(get_file_path ${key})
  filename=$(basename ${key})
  # Backup our old key
  backup_key="$(backup_keys ${ssh_base} ${filename})"

  # Let's get to work on that new key
  gen_key "${key}" 4096

  for (( i=0; i<${#argv[*]}; i++ )); do
    dest=${argv[$i]}
    if [[ ${backup_key} == '' ]]; then
      echo "No current key exists."
      echo "Skipping backup and removal from remote."
      # Push the new key
      push_key "${dest}" '' ${key} 
    else
      # Push the new key
      push_key "${dest}" ${backup_key} ${key} 
      # Clean up the old key from our remote
      remove_remote_key "${dest}" "${backup_key}"
    fi
    echo "Key ${key} updated for ${dest}."
  done
}


#
# Remove main function. One param because the rest are global
#
function key_remove {
  argv=( ${@} )
  for (( i=0; i<${#argv[*]}; i++ )); do
    dest=${argv[$i]}
    remove_remote_key "${dest}" "${key}"
    echo "Key ${key} removed from ${dest}."
  done
}


#
# The main function
#
function main {
  # Parse our script args
  # Believe me, this is a lot better than the alternatives
  parse_args ${@}

  destinations=( $(cat ${manifest}) )
  # Key required
  if [[ ${key} == '' ]]; then
    echo -n "Please specify a key (-k) to ${action}."
    echo
    exit
  fi

  # Let's start doing stuff
  if [[ ${action} == 'push' ]]; then
    key_push ${destinations[*]}
  elif [[ ${action} == 'update' ]]; then
    key_update ${destinations[*]}
  elif [[ ${action} == 'remove' ]]; then
    key_remove ${destinations[*]}
  fi
}

main ${@}