boto3-assist 0.12.0__tar.gz → 0.14.0__tar.gz

{boto3_assist-0.12.0 → boto3_assist-0.14.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: boto3_assist
-Version: 0.12.0
+Version: 0.14.0
 Summary: Additional boto3 wrappers to make your life a little easier
 Author-email: Eric Wilson <boto3-assist@geekcafe.com>
 License-File: LICENSE-EXPLAINED.txt

{boto3_assist-0.12.0 → boto3_assist-0.14.0}/pyproject.toml

@@ -15,7 +15,7 @@ addopts = "-m 'not integration'"
 
 [project]
 name = "boto3_assist"
-version = "0.12.0"
+version = "0.14.0"
 
 authors = [
   { name="Eric Wilson", email="boto3-assist@geekcafe.com" }

boto3_assist-0.14.0/src/boto3_assist/boto3session.py

@@ -0,0 +1,87 @@
+"""
+Geek Cafe, LLC
+Maintainers: Eric Wilson
+MIT License.  See Project Root for the license information.
+"""
+
+from typing import Optional, List, Any
+import boto3
+from botocore.config import Config
+from .session_setup_mixin import SessionSetupMixin
+from .role_assumption_mixin import RoleAssumptionMixin
+
+
+class Boto3SessionManager(SessionSetupMixin, RoleAssumptionMixin):
+    def __init__(
+        self,
+        service_name: str,
+        *,
+        aws_profile: Optional[str] = None,
+        aws_region: Optional[str] = None,
+        assume_role_arn: Optional[str] = None,
+        assume_role_chain: Optional[List[str]] = None,
+        assume_role_session_name: Optional[str] = None,
+        assume_role_duration_seconds: Optional[int] = 3600,
+        config: Optional[Config] = None,
+        aws_endpoint_url: Optional[str] = None,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+        aws_session_token: Optional[str] = None,
+    ):
+        self.service_name = service_name
+        self.aws_profile = aws_profile
+        self.aws_region = aws_region
+        self.config = config
+        self.endpoint_url = aws_endpoint_url
+        self.assume_role_chain = assume_role_chain or (
+            [assume_role_arn] if assume_role_arn else []
+        )
+        self.assume_role_session_name = (
+            assume_role_session_name or f"AssumeRoleSessionFor{service_name}"
+        )
+        self.assume_role_duration_seconds = assume_role_duration_seconds
+        self.aws_access_key_id = aws_access_key_id
+        self.aws_secret_access_key = aws_secret_access_key
+        self.aws_session_token = aws_session_token
+
+        self.__session: Optional[boto3.Session] = None
+        self.__client: Any = None
+        self.__resource: Any = None
+
+        self.__initialize()
+
+    def __initialize(self):
+        base_session = self._create_base_session(
+            self.aws_profile,
+            self.aws_region,
+            self.aws_access_key_id,
+            self.aws_secret_access_key,
+            self.aws_session_token,
+        )
+
+        if self.assume_role_chain:
+            self.__session = self._assume_roles_in_chain(
+                base_session,
+                self.assume_role_chain,
+                self.assume_role_session_name,
+                self.assume_role_duration_seconds,
+                self.aws_region,
+            )
+        else:
+            self.__session = base_session
+
+    @property
+    def client(self) -> Any:
+        if not self.__client:
+            self.__client = self.__session.client(
+                self.service_name, config=self.config, endpoint_url=self.endpoint_url
+            )
+        return self.__client
+
+    @property
+    def resource(self) -> Any:
+        if not self.__resource:
+            self.__resource = self.__session.resource(
+                self.service_name, config=self.config, endpoint_url=self.endpoint_url
+            )
+        return self.__resource

{boto3_assist-0.12.0 → boto3_assist-0.14.0}/src/boto3_assist/connection.py

@@ -4,7 +4,7 @@ Maintainers: Eric Wilson
 MIT License.  See Project Root for the license information.
 """
 
-from typing import Optional
+from typing import Optional, List
 
 from aws_lambda_powertools import Logger
 from boto3_assist.boto3session import Boto3SessionManager
@@ -31,6 +31,8 @@ class Connection:
         aws_secret_access_key: Optional[str] = None,
         aws_end_point_url: Optional[str] = None,
         assume_role_arn: Optional[str] = None,
+        assume_role_chain: Optional[List[str]] = None,
+        assume_role_duration_seconds: Optional[int] = 3600,
     ) -> None:
         self.__aws_profile = aws_profile
         self.__aws_region = aws_region
@@ -40,7 +42,8 @@ class Connection:
         self.__session: Boto3SessionManager | None = None
         self.__assume_role_arn: Optional[str] = assume_role_arn
         self.__service_name: str | None = service_name
-
+        self.__assume_role_chain = assume_role_chain
+        self.__assume_role_duration_seconds = assume_role_duration_seconds
         if self.__service_name is None:
             raise RuntimeError(
                 "Service Name is not available. The service name is required."
@@ -75,6 +78,8 @@ class Connection:
             aws_secret_access_key=self.aws_secret_access_key,
             aws_endpoint_url=self.end_point_url,
             assume_role_arn=self.__assume_role_arn,
+            assume_role_chain=self.__assume_role_chain,
+            assume_role_duration_seconds=self.__assume_role_duration_seconds,
         )
 
         tracker.add(service_name=self.service_name)

{boto3_assist-0.12.0 → boto3_assist-0.14.0}/src/boto3_assist/dynamodb/dynamodb_connection.py

@@ -4,7 +4,7 @@ Maintainers: Eric Wilson
 MIT License.  See Project Root for the license information.
 """
 
-from typing import Optional
+from typing import Optional, List
 from typing import TYPE_CHECKING
 
 from aws_lambda_powertools import Logger
@@ -33,6 +33,8 @@ class DynamoDBConnection(Connection):
         aws_access_key_id: Optional[str] = None,
         aws_secret_access_key: Optional[str] = None,
         assume_role_arn: Optional[str] = None,
+        assume_role_chain: Optional[List[str]] = None,
+        assume_role_duration_seconds: Optional[int] = 3600,
     ) -> None:
         super().__init__(
             service_name="dynamodb",
@@ -42,6 +44,8 @@ class DynamoDBConnection(Connection):
             aws_secret_access_key=aws_secret_access_key,
             aws_end_point_url=aws_end_point_url,
             assume_role_arn=assume_role_arn,
+            assume_role_chain=assume_role_chain,
+            assume_role_duration_seconds=assume_role_duration_seconds,
         )
 
         self.__dynamodb_client: DynamoDBClient | None = None

boto3_assist-0.14.0/src/boto3_assist/role_assumption_mixin.py

@@ -0,0 +1,38 @@
+"""
+Geek Cafe, LLC
+Maintainers: Eric Wilson
+MIT License.  See Project Root for the license information.
+"""
+
+import boto3
+from typing import List
+
+
+class RoleAssumptionMixin:
+    def _assume_roles_in_chain(
+        self,
+        base_session: boto3.Session,
+        role_chain: List[str],
+        session_name: str,
+        duration_seconds: int,
+        region: str,
+    ) -> boto3.Session:
+        session = base_session
+
+        for role_arn in role_chain:
+            sts_client = session.client("sts")
+            response = sts_client.assume_role(
+                RoleArn=role_arn,
+                RoleSessionName=session_name,
+                DurationSeconds=duration_seconds,
+            )
+            creds = response["Credentials"]
+
+            session = boto3.Session(
+                aws_access_key_id=creds["AccessKeyId"],
+                aws_secret_access_key=creds["SecretAccessKey"],
+                aws_session_token=creds["SessionToken"],
+                region_name=region,
+            )
+
+        return session

boto3_assist-0.14.0/src/boto3_assist/session_setup_mixin.py

@@ -0,0 +1,29 @@
+"""
+Geek Cafe, LLC
+Maintainers: Eric Wilson
+MIT License.  See Project Root for the license information.
+"""
+
+import boto3
+from typing import Optional
+
+
+class SessionSetupMixin:
+    def _create_base_session(
+        self,
+        aws_profile: Optional[str],
+        aws_region: Optional[str],
+        aws_access_key_id: Optional[str],
+        aws_secret_access_key: Optional[str],
+        aws_session_token: Optional[str],
+    ) -> boto3.Session:
+        try:
+            return boto3.Session(
+                profile_name=aws_profile,
+                region_name=aws_region,
+                aws_access_key_id=aws_access_key_id,
+                aws_secret_access_key=aws_secret_access_key,
+                aws_session_token=aws_session_token,
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to create boto3 session: {e}") from e

boto3_assist-0.14.0/src/boto3_assist/version.py

@@ -0,0 +1 @@
+__version__ = '0.14.0'

boto3_assist-0.12.0/src/boto3_assist/boto3session.py

@@ -1,190 +0,0 @@
-"""
-Geek Cafe, LLC
-Maintainers: Eric Wilson
-MIT License.  See Project Root for the license information.
-"""
-
-from typing import Any, Optional
-
-import boto3
-from aws_lambda_powertools import Logger
-from botocore.config import Config
-from botocore.exceptions import ProfileNotFound
-from boto3_assist.environment_services.environment_variables import EnvironmentVariables
-
-
-logger = Logger(__name__)
-
-
-class Boto3SessionManager:
-    """Manages Boto3 Sessions"""
-
-    def __init__(
-        self,
-        service_name: str,
-        *,
-        aws_profile: Optional[str] = None,
-        aws_region: Optional[str] = None,
-        assume_role_arn: Optional[str] = None,
-        assume_role_session_name: Optional[str] = None,
-        # cross_account_role_arn: Optional[str] = None,
-        config: Optional[Config] = None,
-        aws_endpoint_url: Optional[str] = None,
-        aws_access_key_id: Optional[str] = None,
-        aws_secret_access_key: Optional[str] = None,
-        aws_session_token: Optional[str] = None,
-    ):
-        self.service_name = service_name
-        self.aws_profile = aws_profile
-        self.aws_region = aws_region
-        self.assume_role_arn = assume_role_arn
-        self.assume_role_session_name = assume_role_session_name
-        self.config = config
-        # # self.cross_account_role_arn = cross_account_role_arn
-        self.endpoint_url = aws_endpoint_url
-        self.aws_access_key_id = aws_access_key_id
-        self.aws_secret_access_key = aws_secret_access_key
-        self.aws_session_token = aws_session_token
-
-        self.__session: Any = None
-        self.__client: Any = None
-        self.__resource: Any = None
-
-        self.__setup()
-
-    def __setup(self):
-        """Setup AWS session, client, and resource."""
-
-        profile = self.aws_profile or EnvironmentVariables.AWS.profile()
-        region = self.aws_region or EnvironmentVariables.AWS.region()
-
-        logger.debug("Connecting without assuming a role.")
-        self.__session = self.__get_aws_session(profile, region)
-
-        if profile:
-            print(f"Connecting with a profile: {profile}")
-
-        if self.assume_role_arn:
-            self.__assume_role()
-
-    def __assume_role(self):
-        """Assume an AWS IAM role."""
-        try:
-            if not self.__session:
-                raise RuntimeError(
-                    "Session must be established before assuming a role."
-                )
-
-            logger.debug(f"Assuming role {self.assume_role_arn}")
-
-            sts_client = self.__session.client("sts")
-            session_name = (
-                self.assume_role_session_name
-                or f"AssumeRoleSessionFor{self.service_name}"
-            )
-
-            assumed_role_response = sts_client.assume_role(
-                RoleArn=self.assume_role_arn,
-                RoleSessionName=session_name,
-            )
-            credentials = assumed_role_response["Credentials"]
-
-            # Now override the session with assumed credentials
-            self.__session = boto3.Session(
-                aws_access_key_id=credentials["AccessKeyId"],
-                aws_secret_access_key=credentials["SecretAccessKey"],
-                aws_session_token=credentials["SessionToken"],
-                region_name=self.aws_region,
-            )
-            logger.debug("Successfully assumed role and created new session.")
-
-        except Exception as e:
-            logger.error(f"Error assuming role: {e}")
-            raise RuntimeError(f"Failed to assume role {self.assume_role_arn}") from e
-
-    def __get_aws_session(
-        self, aws_profile: Optional[str] = None, aws_region: Optional[str] = None
-    ) -> boto3.Session | None:
-        """Get a boto3 session for AWS."""
-        logger.debug({"profile": aws_profile, "region": aws_region})
-        try:
-            self.aws_profile = aws_profile or EnvironmentVariables.AWS.profile()
-            self.aws_region = aws_region or EnvironmentVariables.AWS.region()
-            tmp_access_key_id = self.aws_access_key_id
-            tmp_secret_access_key = self.aws_secret_access_key
-            if not EnvironmentVariables.AWS.display_aws_access_key_id():
-                tmp_access_key_id = (
-                    "None" if tmp_access_key_id is None else "***************"
-                )
-            if not EnvironmentVariables.AWS.display_aws_secret_access_key():
-                tmp_secret_access_key = (
-                    "None" if tmp_secret_access_key is None else "***************"
-                )
-
-            logger.debug(
-                {
-                    "profile": self.aws_profile,
-                    "region": self.aws_region,
-                    "aws_access_key_id": tmp_access_key_id,
-                    "aws_secret_access_key": tmp_secret_access_key,
-                    "aws_session_token": (
-                        "*******" if self.aws_session_token is not None else ""
-                    ),
-                }
-            )
-            logger.debug("Creating boto3 session")
-            session = self.__create_boto3_session()
-        # if self.aws_profile or self.aws_region
-        # else boto3.Session()
-
-        except Exception as e:
-            logger.error(e)
-            raise RuntimeError("Failed to create a boto3 session.") from e
-
-        logger.debug({"session": session})
-        return session
-
-    @property
-    def client(self) -> Any:
-        """Return the boto3 client connection."""
-        if not self.__client:
-            logger.debug(f"Creating {self.service_name} client")
-            self.__client = self.__session.client(
-                self.service_name,
-                config=self.config,
-                endpoint_url=self.endpoint_url,
-            )
-
-        return self.__client
-
-    @property
-    def resource(self) -> Any:
-        """Return the boto3 resource connection."""
-        if not self.__resource:
-            logger.debug(f"Creating {self.service_name} resource")
-            self.__resource = self.__session.resource(
-                self.service_name,
-                config=self.config,
-                endpoint_url=self.endpoint_url,
-            )
-        return self.__resource
-
-    def __create_boto3_session(self) -> boto3.Session | None:
-        try:
-            logger.debug(f"Creating session for {self.service_name}")
-            session = boto3.Session(
-                profile_name=self.aws_profile,
-                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,
-            )
-            return session
-        except ProfileNotFound as e:
-            print(
-                f"An error occurred setting up the boto3 sessions. Profile not found: {e}"
-            )
-            raise e
-        except Exception as e:
-            print(f"An error occurred setting up the boto3 sessions: {e}")
-            raise e

boto3_assist-0.12.0/src/boto3_assist/version.py

@@ -1 +0,0 @@
-__version__ = '0.12.0'