assumerole

#! /bin/bash

SetInitialEnvironment() {
  unset AWS_ACCESS_KEY_ID
  unset AWS_SECRET_ACCESS_KEY
  unset AWS_SESSION_TOKEN
  unset AWS_PROFILE
  unset ASSUMEROLE_ENV
  unset SSHKEY
  unset OTP

  ASSUMEROLE_CODEARTIFACT=0
  [[ -z "${ASSUMEROLE_SKIP_CODEARTIFACT}" && -e ~/.m2/settings.xml ]] && ASSUMEROLE_CODEARTIFACT=1
}

SelectProfile() {
  local saved_ps3
  saved_ps3="${PS3}"
  PS3="Enter a number: "

  plain "Select from these available accounts:"
  plain ""

  select choice in ${@}; do
    if [[ ${REPLY} -gt 0 && ${REPLY} -le ${#} ]]; then
      aws_account=${choice}
      break
    fi
  done
  PS3="${saved_ps3}"
}

SaveCredentials() {
  [[ -d ~/.assumerole.d/cache ]] || mkdir -p ~/.assumerole.d/cache

  {
    echo "export AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}"
    echo "export AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}"
    echo "export AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN}"
    echo "export ROLE=${ROLE}"
    echo "export ACCOUNT=${ACCOUNT}"
    echo "export AWS_ACCOUNT_ID=${ACCOUNT}"
    echo "export aws_account=${aws_account}"
    echo "export AWS_ACCOUNT=${aws_account}"
    echo "export AWS_EXPIRATION=${AWS_EXPIRATION}"
    echo "export SSHKEY=${SSHKEY}"
    echo "export PROFILE=${PROFILE}"
    echo "export AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION}"
    echo "${ASSUMEROLE_ENV}"
  } > "${HOME}/.assumerole.d/cache/${aws_account}"

  chmod 0600 "${HOME}/.assumerole.d/cache/${aws_account}"
}

CheckAndLoadFromCache() {
  local cache_file
  cache_file=~/.assumerole.d/cache/${aws_account}
  if [[ -e ${cache_file} ]]
  then
    ### Cache exists, load it and check if it is still valid
    # shellcheck source=./test/cache-example
    . "${cache_file}"
    if aws sts get-caller-identity >/dev/null 2>&1
    then
      IsVerbose && info "Credentials for ${aws_account} loaded from cache."
      return 0
    else
      IsVerbose && info "Cache found for ${aws_account}, but credentials have expired and will be deleted."
      rm -f "${cache_file}"
      return 1
    fi
  else
    return 1
  fi
}

GetOtpInfo() {
  # Check Env Vars
  unset IS_SET
  if [[ -z ${ONEPASS_ACCOUNT} ]] 
  then
    info "Var ONEPASS_ACCOUNT is not set, reverting to manual MFA token prompt"
    IS_SET=false
  fi
  if [[ -z ${ONEPASS_ENTRY} ]]
    then info "Var ONEPASS_ENTRY is not set, reverting to manual MFA token prompt"
    IS_SET=false
  fi

  # Get OTP Info
  if [[ -z ${IS_SET} ]]
  then
    eval $(op signin --account ${ONEPASS_ACCOUNT})
    OTP=$(op item get "${ONEPASS_ENTRY}" --otp)
    op signout
  fi
}

GetAccountInfo() {
               PROFILE=$(jq    --raw-output ".assume_roles[\"${aws_account}\"][\"aws_profile\"]" "${CONF}")
               ACCOUNT=$(jq    --raw-output ".assume_roles[\"${aws_account}\"][\"aws_account\"]" "${CONF}")
                  ROLE=$(jq    --raw-output ".assume_roles[\"${aws_account}\"][\"aws_role\"]"    "${CONF}")
               MFA_ARN=$(jq    --raw-output ".assume_roles[\"${aws_account}\"][\"aws_mfa_arn\"]" "${CONF}")
  MAX_SESSION_DURATION=$(jq    --raw-output ".assume_roles[\"${aws_account}\"][\"max_session_duration\"]" "${CONF}")
        ASSUMEROLE_ENV=$(jq -j --raw-output ".assume_roles[\"${aws_account}\"] | select(.environment != null) | .environment[] | \"export \", .name, \"=\", .value, \";\n\"" "${CONF}")
                SSHKEY=$(jq    --raw-output ".assume_roles[\"${aws_account}\"] | select(.sshkey != null) | .sshkey" "${CONF}")
                AWS_DEFAULT_REGION=$(jq    --raw-output ".assume_roles[\"${aws_account}\"][\"aws_region\"]" "${CONF}")
                [[ "${AWS_DEFAULT_REGION}" == "null" ]] && AWS_DEFAULT_REGION="eu-central-1"
}

GetCodeartifactInfo() {
  local profile="${1:-NA}"
  local jq_output

  jq_output=$(jq ".codeartifact[] | select(.profile==\"${profile}\")" "${CONF}")
  if [[ -n "${jq_output}" ]]; then
    IsVerbose && info "Found CodeArtifact entry for ${profile} in ${CONF}."
    ASSUMEROLE_CODEARTIFACT_ID=$(jq --raw-output ".codeartifact[] | select(.profile==\"${profile}\") | .id // \"codeartifact\"" "${CONF}")
    ASSUMEROLE_CODEARTIFACT_USER=$(jq --raw-output ".codeartifact[] | select(.profile==\"${profile}\") | .username // \"aws\"" "${CONF}")
    ASSUMEROLE_CODEARTIFACT_DOMAINOWNER=$(jq --raw-output ".codeartifact[] | select(.profile==\"${profile}\") | .domain_owner // \"NOT_FOUND\"" "${CONF}")
    ASSUMEROLE_CODEARTIFACT_DOMAIN=$(jq --raw-output ".codeartifact[] | select(.profile==\"${profile}\") | .domain // \"NOT_FOUND\"" "${CONF}")

    if [[ "${ASSUMEROLE_CODEARTIFACT_DOMAIN}" == "NOT_FOUND" ]]; then
      warning "No codeartifact domain found in ~/.assumerole, settings.xml will not be updated"
      ASSUMEROLE_CODEARTIFACT=0
    fi
    if [[ "${ASSUMEROLE_CODEARTIFACT_DOMAINOWNER}" == "NOT_FOUND" ]]; then
      warning "No codeartifact domain owner found in ~/.assumerole, settings.xml will not be updated"
      ASSUMEROLE_CODEARTIFACT=0
    fi
  else
    warning "No CodeArtifact entry found for ${profile} in ${CONF}."
    ASSUMEROLE_CODEARTIFACT=0
  fi
}

CreateCredentials() {
  local aws_token
  local role_session_name
  local json_output

  ### Check config before asking for the MFA token
  GetAccountInfo

  [[ "${PROFILE}" = "null" ]] && { IsVerbose && fail "aws_profile missing for account ${aws_account} in ${CONF}"; }
  [[ "${ACCOUNT}" = "null" ]] && { IsVerbose && fail "aws_account missing for account ${aws_account} in ${CONF}"; }
  [[ "${ROLE}" = "null" ]]    && { IsVerbose && fail "aws_role missing for account ${aws_account} in ${CONF}"; }
  [[ "${MFA_ARN}" = "null" ]] && { IsVerbose && fail "aws_mfa_arn missing for account ${aws_account} in ${CONF}"; }
  [[ "${MAX_SESSION_DURATION}" = "null" ]] && MAX_SESSION_DURATION=${AWS_STS_DURATION_SECONDS:-3600}

  ### Get MFA token from commandline or ask for it
  if [[ "${OP_INSTALLED}" == "true" ]] && [[ -z ${2} ]]; then GetOtpInfo; fi

  if [[ -n ${2} ]]
  then
    aws_token=${2}
  elif [[ ${OTP} ]]
  then
    aws_token=${OTP}
  else
    printf "MFA token: "
    read -r aws_token
  fi

  unset AWS_ACCESS_KEY_ID
  unset AWS_SECRET_ACCESS_KEY
  unset AWS_SESSION_TOKEN
  unset AWS_ACCOUNT
  unset AWS_ACCOUNT_ID

  role_session_name=$(echo "${ROLE:0:40}${$}" | tr '/' '-' | tr ':' '-')

  json_output=$(aws sts assume-role \
           --role-arn "arn:aws:iam::${ACCOUNT}:role/${ROLE}" \
           --role-session-name "${role_session_name}" \
           --duration-seconds "${MAX_SESSION_DURATION}" \
           --serial-number "${MFA_ARN}" \
           --token-code "${aws_token}" \
           --profile "${PROFILE}" 2>&1)

  if [[ $? -ne 0 ]]; then
    IsVerbose && plain "${json_output}"
    exit 1
  fi 

      AWS_ACCESS_KEY_ID=$(echo "${json_output}" | jq --raw-output ".Credentials[\"AccessKeyId\"]")
  AWS_SECRET_ACCESS_KEY=$(echo "${json_output}" | jq --raw-output ".Credentials[\"SecretAccessKey\"]")
      AWS_SESSION_TOKEN=$(echo "${json_output}" | jq --raw-output ".Credentials[\"SessionToken\"]")
         AWS_EXPIRATION=$(echo "${json_output}" | jq --raw-output ".Credentials[\"Expiration\"]")
            AWS_ACCOUNT="${aws_account}"
         AWS_ACCOUNT_ID="${ACCOUNT}"

  export AWS_ACCESS_KEY_ID="${AWS_ACCESS_KEY_ID}"
  export AWS_SECRET_ACCESS_KEY="${AWS_SECRET_ACCESS_KEY}"
  export AWS_SESSION_TOKEN="${AWS_SESSION_TOKEN}"
  export AWS_ACCOUNT
  export AWS_ACCOUNT_ID
  export AWS_DEFAULT_REGION

  SaveCredentials
}

SetEnvironment() {
  [[ -n "${ASSUMEROLE_ENV}" ]] && {
    IsVerbose && info "Setting environment for profile ${PROFILE}"
    eval "$(echo "${ASSUMEROLE_ENV}")"
  }
}

AddSshKey() {
  [[ -n ${SSHKEY} ]] && {
    if ! ssh-add -l | grep -q "$(basename "${SSHKEY}")"
    then
      IsVerbose && info "Adding key ${SSHKEY} for profile ${PROFILE}"
      ssh-add "${SSHKEY}"
    fi
  }
}

GetAccountName() {
  if [[ -n ${1} && ${AVAILABLE_PROFILES} == *${1}* ]]
  then
    # Argument passed on commandline is a valid profile
    IsVerbose && info "The profile ${1} passed on the commandline is a valid profile."
    aws_account=${1}
  else
    SelectProfile ${AVAILABLE_PROFILES}
  fi
}

IsVerbose() {
  if [[ -z ${ASSUMEROLE_COMMAND} && ${ASSUMEROLE_QUIET:-0} -eq 0 ]]; then
    return 0
  else
    return 1
  fi
}

PrintAccountInfo() {
  GetAccountInfo
  
  IsVerbose && plain "Account Name: ${aws_account}"
  IsVerbose && plain "Account ID: ${ACCOUNT}"
  IsVerbose && plain "Assumed Role: ${ROLE}"
}

CheckCommand() {
  command -v "${1:-not_present}" >/dev/null 2>&1
  return ${?}
}

DoCodeArtifactStuff() {
  local assumerole_codeartifact_auth_token

  GetCodeartifactInfo "${PROFILE}"
  export CODEARTIFACT_AUTH_TOKEN=""
  if [[ "${ASSUMEROLE_CODEARTIFACT}" -eq "1" ]]; then
    ### Try to get a authentication token
    if assumerole_codeartifact_auth_token=$(aws codeartifact get-authorization-token \
       --domain "${ASSUMEROLE_CODEARTIFACT_DOMAIN}" \
       --domain-owner "${ASSUMEROLE_CODEARTIFACT_DOMAINOWNER}" \
       --query authorizationToken \
       --output text 2>/dev/null); then
      IsVerbose && success "Successfully retrieved a AWS CodeArtifact authentication token"
      ### Set the auth token as envvar
      export CODEARTIFACT_AUTH_TOKEN="${assumerole_codeartifact_auth_token}"
      ### Add to settings.xml
      if xml sel -t -v "/settings/servers/server[id='${ASSUMEROLE_CODEARTIFACT_ID}']/id" ~/.m2/settings.xml >/dev/null 2>&1; then
        IsVerbose && info "settings.xml already contains entry for ID ${ASSUMEROLE_CODEARTIFACT_ID}, updating the password"
        ### Entry exists, replace the password
        xml ed --inplace --update "/settings/servers/server[id='codeartifact']/password" \
                         --value "${assumerole_codeartifact_auth_token}" \
            ~/.m2/settings.xml
      else
        ### Add server entry for codeartifact
        IsVerbose && info "settings.xml does not contain entry for ID ${ASSUMEROLE_CODEARTIFACT_ID}, adding it ..."
        xml ed --inplace --subnode "/settings/servers" --type elem -n "server" --value "" \
                         --subnode "//server[last()]" --type elem -n "id" --value "${ASSUMEROLE_CODEARTIFACT_ID}" \
                         --subnode "//server[last()]" --type elem -n "username" --value "${ASSUMEROLE_CODEARTIFACT_USER}" \
                         --subnode "//server[last()]" --type elem -n "password" --value "${assumerole_codeartifact_auth_token}" \
            ~/.m2/settings.xml
      fi
    else
      # This should not cause the rest to fail
      warning "Unable to obtain a AWS CodeArtifact authentication token."
      warning "The role you are assuming does not have the permissions for the CodeArtifact"
      warning "GetAuthorizationToken operation."
    fi
  else
    info "Skipping Maven/AWS CodeArtifact setup"
  fi
}

gray='\033[37m'
blue='\033[36m'
red='\033[31m'
green='\033[32m'
orange='\033[33m'
reset='\033[0m'

info()    { echo -e "${blue}INFO: $*${reset}"; }
plain()   { echo -e "${green}$*${reset}"; }
warning() { echo -e "${orange}WARN: $*${reset}"; }
error()   { echo -e "${red}ERROR: $*${reset}"; }
success() { echo -e "${green}✔ $*${reset}"; }
fail()    { echo -e "${red}✖ $*${reset}"; exit 1; }
debug()   { [[ "${DEBUG}" == "true" ]] && echo -e "${gray}DEBUG: $*${reset}" || true; }

# Start with unsetting the current AWS_* envvars to avoid namespace pollution
SetInitialEnvironment

# Perform some checks
CheckCommand jq || fail "jq is required by $(basename "${0}"), please install jq and try again."
if [[ "${ASSUMEROLE_CODEARTIFACT}" -eq "1" ]]; then
  CheckCommand xml || fail "xml is required by $(basename "${0}"), please install xmlstarlet and try again"
fi

export CONF="${HOME}/.assumerole"
AVAILABLE_PROFILES=$(jq --raw-output ".assume_roles | to_entries[] | .key"  "${CONF}")

if [[ -n ${1} && ${1} == accountlist ]]
then
  IsVerbose && plain "info"
  IsVerbose && plain "accountlist"
  IsVerbose && plain "${AVAILABLE_PROFILES}"
  exit 0
elif [[ -n ${1} && ${1} == accountlist_text ]]
then
  echo "${AVAILABLE_PROFILES}"
  exit 0
elif [[ -n ${1} && ${1} == info ]]
then
  shift 1
  GetAccountName "$@"
  PrintAccountInfo
  exit 0
fi

# Check if op is installed
if [[ $(type op) ]] &>/dev/null
 then
   success "1Password CLI is installed"
   OP_INSTALLED=true
 else
   info "1Password CLI is not installed, install 1Password CLI to enable this feature"
   OP_INSTALLED=false
fi

GetAccountName "$@"

if ! CheckAndLoadFromCache
then
  CreateCredentials "$@"
fi

SetEnvironment
AddSshKey
DoCodeArtifactStuff

AWS_CALLER_IDENTITY=$(aws sts get-caller-identity)
JSONCALLERIDENTITYACCOUNT=$(echo "${AWS_CALLER_IDENTITY}" | jq --raw-output '.Account')
JSONCALLERIDENTITYROLEARN=$(echo "${AWS_CALLER_IDENTITY}" | jq --raw-output '.Arn')
ROLEBASENAME="$(basename "${ROLE}")"

if [[ ${JSONCALLERIDENTITYACCOUNT} == "${ACCOUNT}" ]]
then
  IsVerbose && info "Account of assumed role ${JSONCALLERIDENTITYACCOUNT} matches desired account ${ACCOUNT}"
  if [[ ${JSONCALLERIDENTITYROLEARN} == */${ROLEBASENAME}/* ]]
  then
    IsVerbose && info "Assumed role ${JSONCALLERIDENTITYROLEARN} matches desired role ${ROLE}"
    IsVerbose && info "The temporary credentials expire on ${AWS_EXPIRATION}"
    IsVerbose && info "Copy paste following commands to have the same credentials in"
    IsVerbose && info "  another shell"
    IsVerbose && plain "export AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}"
    IsVerbose && plain "export AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}"
    IsVerbose && plain "export AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN}"
    IsVerbose && plain "export CODEARTIFACT_AUTH_TOKEN=\"${CODEARTIFACT_AUTH_TOKEN}\""
    IsVerbose && plain "export AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION}"
    export aws_account
    if [[ -n ${ASSUMEROLE_COMMAND} ]]; then
      IsVerbose && info "Running command ${ASSUMEROLE_COMMAND}"
      ${ASSUMEROLE_COMMAND}
    else
      IsVerbose && info "Starting a new shell"
      ${SHELL}
    fi
  else
    IsVerbose && error "Assumed role ${JSONCALLERIDENTITYROLEARN} does not match desired role ${ROLE}"
    IsVerbose && error "Unsetting environment"
    SetInitialEnvironment
  fi
else
  IsVerbose && error "Account of assumed role ${JSONCALLERIDENTITYACCOUNT} does not match desired account ${ACCOUNT}"
  IsVerbose && error "Unsetting environment"
  SetInitialEnvironment
fi