feat: Added ability to assume a Role using STS (#33)

Author: AlastairDewar

odd_collector_aws/aws/aws_client.py

@@ -3,6 +3,8 @@
 
 from odd_collector_aws.domain.plugin import AwsPlugin
 
+from logging import logger
+
 
 class AwsClient:
     session: Session
@@ -13,13 +15,32 @@ def __init__(self, config: AwsPlugin):
         self._init_session()
 
     def _init_session(self):
+
         self.session = Session(
             aws_access_key_id=self._config.aws_access_key_id,
             aws_secret_access_key=self._config.aws_secret_access_key,
             aws_session_token=self._config.aws_session_token,
             region_name=self._config.aws_region,
         )
 
+        if self._config.aws_role_arn and self._config.aws_role_session_name:
+            try:
+                assumed_role_response = self.session.client("sts").assume_role(
+                    RoleArn=self._config.aws_role_arn,
+                    RoleSessionName=self._config.aws_role_session_name
+                )
+
+                if assumed_role_response.Credentials:
+                    self.session = Session(
+                        aws_access_key_id=assumed_role_response.Credentials.AccessKeyId,
+                        aws_secret_access_key=assumed_role_response.Credentials.SecretAccessKey,
+                        aws_session_token=assumed_role_response.Credentials.SessionToken,
+                        region_name=self._config.aws_region
+                    )
+            except Exception:
+                logger.debug(
+                    "Error assuming AWS Role", exc_info=True)
+
     def get_client(self, service_name: str) -> BaseClient:
         return self.session.client(service_name, endpoint_url=self._config.endpoint_url)
 

odd_collector_aws/domain/plugin.py

@@ -15,6 +15,8 @@ class AwsPlugin(Plugin):
     aws_session_token: Optional[str]
     aws_account_id: Optional[str]
     endpoint_url: Optional[str]
+    aws_role_arn: Optional[str]
+    aws_role_session_name: Optional[str]
 
 
 class GluePlugin(AwsPlugin):