Last active
July 19, 2024 15:53
-
-
Save dceoy/2b81705d25f5b18d8cece4b409fb2be0 to your computer and use it in GitHub Desktop.
[Python] Assume an AWS IAM role and make a session with the role
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/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, | |
| ) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment