aws-env.sh

#!/bin/bash

# Check argument for unauth
if [[ "$1" == "unauth" ]]; then
  echo "==> Clearing AWS session credentials..."
  unset AWS_ACCESS_KEY_ID
  unset AWS_SECRET_ACCESS_KEY
  unset AWS_SESSION_TOKEN
  unset AWS_REGION
  echo "==> Done."
  return 0 2>/dev/null || exit 0
fi

# Check if script is sourced (only enforce sourcing for auth, unauth works either way but better sourced)
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
  echo "Error: This script must be sourced, not executed directly."
  echo "Usage:" 
  echo "  source ./aws-env.sh [code]   # Authenticate"
  echo "  source ./aws-env.sh unauth   # Clear credentials"
  exit 1
fi

# Load .env file if it exists
if [[ -f .env ]]; then
  set -a
  source .env
  set +a
fi

# Check for MFA_DEVICE_ARN
if [[ -z "${MFA_DEVICE_ARN}" ]]; then
  echo "Error: MFA_DEVICE_ARN is not set. Please set it in .env or environment."
  return 1
fi

MFA_CODE="$1"

# Interactive prompt for MFA code if not provided
if [[ -z "${MFA_CODE}" ]]; then
  # Use printf to ensure compatibility with Zsh (where read -p means read from coprocess)
  printf "Enter MFA Code: "
  read MFA_CODE
fi

if [[ -z "${MFA_CODE}" ]]; then
  echo "Error: MFA Code is required."
  return 1
fi

# Get session token
echo "==> Authenticating with MFA..."

# Unset existing session credentials if any, otherwise STS call will fail
# because you cannot call GetSessionToken using temporary credentials.
unset AWS_SESSION_TOKEN
unset AWS_ACCESS_KEY_ID
unset AWS_SECRET_ACCESS_KEY

CREDENTIALS=$(aws sts get-session-token \
  --serial-number "${MFA_DEVICE_ARN}" \
  --token-code "${MFA_CODE}" \
  --output json)

if [[ $? -ne 0 ]]; then
  echo "Error: Failed to get session token."
  return 1
fi

# Parse credential
ACCESS_KEY=$(echo "${CREDENTIALS}" | jq -r '.Credentials.AccessKeyId')
SECRET_KEY=$(echo "${CREDENTIALS}" | jq -r '.Credentials.SecretAccessKey')
SESSION_TOKEN=$(echo "${CREDENTIALS}" | jq -r '.Credentials.SessionToken')

# Export variables
export AWS_ACCESS_KEY_ID="${ACCESS_KEY}"
export AWS_SECRET_ACCESS_KEY="${SECRET_KEY}"
export AWS_SESSION_TOKEN="${SESSION_TOKEN}"
export AWS_REGION="${REGION}"

echo
echo "==> Success! AWS session variables exported."
echo
echo "==> Access Key: ${ACCESS_KEY}"
echo "==> Region: ${REGION}"
echo
echo "==> You can run following commands to verify the session token:"
echo "echo && echo AWS_ACCESS_KEY_ID: \$AWS_ACCESS_KEY_ID && echo && echo AWS_SECRET_ACCESS_KEY: \$AWS_SECRET_ACCESS_KEY && echo && echo AWS_REGION: \$AWS_REGION && echo && echo AWS_SESSION_TOKEN: \$AWS_SESSION_TOKEN"
echo
echo "==> You can now run ./sync-dynamodb.sh"