#!/usr/bin/env bash
set -euo pipefail

source_profile="${1:?source AWS profile is required}"
role_arn="${2:?runtime-account cp9-developer role ARN is required}"
session_tag="${CP9_SESSION_NAME:-${3:-cp9-local}}"
# SSO profiles already resolve to assumed-role credentials. AWS role chaining caps
# the next AssumeRole session at 1 hour, even when the target role allows longer.
duration_seconds="${4:-3600}"
owner_arg="${5:-}"

# credential_process may inherit stale temporary credentials from the parent
# shell. Clear them so the source profile is the only credential source used to
# assume cp9-developer.
unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN AWS_SECURITY_TOKEN AWS_PROFILE

caller_arn="$(AWS_PAGER= aws sts get-caller-identity \
  --profile "${source_profile}" \
  --query 'Arn' \
  --output text)"

owner="${CP9_OWNER:-${owner_arg:-${caller_arn##*/}}}"
owner="$(printf '%s' "${owner}" | tr '[:upper:]' '[:lower:]')"
role_session_name="${CP9_ROLE_SESSION_NAME:-${owner}}"

assume_output="$(AWS_PAGER= aws sts assume-role \
  --profile "${source_profile}" \
  --role-arn "${role_arn}" \
  --role-session-name "${role_session_name}" \
  --duration-seconds "${duration_seconds}" \
  --tags "Key=cp9-owner,Value=${owner}" "Key=cp9-session,Value=${session_tag}" \
  --output json)"

ASSUME_OUTPUT="${assume_output}" python3 - <<'PY'
import json
import os

payload = json.loads(os.environ["ASSUME_OUTPUT"])
creds = payload["Credentials"]
print(json.dumps({
    "Version": 1,
    "AccessKeyId": creds["AccessKeyId"],
    "SecretAccessKey": creds["SecretAccessKey"],
    "SessionToken": creds["SessionToken"],
    "Expiration": creds["Expiration"],
}))
PY
