main.py

#!/usr/bin/env python3

import json
import os
import subprocess
import sys
import time
import urllib.error
import urllib.request

ASSUME_ROLE_RETRY_COUNT = 5
ASSUME_ROLE_RETRY_BACKOFF_MILLISECONDS = 500


def exit_error(message: str):
    print("::error ::" + message.strip().replace("\n", "%0A"))
    sys.exit(1)


def mask_value(value: str):
    print("::add-mask::" + value)


def read_inputs() -> tuple[str, str, str, str, str]:
    def _env(key: str) -> str:
        return os.environ.get(key, "").strip()

    web_identity_role_arn = _env("INPUT_WEB_IDENTITY_ROLE_ARN")
    if web_identity_role_arn == "":
        exit_error("input web identity role ARN must be provided")

    assume_role_duration = _env("INPUT_ASSUME_ROLE_DURATION_SECONDS")
    if not assume_role_duration.isdigit():
        exit_error("input assume role duration seconds must be numeric")

    assume_role_session_name = _env("INPUT_ASSUME_ROLE_SESSION_NAME")
    if assume_role_session_name == "":
        exit_error("input assume role session name must be provided")

    aws_region = _env("INPUT_AWS_REGION")
    if aws_region == "":
        exit_error("input AWS region must be provided")

    return (
        web_identity_role_arn,
        _env("INPUT_ASSUME_ROLE_ARN"),
        assume_role_duration,
        assume_role_session_name,
        aws_region,
    )


def fetch_oidc_jwt() -> str:
    # fetch GitHub environment variables to make HTTP token fetch request
    req_url = os.environ.get("ACTIONS_ID_TOKEN_REQUEST_URL")
    if req_url is None:
        exit_error(
            "expected ACTIONS_ID_TOKEN_REQUEST_URL environment variable not found"
        )

    req_token = os.environ.get("ACTIONS_ID_TOKEN_REQUEST_TOKEN")
    if req_token is None:
        exit_error(
            "expected ACTIONS_ID_TOKEN_REQUEST_TOKEN environment variable not found"
        )

    # build HTTP request and execute
    request = urllib.request.Request(
        headers={"Authorization": "bearer " + req_token}, url=req_url
    )

    try:
        response = urllib.request.urlopen(request)
    except urllib.error.HTTPError as err:
        exit_error(
            "unexpected error fetching OIDC web identity token: " + str(err.read())
        )

    # parse response, return `value` property - containing the desired web identity JWT
    try:
        token_data = json.load(response)
    except json.decoder.JSONDecodeError:
        exit_error("unable to fetch OIDC web identity token - malformed HTTP response")

    response.close()
    return token_data.get("value", "")


