Assume role wrapper scripts in bash and PowerShell (Using Role ARN from environment variable "ROLE_ARN")

assume_role.ps1

if (Get-Command "powershell" -ErrorAction SilentlyContinue) {
  $powershell = "powershell"
} else {
  $powershell = "pwsh"
}

& $powershell -Command {
  Set-StrictMode -Version 2.0

  Import-Module AWSPowerShell.NetCore

  $CurrentArn = (Get-STSCallerIdentity).Arn
  $RoleArn = (Get-Item Env:ROLE_ARN).Value

  $RoleCreds = (Use-STSRole -RoleArn $RoleArn -DurationSeconds 10800 -RoleSessionName $($CurrentArn.Substring($CurrentArn.LastIndexOf("/") + 1))).Credentials

  Get-ChildItem -Path Env: | Where-Object {($_.Name -Like "AWS_*") -and ($_.Name -NotMatch "^AWS(?:_DEFAULT)?_REGION$")} | Remove-Item

  $Env:AWS_ACCESS_KEY_ID = $RoleCreds.AccessKeyId
  $Env:AWS_SECRET_ACCESS_KEY = $RoleCreds.SecretAccessKey
  $Env:AWS_SESSION_TOKEN = $RoleCreds.SessionToken

  Invoke-Expression "$args"
  
  exit $LastExitCode
} -args $args

exit $LastExitCode

assume_role.sh

#!/bin/bash

set -euo pipefail

creds="$(aws sts assume-role --query 'Credentials' --output json --role-arn "$ROLE_ARN" --duration-seconds 10800 --role-session-name "$(aws sts get-caller-identity --query 'Arn' --output text | awk -F/ '{ print $NF }')")"

# Unset existing credentials except for region
unset $(compgen -v AWS_ | grep -Pv '^AWS_(?:DEFAULT_)?REGION')

export AWS_ACCESS_KEY_ID="$(echo "$creds" | jq -r '.AccessKeyId')"
export AWS_SECRET_ACCESS_KEY="$(echo "$creds" | jq -r '.SecretAccessKey')"
export AWS_SESSION_TOKEN="$(echo "$creds" | jq -r '.SessionToken')"

"$@"