Created
June 15, 2022 16:14
-
-
Save vicyap/1dc1e61b547ae8491978fb046cace519 to your computer and use it in GitHub Desktop.
A boto3 Session Factory that creates sessions that auto refresh their credentials. Primarily intended for assuming a role, since assume role credentials expire.
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
| from boto3 import Session | |
| from botocore.credentials import RefreshableCredentials | |
| from botocore.session import get_session | |
| class AutoRefreshSessionFactory: | |
| """A Factory that creates boto3.Session sessions that auto refresh their credentials. | |
| Primarily intended for assuming a role, since assume role credentials expire. | |
| """ | |
| def __init__(self, role_arn: str, role_session_name: str, duration_seconds: int = 900): | |
| self.role_arn = role_arn | |
| self.role_session_name = role_session_name | |
| self.duration_seconds = duration_seconds | |
| self._sts_client = Session().client("sts") | |
| def _refresh(self): | |
| credentials_response = self._sts_client.assume_role( | |
| RoleArn=self.role_arn, | |
| RoleSessionName=self.role_session_name, | |
| DurationSeconds=self.duration_seconds, | |
| ).get("Credentials") | |
| credentials = { | |
| "access_key": credentials_response.get("AccessKeyId"), | |
| "secret_key": credentials_response.get("SecretAccessKey"), | |
| "token": credentials_response.get("SessionToken"), | |
| "expiry_time": credentials_response.get("Expiration").isoformat(), | |
| } | |
| return credentials | |
| def get_session(self) -> Session: | |
| refreshable_credentials = RefreshableCredentials.create_from_metadata( | |
| metadata=self._refresh(), | |
| refresh_using=self._refresh, | |
| method="sts-assume-role", | |
| ) | |
| session = get_session() | |
| session._credentials = refreshable_credentials | |
| autorefresh_session = Session(botocore_session=session) | |
| return autorefresh_session | |
| assume_role_arn = "arn:aws:iam::account:role/role-name-with-path" | |
| session = AutoRefreshSessionFactory(assume_role_arn, "my_session_name").get_session() | |
| caller_identity = session.client("sts").get_caller_identity() | |
| print(f"successfully assumed role {caller_identity.get('Arn')}") |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment