[Python] Assume an AWS IAM role and make a session with the role

aws_sts_assume_role.py

#!/usr/bin/env python

import logging
from typing import Optional

import boto3
from pydantic import BaseModel, model_validator
from typing_extensions import Self


class Boto3SessionMaker(BaseModel):
    iam_role_name: str
    logger: logging.Logger = logging.getLogger(__name__)
    aws_region: Optional[str] = None
    aws_account_id: Optional[str] = None
    aws_access_key_id: Optional[str] = None
    aws_secret_access_key: Optional[str] = None
    aws_session_token: Optional[str] = None

    @model_validator(mode="after")
    def sts_assume_role(self) -> Self:
        if not self.aws_region:
            self.aws_region = boto3.session.Session().region_name
            self.logger.info(f"The default region is used: {self.aws_region}")
        self.logger.info(
            f"Retrieve the caller identity in the region: {self.aws_region}"
        )
        sts_client = boto3.client("sts", region_name=self.aws_region)
        caller_identity = sts_client.get_caller_identity()
        self.logger.debug(f"caller_identity: {caller_identity}")
        if not self.aws_account_id:
            self.aws_account_id = caller_identity["Account"]
            self.logger.info(
                f"The default account ID is used: {self.aws_account_id}"
            )
        current_role_or_user = caller_identity["Arn"].split("/")[-1]
        iam_role_arn = (
            f"arn:aws:iam::{self.aws_account_id}:role/{self.iam_role_name}"
        )

        self.logger.info(f"Assume role: {iam_role_arn}")
        response = sts_client.assume_role(
            RoleArn=iam_role_arn,
            RoleSessionName=(
                f"{self.iam_role_name}-session-{current_role_or_user}"
            ),
        )
        self.logger.debug(f"response: {response}")
        self.aws_access_key_id = response["Credentials"]["AccessKeyId"]
        self.aws_secret_access_key = response["Credentials"]["SecretAccessKey"]
        self.aws_session_token = response["Credentials"]["SessionToken"]
        return self

    def create_session(self) -> boto3.session.Session:
        self.logger.info(f"Create a boto3 session for the region: {self.aws_region}")
        return boto3.session.Session(
            region_name=self.aws_region,
            aws_access_key_id=self.aws_access_key_id,
            aws_secret_access_key=self.aws_secret_access_key,
            aws_session_token=self.aws_session_token,
        )

    def setup_default_session(self) -> None:
        self.logger.info(f"Setup the default boto3 session for the region: {self.aws_region}")
        boto3.setup_default_session(
            region_name=self.aws_region,
            aws_access_key_id=self.aws_access_key_id,
            aws_secret_access_key=self.aws_secret_access_key,
            aws_session_token=self.aws_session_token,
        )