qlaws

#!/bin/bash

SCRIPT_NAME=$(basename "$0")
DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"

# Initialize variables
iam="root-iam"
export_flag=false

# Function to display help
show_help() {
  echo "Usage: $0 [-e iam] [-x] [-h]"
  echo "  -a role               Assume iam role"
  echo "  -x                    Export only (no new session)"
  echo "  -h                    Display help"
  echo
  echo "After logging in:"
  echo "  loadauth              When in export only mode, loads the credentials into your session."
  echo "  assume [role iam]     Assumes a iam role and updates credentials so you can connects."
  echo "  leave                 Goes back to parent role (root)."
  echo "  connect [hostname]    Connects to a server (assuming you're in the right account)."
  echo "  deploy [hostname]     To do: will trigger an app deployment to the specified server."
}

# Parse arguments
while getopts ":a:xh" opt; do
  case ${opt} in
    a )
      iam=$OPTARG 
      ;;
    x )
      export_flag=true
      ;;
    h )
      show_help
      exit 0
      ;;
    \? )
      echo "Invalid option: -$OPTARG" >&2
      show_help
      exit 1
      ;;
    : )
      echo "Option -$OPTARG requires an argument." >&2
      show_help
      exit 1
      ;;
  esac
done

# Files
setup_files() {
  SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
  AWS_CONFIG="$HOME/.aws/config"
  AWS_KEYCHAIN="$HOME/Library/Keychains/aws-vault.keychain-db"
  SHELL_PROFILE="$HOME/.zshrc"
  PARENT_CREDS="$SCRIPT_DIR/.$SCRIPT_NAME.root"
  ENVIRONMENT_FILE="$SCRIPT_DIR/.$SCRIPT_NAME.env"
  CONFIG_FILE="$SCRIPT_DIR/.$SCRIPT_NAME.cfg"
  KEYCHAIN_P="$SCRIPT_DIR/.secret.keychain"
  TOKEN_2FA="$SCRIPT_DIR/.secret.token"

  echo "# $SCRIPT_NAME ENVIRONMENT" > $ENVIRONMENT_FILE
  export_to_env "AWS_ACCOUNT_PROMPT=$iam"
  chmod a+x $ENVIRONMENT_FILE

  if ! [[ -f "$CONFIG_FILE" ]]; then
    echo "# $SCRIPT_NAME CONFIG" > $CONFIG_FILE
    echo 'export PATH='$SCRIPT_DIR':$PATH' >> $CONFIG_FILE
    echo "export PARENT_CREDS=$PARENT_CREDS" >> $CONFIG_FILE
    echo "export AWS_CONFIG=$AWS_CONFIG" >> $CONFIG_FILE
    echo "export ENVIRONMENT_FILE=$ENVIRONMENT_FILE" >> $CONFIG_FILE
    echo "export CONFIG_FILE=$CONFIG_FILE" >> $CONFIG_FILE
    add_prompt_to_config
    add_methods_to_config
    chmod a+x $CONFIG_FILE
  fi
  source $CONFIG_FILE
}

# Colors
lcl="\033[38;2;85;172;243m" # Blue
gcl="\033[1;32m" # Green
rcl="\033[1;31m" # Red
ycl="\033[0;33m" # Yellow
wcl="\033[35;96m" # White
ncl="\033[0m" # No color

# Gfx side lines and stuff
sls="${lcl}┏╾┈" # Side line start
slm="${lcl}┃ " # Side line mid
sle="${lcl}┗╾┈${ncl} " # Side line end
rsls="${lcl}┈╼┓"
rsle="${lcl}┈╼┛"
spinner=( ∙∙∙ ●∙∙ ∙●∙ ∙∙● )
spc="${gcl}${spinner[0]}" # Spinner characters
rws="\033[${#spinner[0]}D" # Rewind spinner characters
rwc="\033[1D" # Rewind cursor 1 character
mcu="\033[1A"
bul="${lcl}▬${ncl} " # Bullet
chk="❤" # Checkmark

