#!/usr/bin/env bash

# Generates temporary AWS credentials using your MFA device (12 hour duration).
# Usage: ./aws-session-login [profile_name]
#  - or to be prompted to enter the profile -
# ./aws-session-login

txtred=$(tput setaf 1)
txtgreen=$(tput setaf 2)
txtblue=$(tput setaf 4)
txtpurple=$(tput setaf 5)
txtreset=$(tput sgr0)

# Reset staged credentials
unset AWS_ACCESS_KEY_ID
unset AWS_SECRET_ACCESS_KEY
unset AWS_SESSION_TOKEN
unset AWS_DEFAULT_REGION

# Load the known aws profiles
if [ ! -f "${HOME}"/.aws/credentials ]; then
    echo "There is no credentials file located in ${HOME}/.aws/credentials"
    exit 1
fi

# Older version of this tool moved your crentials into a temp file ~/.aws/.store
# This version of the script won't overwrite your local files, and will attempt
# to correct the old behavior
if [ -f "${HOME}"/.aws/.store ]; then
  echo "${txtred}A previous version of this script moved your credentials into a temporary file at ~/.aws/.store${txtreset}"
  read -r -p "Permantently move ~/.aws/.store back to ~/.aws/credentials? (Y/n): " overwrite_dot_store
  if [ "$overwrite_dot_store" == 'Y' ]; then
    mv "${HOME}"/.aws/credentials "${HOME}"/.aws/credentials.bak
    mv "${HOME}"/.aws/.store "${HOME}"/.aws/credentials
    echo "Backed up old credentials file to ~/.aws/credentials.bak"
    echo "Moved ~/.aws/.store to ~/.aws/credentials"
  else
    echo "WARNING: ~/.aws/.store file is present. This may mean that the contents of your credentials file are corrupted."
  fi
fi

TMP_PROFILE_LIST=( $(grep "^\[[a-z]*\]" "$HOME/.aws/credentials" | sed 's/^\[//g' | sed 's/\]$//g' | uniq | tr '\n' ' ') )

# Check if the profile name was passed as an argument
if [ -n "${1}" ]; then
  FOUND_IN_LIST=0
  for i in "${TMP_PROFILE_LIST[@]}"; do
    if [ "${i}" = "${1}" ]; then
      FOUND_IN_LIST=true
    fi
  done
  if [ "${FOUND_IN_LIST}" != true ]; then
    echo "AWS profile '${1}' not found."
    exit 1
  fi
  TMP_AWS_PROFILE="${1}"
fi

if [ -z "${TMP_AWS_PROFILE}" ]; then
  # Get the AWS profile from user input
  echo "Select an AWS profile:"
  COUNTER=0
  for i in "${TMP_PROFILE_LIST[@]}"; do
      ((COUNTER=COUNTER+1))
      echo -e "$COUNTER) ${txtgreen}$i${txtreset}"
  done
  echo "-------------"
  read -r -p "Enter profile number: " profile_index
  ((profile_index=profile_index-1))
  if [ "$profile_index" -lt 0 ] || [ -z "${TMP_PROFILE_LIST[$profile_index]}" ]; then
      echo "Selected profile does not exist"
      exit 1
  fi
  TMP_AWS_PROFILE="${TMP_PROFILE_LIST[$profile_index]}"
fi
export AWS_PROFILE="${TMP_AWS_PROFILE}"

# Query IAM to get the username associated with the keys in the credentials file
username=$(aws iam get-user --query 'User.[UserName]' --output text)
if [ -z "${username}" ]
then
  echo "Cannot get profile name."
  exit 1
fi
echo "Profile found for:    ${txtgreen}${username}${txtreset}"

# Query IAM to get the user's MFA device
device=$(aws iam list-mfa-devices --user-name "${username}" --query 'MFADevices[*].SerialNumber' --output text)
if [ -z "${device}" ]; then
  echo "Can not find any MFA device for you."
  exit 1
fi
echo "Your MFA device is:  " "${device}"
read -s -r -p "MFA code: " code
echo ''

# Use the MFA device to get a session token and query out the necessary credentials.
output=$(
  aws sts get-session-token \
    --serial-number "${device}" \
    --query 'Credentials.[SecretAccessKey,AccessKeyId,SessionToken,Expiration]' \
    --output text \
    --token-code "${code}" | tr '\t' ' ')
secret=$(echo "$output" | cut -f1 -d ' ')
access=$(echo "$output" | cut -f2 -d ' ')
session=$(echo "$output" | cut -f3 -d ' ')
expiration=$(echo "$output" | cut -f4 -d ' ')

if [ -z "${secret}" ] || [ -z "${session}" ] || [ -z "${access}" ]; then
  echo "Unable to get temporary credentials."
  echo "Output: ${output}"
  exit 1
fi

SESSION_CREDENTIALS="export AWS_ACCESS_KEY_ID=${access} AWS_SECRET_ACCESS_KEY=${secret} AWS_SESSION_TOKEN=${session} AWS_DEFAULT_REGION=us-east-1"
echo -e "\n${txtgreen}Login Sucessful!${txtreset}\n"
echo -e "Copy and paste the following into your terminal:\n"
echo -e "${txtpurple}${SESSION_CREDENTIALS}${txtreset}\n"

SESSION_FILE="${HOME}/.aws/${TMP_AWS_PROFILE}.session.env"
echo "# expires ${expiration}" > "$SESSION_FILE"
echo "${SESSION_CREDENTIALS}" >> "$SESSION_FILE"
echo "Credentials saved in file ${txtblue}${SESSION_FILE}${txtreset}"

echo "To load credentials in a new terminal, run the command: \"${txtblue}. ${SESSION_FILE}${txtreset}\""
echo "Created at: $(date -u +%FT%T%Z). ${txtred}Valid until: ${expiration}${txtreset}"