def aws_sts_assume_role(
    cmd_name: str,
    role_arn: str,
    role_session_name: str,
    role_duration: str,
    web_identity_token: str = "",
    env_var_collection: dict[str, str] = {},
    retry_error_match_list: list[str] = [],
) -> tuple[str, str, str]:
    # build command argument list and environment variables to pass
    arg_list = [
        "aws",
        "sts",
        cmd_name,
        "--role-arn",
        role_arn,
        "--role-session-name",
        role_session_name,
        "--duration-seconds",
        role_duration,
        "--output",
        "json",
    ]

    if web_identity_token != "":
        arg_list += ["--web-identity-token", web_identity_token]

    # setting `AWS_EC2_METADATA_DISABLED` stops the AWS CLI from reaching out
    # to (a non-existent) metadata endpoint on GitHub hosted runners
    env_var_collection["AWS_EC2_METADATA_DISABLED"] = "true"
    env_var_collection["PATH"] = os.environ.get("PATH", "")

    retry_remain = ASSUME_ROLE_RETRY_COUNT
    retry_backoff_milliseconds = 0
    result_stdout = ""

    def allow_retry(result_stderr) -> bool:
        if retry_error_match_list and (retry_remain > 0):
            # if error message text contains item in error match list - allow retry
            for item in retry_error_match_list:
                if item in result_stderr:
                    return True

        return False

    while True:
        # execute AWS CLI command
        retry_remain -= 1
        try:
            result = subprocess.run(
                arg_list,
                encoding="utf-8",
                env=env_var_collection,
                stderr=subprocess.PIPE,
                stdout=subprocess.PIPE,
            )
        except FileNotFoundError as ex:
            exit_error("unable to assume role, AWS CLI installed?")

        if result.returncode == 0:
            # hold result of successful execution, exit retry loop
            result_stdout = result.stdout
            break

        # command execution resulted in error
        result_stderr = result.stderr.strip()
        if allow_retry(result_stderr):
            # backoff a little, move onto a retry attempt
            retry_backoff_milliseconds += ASSUME_ROLE_RETRY_BACKOFF_MILLISECONDS
            time.sleep(retry_backoff_milliseconds / 1000)
            continue

        exit_error("unable to assume role: \n" + result_stderr)

    # parse JSON response from AWS CLI assume role call
    try:
        assume_data = json.loads(result_stdout)
    except json.decoder.JSONDecodeError:
        exit_error("unable to assume role - malformed AWS CLI response")

    # pull out generated session credentials
    def credential_part(key: str) -> str:
        return assume_data.get("Credentials", {}).get(key, "")

    access_key_id = credential_part("AccessKeyId")
    secret_access_key = credential_part("SecretAccessKey")
    session_token = credential_part("SessionToken")
    if (access_key_id == "") or (secret_access_key == "") or (session_token == ""):
        exit_error("unable to assume role, missing expected response credentials")

    return (access_key_id, secret_access_key, session_token)


def write_aws_env_var_collection(
    access_key_id: str,
    secret_access_key: str,
    session_token: str,
    aws_region: str,
):
    # fetch and ensure GITHUB_ENV environment variable exists
    env_export_file_path = os.environ.get("GITHUB_ENV")
    if env_export_file_path is None:
        exit_error("expected GITHUB_ENV environment variable not found")

    # write AWS session credentials to GitHub environment file for job steps which follow
    fh = open(env_export_file_path, "w")
    fh.write(
        f"AWS_ACCESS_KEY_ID={access_key_id}\n"
        + f"AWS_SECRET_ACCESS_KEY={secret_access_key}\n"
        + f"AWS_SESSION_TOKEN={session_token}\n"
        + f"AWS_REGION={aws_region}\n"
    )

    fh.close()

    # mask any AWS session credential values from GitHub Actions logs if echoed in job steps which follow
    mask_value(access_key_id)
    mask_value(secret_access_key)
    mask_value(session_token)


def main():
    # read inputs passed to action
    (
        web_identity_assume_role_arn,
        assume_role_arn,
        assume_role_duration,
        assume_role_session_name,
        aws_region,
    ) = read_inputs()

    # assume IAM role ARN via OpenID Connect (OIDC)
    wi_token = fetch_oidc_jwt()
    (access_key_id, secret_access_key, session_token) = aws_sts_assume_role(
        "assume-role-with-web-identity",
        role_arn=web_identity_assume_role_arn,
        role_session_name=assume_role_session_name,
        role_duration=assume_role_duration,
        web_identity_token=wi_token,
        retry_error_match_list=[
            "Couldn't retrieve verification key from your identity provider",
        ],
    )

    if assume_role_arn != "":
        # from the OIDC IAM role, assume *another* final IAM role
        (access_key_id, secret_access_key, session_token) = aws_sts_assume_role(
            "assume-role",
            role_arn=assume_role_arn,
            role_session_name=assume_role_session_name,
            role_duration=assume_role_duration,
            env_var_collection={
                "AWS_ACCESS_KEY_ID": access_key_id,
                "AWS_SECRET_ACCESS_KEY": secret_access_key,
                "AWS_SESSION_TOKEN": session_token,
            },
        )

    write_aws_env_var_collection(
        access_key_id=access_key_id,
        secret_access_key=secret_access_key,
        session_token=session_token,
        aws_region=aws_region,
    )


if __name__ == "__main__":
    main()