# Print colored symbol and play a sound
print_symbol() {
  printf "$2$1${ncl}%2s"
  afplay "$DIR/sounds/step.mp3" &
}

# Play the decoding sound
play_decoding_sound() {
  while :; do afplay "$DIR/sounds/interface.mp3"; done
}

# Spinner animation
start_spinner() {
  trap "clean_up" INT # Make sure spinner is killed if script is interrupted
  stty -echo # Disable keyboard
  printf " $spc"
  spinner_process &
  spinner_pid=$!
  disown
}

# Stop spinner animation
stop_spinner() {
  kill $spinner_pid >/dev/null 2>&1
  trap - INT
  printf "${rws}" # rewinds cursor position
  print_symbol "$chk" "$gcl"
  stty echo # Reenable keyboard
}

# Spinner animation process (detached)
spinner_process() {
  if [[ "$1" == "" ]]; then
    while true; do for i in "${spinner[@]}"; do printf "${rws}$i"; sleep 0.05; done; done
  else
    for i in "${spinner[@]}"; do printf "${rws}$i"; sleep 0.05; done
  fi
}

# Spinner clean up
clean_up() {
  kill $spinner_pid >/dev/null 2>&1
  kill $sound_pid >/dev/null 2>&1
  stty echo
  exit
}

shuffle_and_scroll() {
  trap "clean_up" INT # Make sure stuff is killed if script is interrupted
  play_decoding_sound &
  sound_pid=$!
  disown
  local text="$1" length="$2" chars="$3"
  for (( i=0; i<length && i<${#text}; i++ )); do
    printf "%s" "${chars:RANDOM%${#chars}:1}"
    sleep 0.005
    printf "\b%s" "${text:$i:1}"
  done
  kill $sound_pid >/dev/null 2>&1
  for (( i=1; i<=${#text} - length; i++ )); do
    printf "\b%.0s" $(seq 1 $length)
    printf "%.*s" "$length" "${text:i}"
    afplay "$DIR/sounds/tick.mp3" &
    sleep 0.2
  done
}

# Deshuffle animation
fake_deshuffle() {
  trap "clean_up" INT # Make sure stuff is killed if script is interrupted
  play_decoding_sound &
  sound_pid=$!
  disown
  word_len="${#1}"
  for ((k=0; k<$word_len; k++)); do
    len_left=$((word_len-k))
    for ((i=0; i<3; i++)); do
      for ((j=0; j<$len_left; j++)); do
        printf "${2:$(( RANDOM % ${#2} )):1}"
      done
      sleep $3
      printf "\033[${len_left}D"
    done
    printf "${1:$k:1}"
  done
  kill $sound_pid >/dev/null 2>&1
}

# Add prompt configuration to config file
add_prompt_to_config() {
  echo 'if [ ! -z "$AWS_VAULT" ]; then\
  ZSH_THEME_GIT_PROMPT_PREFIX="%{$fg[yellow]%}("
  ZSH_THEME_GIT_PROMPT_CLEAN="%{$fg[green]%}✓%{$reset_color%}"
  ZSH_THEME_GIT_PROMPT_DIRTY="%{$fg[red]%}✗%{$reset_color%}"
  ZSH_THEME_GIT_PROMPT_SUFFIX="%{$fg[yellow]%})%{$reset_color%}"
  PROMPT='"'"'%{$fg[$NCOLOR]%}${PERSONAL_IAM_PROFILE}%{$reset_color%}@%{$fg[cyan]%}${AWS_ACCOUNT_PROMPT}\
%{$reset_color%}:%{$fg[magenta]%}%c\
$(git_prompt_info) \
%{$fg[red]%}%(!.#.»)%{$reset_color%} '"'"'
  fi' >> $CONFIG_FILE
}

# Add additional methods to config file
add_methods_to_config() {
  echo '
  # Export to env file
  export_to_env() {
    echo "export $1" >> $ENVIRONMENT_FILE
  }
  # Get role arn from .aws/config
  get_role_arn() {
    echo $(awk -v profile="profile $1" '"'"'
    BEGIN {found = 0}
    /^\[/{found = 0}
    $0 ~ profile {found = 1}
    found && /role_arn/ {print $3; exit}
  '"'"' "$AWS_CONFIG")
  }
  # EC2 id translator - this will look for the settings file for the specified server, i.e. my.server.com and then return the ec2 id, i.e. i-1234567890
  # You should write this method according to your server provisioning structure
  find_ec2_by_host() {
    host_file=$(find "$HOME/Dovetail/server-provisioning/hosts/" -type f -exec grep -l -w " $1" {} +)
    echo $(cat $host_file | grep "ansible_host:" | awk -F": " '"'"'{print $2}'"'"') 
  }
  # Load credentials to current session
  loadauth() {
    source $ENVIRONMENT_FILE
    source $CONFIG_FILE
    printf "Loaded credentials for $AWS_ACCOUNT_PROMPT into environment\n"
  }
  # EC2 connect
  connect() {
    ssh -o ProxyCommand="sh -c '"'"'aws ssm start-session --target %h --document-name AWS-StartSSHSession --parameters portNumber=%p'"'"'" ubuntu@$(find_ec2_by_host "$1")
  }
  # Assume a role
  assume() {
    echo "$(env | grep ^AWS_)" > $PARENT_CREDS
    role_arn="$(get_role_arn $1)"
    ASSUMED_CREDS=$(aws sts assume-role --role-arn "$role_arn" --role-session-name "ubuntu")
    export_to_env "AWS_ACCESS_KEY_ID=\"$(echo "$ASSUMED_CREDS" | jq '.Credentials.AccessKeyId')\""
    export_to_env "AWS_SESSION_TOKEN=\"$(echo "$ASSUMED_CREDS" | jq '.Credentials.SessionToken')\""
    export_to_env "AWS_SECURITY_TOKEN=\"$(echo "$ASSUMED_CREDS" | jq '.Credentials.SessionToken')\""
    export_to_env "AWS_SECRET_ACCESS_KEY=\"$(echo "$ASSUMED_CREDS" | jq '.Credentials.SecretAccessKey')\""
    export_to_env "AWS_SESSION_EXPIRATION=\"$(echo "$ASSUMED_CREDS" | jq '.Credentials.Expiration')\""
    export_to_env "AWS_ACCOUNT_PROMPT=$1"
    source $ENVIRONMENT_FILE
    printf "Loaded credentials for $AWS_ACCOUNT_PROMPT into environment\n"
  }
  # Leave role
  leave() {
    if [[ "$AWS_ACCOUNT_PROMPT" == "root-iam" ]]; then
      unset AWS_VAULT
      printf "Logged out aws-vault\n"
    else
      cat "$PARENT_CREDS" | while read line ; do export_to_env "$line"; done
      export_to_env "AWS_ACCOUNT_PROMPT=root-iam"
      source $ENVIRONMENT_FILE
      printf "Left assume role back into $AWS_ACCOUNT_PROMPT\n"
    fi
    source $CONFIG_FILE
  } ' >> $CONFIG_FILE
}

# Config
run_config() {
  if [[ -z "$PERSONAL_IAM_PROFILE" ]]; then
    PERSONAL_IAM_PROFILE=$(awk -F'[\\[\\] ]+' '/^\[profile / {print $3; exit}' "$AWS_CONFIG")
    if [[ -n "$PERSONAL_IAM_PROFILE" ]]; then
      printf "Found aws profile: $PERSONAL_IAM_PROFILE \n"
      printf "Is is correct? (y/n): "
      read -n 1 profile_correct
      printf "\n"
    fi
    
    if [[ -z "$PERSONAL_IAM_PROFILE" || "$profile_correct" != "y" ]]; then
      printf "Enter the IAM profile you want to use: "
      read PERSONAL_IAM_PROFILE
    fi
    echo "export PERSONAL_IAM_PROFILE=$PERSONAL_IAM_PROFILE" >> $CONFIG_FILE
  fi

  if ! [[ -f "$KEYCHAIN_P" ]]; then
    printf "I can unlock you aws-vault keychain for you as well as generate 2FA tokens\n"
    if [[ -f "$AWS_KEYCHAIN" ]]; then
      printf "I found your aws-vault keychain here: $AWS_KEYCHAIN\n"
      printf "Enter the keychain password: "
      read -s keychain_pwd
      printf "\n"
      enc_dec "$keychain_pwd" "$KEYCHAIN_P"
    else
      printf "Unable to find your aws-vault keychain in the usual place.\n"
      printf "Do you want to specify he path? (y/n): "
      read -n 1 specify_path
      printf "\n"
    fi

    if [[ "$specify_path" == "y" ]]; then
      printf "Enter the path to your keychain file: "
      read AWS_KEYCHAIN
      echo "export KEYCHAIN_FILE_PATH=$AWS_KEYCHAIN" >> $CONFIG_FILE
    fi
  fi

  if ! [[ -f "$TOKEN_2FA" ]]; then
    printf "Do you want me to automate 2FA? "
    read -n 1 do_2fa
    printf "\n"

    if [[ "$do_2fa" == "y" ]]; then
      printf "Enter your totp secret key (find it in your current authenticator app): "
      read -s totp_secret
      printf "\n"
      enc_dec "$totp_secret" "$TOKEN_2FA"
    fi
  fi
}

# Source the config file in the user's shell profile (if it's not there yet)
add_config_to_shell_profile() {
  line="source $CONFIG_FILE > /dev/null 2>&1"
  if ! grep -Fxq "$line" "$SHELL_PROFILE"; then
    echo "$line" >> "$SHELL_PROFILE"
  fi
}

# Print box
print_box_start() {
  printf "${sls}%37s${rsls}"
}

print_box_mid() {
  printf "\n${slm}%40s${slm}"
  printf "%.0s${rwc}" {1..42}
}

print_box_end() {
  printf "\n${sle}%36s${rsle}\n"
}

# Show commands
show_commands() {
  print_box_mid;printf "${bul}Available commands${gcl}:"
  print_box_mid;printf "  ${gcl}▪${lcl} connect [host] ${gcl}▪${lcl} assume [role-iam]"
  print_box_mid;printf "  ${gcl}▪${lcl} deploy [host]  ${gcl}▪${lcl} leave (to root)"
}

# Logo string
logo_string="-
░▄▀▀▄░█░░░▒░▄▀▀▄░█░░▐░▄▀▀▀░▒░▀▀▀▀▀░▀░◣░
░█░░█░█░░░▒░█▀▀█░█ █▐░▀▀▀█░▒░▀▀▀▀░▀░◣░░
░░▀▀▄░▀▀▀░▒░▀░░▀░▀▀░▀░▀▀▀▀░▒░▀▀▀░▀░◣░░░
▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀"

# Logo animation
print_logo() {
  print_box_start
  IFS=$'\n' read -rd '' -a lines <<< "$logo_string"
  for ((l=1; l<${#lines[@]}; l++)); do
    color="$(echo ${!l})"
    print_box_mid
    printf "$color"
    fake_deshuffle "${lines[l]}" "▁▂▃▄▅▆▇█▉▊▋▌▍▎▏▐░▒▓▕▖▗▘▙▚▛▜▝▞▟" 0.0002
  done
}

# Outro
print_outro() {
  print_box_mid
  printf "${bul}"
  fake_deshuffle "Quick Login · Amazon Web Services " "абвгдеёжзийклмнопрстуфхцчшщъыьэюя" 0.0005; print_symbol "$chk" "$gcl"
}

# Finish action
finish() {
  print_box_end
  afplay "$SCRIPT_DIR/sounds/success.mp3" &
}

# Encrypt and decrypt
enc_dec() {
  auth_openssl="-pass pass:$SCRIPT_NAME"
  if [[ "$2" == "" ]]; then
    openssl enc -d -aes-256-cbc -iter 1000 $auth_openssl -in "$1"
  else
    echo "$1" | openssl enc -aes-256-cbc -iter 1000 -salt $auth_openssl -out "$2"
  fi
}

# 2FA totp generator
totp(){
  local key="${1}"
  local timestamp="${2:-$(date +%s)}"
  local key_hex=$(echo -n "${key}" | base32 -d | xxd -p | xargs | tr -d ' ')
  local counter=$(printf %016x "$((timestamp / 30))")
  local hmac=$(echo -n "${counter}" | xxd -r -p | openssl dgst -sha1 -mac hmac -macopt hexkey:"${key_hex}" | awk '{print $2}')
  local offset=$((0x${hmac:39:1}))
  local truncated=${hmac:2*offset:8}
  local otp=$((0x${truncated} & 0x7fffffff))
  printf "%06d\n" $((otp % 1000000))
}

# Prepare credentials
prepare_credentials() {
  print_box_mid;printf "${bul}Decrypting credentials"
  start_spinner
  keychain_pwd="$(enc_dec $KEYCHAIN_P)"
  secret="$(enc_dec $TOKEN_2FA)"
  stop_spinner
  print_box_mid;printf "${bul}Generating 2FA token "
  printf "${spc}"
  spinner_process 2
  printf "${rws}▸ "
  token="$(totp ${secret})"
  fake_deshuffle "$token" "1234567890" 0.02
}

# Unlock aws-vault keychain
unlock_keychain() {
  security -q -v unlock-keychain -p "$keychain_pwd" $AWS_KEYCHAIN > /dev/null 2>&1  
}

# Check if already in aws vault session
check_existing_aws_session() {
  if ! [ -z ${AWS_VAULT+x} ]; then
    printf "Already in a session on $AWS_VAULT.\n"
    printf "Exit or ${gcl}leave${ncl} this session before running $SCRIPT_NAME again.\n"
    exit 1
  fi
}

# Execute login to aws
login_to_aws() {
  print_box_mid;printf "${bul}"
  shuffle_and_scroll "Logging in to $iam" 34 "абвгдеёжзийклмнопрстуфхцчшщъыьэюя"
  start_spinner
  aws_env=$(aws-vault exec "$PERSONAL_IAM_PROFILE" -t "$token" -- env) > /dev/null 2>&1
  echo "$aws_env" | grep "^AWS_" | while read line ; do export_to_env "$line"; done
  stop_spinner
  source $ENVIRONMENT_FILE
  source $CONFIG_FILE

  if [[ "$iam" != "root-iam" ]]; then
    print_box_mid;printf "${bul}Assuming role"
    start_spinner
    assume "$iam" > /dev/null 2>&1
    stop_spinner
  fi

  if ! grep -q "AWS_VAULT" "$ENVIRONMENT_FILE"; then
    print_box_mid;printf "\r${slm}${bul}Login failed "
    print_symbol "✗" "${rcl}"
    finish
  elif [[ "$export_flag" == false ]]; then
    print_box_mid;printf "\r${slm}${bul}New AWS session started $spc$gcl"
    spinner_process 2
    stop_spinner
    show_commands
    finish
    bash -c ". $ENVIRONMENT_FILE; exec zsh -i" #this is to start a new session in terminal
    rm $ENVIRONMENT_FILE &> /dev/null
    exit
  else
    print_box_mid;printf "\r${slm}${bul}Use ${gcl}loadauth${ncl} to load credentials "
    print_symbol "✦" "$ycl"
    show_commands
    finish
  fi  
}

# Export to env file
export_to_env() {
  echo "export $1" >> $ENVIRONMENT_FILE
}

# Main script logic

setup_files

run_config

add_config_to_shell_profile

check_existing_aws_session

print_logo "$wcl" "" "" "$rcl" # Colors per line

print_outro

prepare_credentials

unlock_keychain

login_to_aws