#!/usr/bin/env bash

set -euo pipefail

# Usage function
usage() {
  echo "Usage: $0 --account-no <account_number> --role-name <role_name>"
  exit 1
}

# Parse named arguments
ACCOUNT_NO=""
ROLE_NAME=""

while [[ $# -gt 0 ]]; do
  case "$1" in
  --account-no)
    ACCOUNT_NO="$2"
    shift 2
    ;;
  --role-name)
    ROLE_NAME="$2"
    shift 2
    ;;
  *)
    echo "Unknown argument: $1"
    usage
    ;;
  esac
done

# Validate required arguments
if [[ -z $ACCOUNT_NO || -z $ROLE_NAME ]]; then
  echo "Error: --account-no and --role-name are required."
  usage
fi

# Assume role
ROLE_ARN="arn:aws:iam::${ACCOUNT_NO}:role/${ROLE_NAME}"
STS="$(aws sts assume-role --role-arn "$ROLE_ARN" --role-session-name "${ROLE_NAME}-session" --no-verify-ssl)"

# Extract credentials
AWS_ACCESS_KEY_ID="$(echo "$STS" | jq -r '.Credentials.AccessKeyId')"
AWS_SECRET_ACCESS_KEY="$(echo "$STS" | jq -r '.Credentials.SecretAccessKey')"
AWS_SESSION_TOKEN="$(echo "$STS" | jq -r '.Credentials.SessionToken')"
export AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN

# Output credentials (optional)
echo "Assumed role: $ROLE_ARN"
echo "AccessKeyId: $AWS_ACCESS_KEY_ID"
echo "SecretAccessKey: $AWS_SECRET_ACCESS_KEY"
echo "SessionToken: $AWS_SESSION_TOKEN"

# Verify identity
aws sts get-caller-identity --no-verify-ssl
