boto3-assist 0.32.0__py3-none-any.whl

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/cloudwatch/cloudwatch_connection.py

@@ -0,0 +1,84 @@
+"""
+Geek Cafe, LLC
+Maintainers: Eric Wilson
+MIT License.  See Project Root for the license information.
+"""
+
+from typing import Optional
+from typing import TYPE_CHECKING
+
+from aws_lambda_powertools import Logger
+
+
+from boto3_assist.cloudwatch.cloudwatch_connection_tracker import (
+    CloudWatchConnectionTracker,
+)
+from boto3_assist.connection import Connection
+
+if TYPE_CHECKING:
+    from mypy_boto3_cloudwatch import CloudWatchClient, CloudWatchServiceResource
+else:
+    CloudWatchClient = object
+    CloudWatchServiceResource = object
+
+
+logger = Logger()
+tracker: CloudWatchConnectionTracker = CloudWatchConnectionTracker()
+
+
+class CloudWatchConnection(Connection):
+    """CW Environment"""
+
+    def __init__(
+        self,
+        *,
+        aws_profile: Optional[str] = None,
+        aws_region: Optional[str] = None,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+    ) -> None:
+        super().__init__(
+            service_name="cloudwatch",
+            aws_profile=aws_profile,
+            aws_region=aws_region,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+        )
+
+        self.__client: CloudWatchClient | None = None
+        self.__resource: CloudWatchServiceResource | None = None
+
+        self.raise_on_error: bool = True
+
+    @property
+    def client(self) -> CloudWatchClient:
+        """CloudWatch Client Connection"""
+        if self.__client is None:
+            logger.info("Creating CloudWatch Client")
+            self.__client = self.session.client
+
+        if self.raise_on_error and self.__client is None:
+            raise RuntimeError("CloudWatch Client is not available")
+        return self.__client
+
+    @client.setter
+    def client(self, value: CloudWatchClient):
+        logger.info("Setting CloudWatch Client")
+        self.__client = value
+
+    @property
+    def resource(self) -> CloudWatchServiceResource:
+        """CloudWatch Resource Connection"""
+        if self.__resource is None:
+            logger.info("Creating CloudWatch Resource")
+            self.__resource = self.session.resource
+
+        if self.raise_on_error and self.__resource is None:
+            raise RuntimeError("CloudWatch Resource is not available")
+
+        return self.__resource
+
+    @resource.setter
+    def resource(self, value: CloudWatchServiceResource):
+        logger.info("Setting CloudWatch Resource")
+        self.__resource = value

boto3_assist/cloudwatch/cloudwatch_connection_tracker.py

@@ -0,0 +1,17 @@
+"""
+Geek Cafe, LLC
+Maintainers: Eric Wilson
+MIT License.  See Project Root for the license information.
+"""
+
+from boto3_assist.connection_tracker import ConnectionTracker
+
+
+class CloudWatchConnectionTracker(ConnectionTracker):
+    """
+    Tracks CloudWatch Connection Requests.
+    Useful in for performance tuning and debugging.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()

boto3_assist/cloudwatch/cloudwatch_log_connection.py

@@ -0,0 +1,62 @@
+"""
+Geek Cafe, LLC
+Maintainers: Eric Wilson
+MIT License.  See Project Root for the license information.
+"""
+
+from typing import Optional
+from typing import TYPE_CHECKING
+
+from aws_lambda_powertools import Logger
+
+from boto3_assist.cloudwatch.cloudwatch_connection_tracker import (
+    CloudWatchConnectionTracker,
+)
+from boto3_assist.connection import Connection
+
+if TYPE_CHECKING:
+    from mypy_boto3_logs import CloudWatchLogsClient
+else:
+    CloudWatchLogsClient = object
+
+
+logger = Logger()
+tracker: CloudWatchConnectionTracker = CloudWatchConnectionTracker()
+
+
+class CloudWatchConnection(Connection):
+    """CW Logs Environment"""
+
+    def __init__(
+        self,
+        *,
+        aws_profile: Optional[str] = None,
+        aws_region: Optional[str] = None,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+    ) -> None:
+        super().__init__(
+            service_name="logs",
+            aws_profile=aws_profile,
+            aws_region=aws_region,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+        )
+
+        self.__client: CloudWatchLogsClient | None = None
+
+    @property
+    def client(self) -> CloudWatchLogsClient:
+        """CloudWatch Client Connection"""
+        if self.__client is None:
+            logger.debug("Creating CloudWatch Client")
+            self.__client = self.session.client
+
+        if self.raise_on_error and self.__client is None:
+            raise RuntimeError("CloudWatch Client is not available")
+        return self.__client
+
+    @client.setter
+    def client(self, value: CloudWatchLogsClient):
+        logger.debug("Setting CloudWatch Client")
+        self.__client = value

boto3_assist/cloudwatch/cloudwatch_logs.py

@@ -0,0 +1,39 @@
+"""
+Geek Cafe, LLC
+Maintainers: Eric Wilson
+MIT License.  See Project Root for the license information.
+"""
+
+from typing import Optional, List, Dict, Any
+from boto3_assist.cloudwatch.cloudwatch_log_connection import CloudWatchConnection
+
+
+class CloudWatchLogs(CloudWatchConnection):
+    def __init__(
+        self,
+        *,
+        aws_profile: Optional[str] = None,
+        aws_region: Optional[str] = None,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+    ) -> None:
+        super().__init__(
+            aws_profile=aws_profile,
+            aws_region=aws_region,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+        )
+
+    def list_log_groups(self):
+        """Retrieve all log groups in the AWS account."""
+        log_groups: List[Dict[str, Any]] = []
+        paginator = self.client.get_paginator("describe_log_groups")
+        for page in paginator.paginate():
+            log_groups.extend(page["logGroups"])  # type: ignore[arg-type]
+        return log_groups
+
+
+def main():
+    query: CloudWatchLogs = CloudWatchLogs()
+    result = query.list_log_groups()
+    print(result)

boto3_assist/cloudwatch/cloudwatch_query.py

@@ -0,0 +1,191 @@
+"""
+Geek Cafe, LLC
+Maintainers: Eric Wilson
+MIT License.  See Project Root for the license information.
+"""
+
+import os
+from datetime import datetime, timedelta, UTC
+from typing import Optional, Dict, Any, List
+from boto3_assist.cloudwatch.cloudwatch_connection import CloudWatchConnection
+from boto3_assist.cloudwatch.cloudwatch_logs import CloudWatchLogs
+
+
+class CloudWatchQuery(CloudWatchConnection):
+    """Query Cloud Watch"""
+
+    def __init__(
+        self,
+        *,
+        aws_profile: Optional[str] = None,
+        aws_region: Optional[str] = None,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+    ) -> None:
+        super().__init__(
+            aws_profile=aws_profile,
+            aws_region=aws_region,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+        )
+
+        self.__cw_logs: CloudWatchLogs | None = None
+
+    @property
+    def cw_logs(self) -> CloudWatchLogs:
+        """CloudWatch Logs Connection"""
+        if self.__cw_logs is None:
+            self.__cw_logs = CloudWatchLogs(
+                aws_profile=self.aws_profile,
+                aws_region=self.aws_region,
+                aws_access_key_id=self.aws_access_key_id,
+                aws_secret_access_key=self.aws_secret_access_key,
+            )
+        return self.__cw_logs
+
+    def get_log_group_size(
+        self, log_group_name: str, start_time: datetime, end_time: datetime
+    ) -> Dict[str, Any]:
+        """
+        Get the log group size for a given period of time
+        Args:
+            log_group_name (str): _description_
+            start_time (datetime): _description_
+            end_time (datetime): _description_
+
+        Returns:
+            _type_: _description_
+        """
+        response = self.client.get_metric_data(
+            MetricDataQueries=[
+                {
+                    "Id": "storedBytes",
+                    "MetricStat": {
+                        "Metric": {
+                            "Namespace": "AWS/Logs",
+                            # "MetricName": "StoredBytes",
+                            "MetricName": "IncomingBytes",
+                            "Dimensions": [
+                                {"Name": "LogGroupName", "Value": log_group_name}
+                            ],
+                        },
+                        "Period": 86400,  # Daily data
+                        "Stat": "Sum",
+                    },
+                    "ReturnData": True,
+                },
+            ],
+            StartTime=start_time,
+            EndTime=end_time,
+        )
+
+        # Extract the total size in bytes for the period
+        size: float = 0.0
+        if response["MetricDataResults"]:
+            # Access the first MetricDataResult
+            metric_data_result = response["MetricDataResults"][0]
+            # Sum the values if they exist
+            size = (
+                sum(metric_data_result["Values"]) if metric_data_result["Values"] else 0
+            )
+        else:
+            size = 0
+
+        size_mb = size / (1024 * 1024)
+        size_gb = size_mb / 1024
+        resp: Dict[str, Any] = {
+            "LogGroupName": log_group_name,
+            "Size": {
+                "Bytes": size,
+                "MB": size_mb,
+                "GB": size_gb,
+            },
+            "StartDate": start_time.isoformat(),
+            "EndDate": end_time.isoformat(),
+        }
+
+        return resp
+
+    def get_log_sizes(
+        self,
+        start_date_time: datetime | None = None,
+        end_date_time: datetime | None = None,
+        days: int | None = 7,
+        top: int = 0,
+    ) -> List[Dict[str, Any]]:
+        """
+        Gets the log sizes for all log groups
+
+        Args:
+            start_date_time (datetime | None, optional): The Start Date. Defaults to None.
+                If None it's set to now in UTC time - the days field
+            end_date_time (datetime | None, optional): he Start Date. Defaults to None.
+                If None it's set to not in UTC time
+            days (int | None, optional): The days offset. Defaults to 7.
+            top (int, optional): If greater than zero it will return the top x after sorting
+                Defaults to 0.
+
+        Returns:
+            list: _description_
+        """
+        if not days:
+            days = 7
+        start_time = start_date_time or (datetime.now(UTC) - timedelta(days=days))
+        end_time = end_date_time or datetime.now(UTC)
+
+        # Step 1: List all log groups
+        log_groups = self.cw_logs.list_log_groups()
+        log_group_sizes = []
+
+        # Step 2: Get sizes for each log group
+        for log_group in log_groups:
+            log_group_name = log_group["logGroupName"]
+
+            size_info = self.get_log_group_size(log_group_name, start_time, end_time)
+            log_group_sizes.append(size_info)
+
+        # Step 3: Sort by size
+        # top_log_groups = sorted(log_group_sizes, key=lambda x: x[1], reverse=True)
+        top_log_groups = sorted(
+            log_group_sizes,
+            key=lambda x: x.get("Size", {}).get("Bytes", 0),
+            reverse=True,
+        )
+        if top and top > 0:
+            # find the top x if provided
+            top_log_groups = top_log_groups[:top]
+
+        return top_log_groups
+
+
+def main():
+    log_group = os.environ.get("LOG_GROUP_QUERY_SAMPLE", "<enter-log-group-here>")
+    start = datetime.now() - timedelta(days=7)  # Last 30 days
+    end = datetime.now()
+    cw_query: CloudWatchQuery = CloudWatchQuery()
+    result = cw_query.get_log_group_size(log_group, start, end)
+    print(result)
+
+    top = 25
+    days = 7
+    top_log_groups = cw_query.get_log_sizes(top=top, days=days)
+    print(f"Top {top} log groups by size for the last week:")
+
+    for top_log_group in top_log_groups:
+        log_group_name = top_log_group["LogGroupName"]
+        size_in_bytes = top_log_group.get("Size", {}).get("Bytes", 0)
+        size_in_megs = top_log_group.get("Size", {}).get("MB", 0)
+        size_in_gigs = top_log_group.get("Size", {}).get("GB", 0)
+        size: str = ""
+        if size_in_gigs > 1:
+            size = f"{size_in_gigs:.2f} GB"
+        elif size_in_megs > 1:
+            size = f"{size_in_megs:.2f} MB"
+        else:
+            size = f"{size_in_bytes} bytes"
+
+        print(f"{size}: {log_group_name}")
+
+
+if __name__ == "__main__":
+    main()

boto3_assist/cognito/cognito_authorizer.py

@@ -0,0 +1,169 @@
+"""
+Geek Cafe, LLC
+Maintainers: Eric Wilson
+MIT License.  See Project Root for the license information.
+"""
+
+import time
+from typing import Any, Dict, List
+
+import jwt  # PyJWT
+from aws_lambda_powertools import Logger
+from jwt import InvalidTokenError, PyJWKClient
+
+from boto3_assist.boto3session import Boto3SessionManager
+from boto3_assist.cognito.jwks_cache import JwksCache
+
+logger = Logger()
+
+jwks_cache = JwksCache()
+
+
+class CognitoCustomAuthorizer:
+    """Cognito Custom Authorizer"""
+
+    def __init__(self):
+        self.__client_connections: Dict[str, Any] = {}
+
+    def __get_client_connection(
+        self, user_pool_id: str, refresh_client: bool = False
+    ) -> Any:
+        """Get the client connection to cognito"""
+        region = user_pool_id.split("_")[0]
+        client = self.__client_connections.get(region)
+        if refresh_client:
+            client = None
+        if not client:
+            session = Boto3SessionManager(service_name="cognito-idp", aws_region=region)
+            client = session.client
+            # boto3.client("cognito-idp", region_name=region)
+            self.__client_connections[region] = client
+
+        return client
+
+    def generate_policy(
+        self, user_pools: str | List[str], event: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Generates the policy for the authorizer"""
+
+        token = event["authorizationToken"]
+        user_pools = self.__to_list(user_pools=user_pools)
+        for user_pool_id in user_pools:
+            try:
+                if not user_pool_id:
+                    continue
+                # up_id = self.__to_id(user_pool_id=user_pool_id)
+                # Decode the token, assuming RS256 (used by Cognito)
+                # decoded_token = self.decode_jwt(token=token, user_pool_id=up_id)
+                issuer = self.build_issuer_url(user_pool_id)
+                claims = self.decode_jwt(token, issuer)
+                # Token is valid, return an IAM policy
+                return self.__generate_policy_doc(
+                    principal_id=claims["sub"],
+                    effect="Allow",
+                    method_arn=event["methodArn"],
+                )
+
+            except InvalidTokenError as e:
+                # Token is not valid for this user pool, try the next one
+                logger.debug(str(e))
+                continue
+            except Exception as e:  # pylint: disable=w0718
+                logger.error(str(e))
+
+        # if we get here we deny it
+        return self.__generate_policy_doc(
+            principal_id="user",
+            effect="Deny",
+            method_arn=event["methodArn"],
+        )
+
+    def __generate_policy_doc(self, *, principal_id, effect, method_arn):
+        """Generate the policy doc"""
+        auth_response: Dict[str, Any] = {"principalId": principal_id}
+
+        if effect and method_arn:
+            policy_document = {
+                "Version": "2012-10-17",
+                "Statement": [
+                    {
+                        "Action": "execute-api:Invoke",
+                        "Effect": effect,
+                        "Resource": method_arn,
+                    }
+                ],
+            }
+            auth_response["policyDocument"] = policy_document
+
+        return auth_response
+
+    def build_issuer_url(self, user_pool_id: str) -> str:
+        """Build the issuer URL"""
+
+        # Extract region from user pool ID format, e.g., "us-east-1_ABC123"
+        region = user_pool_id.split("_")[0]
+        return f"https://cognito-idp.{region}.amazonaws.com/{user_pool_id}"
+
+    def __to_list(self, user_pools: str | List[str]) -> List[str]:
+        if isinstance(user_pools, str):
+            user_pools = str(user_pools).replace(";", ",").replace(" ", "")
+            user_pools = str(user_pools).split(",")
+        elif isinstance(user_pools, list):
+            pass
+        else:
+            logger.warning(
+                f"Missing/ Invalid user pool: {user_pools}, type: {type(user_pools)}"
+            )
+
+        return user_pools
+
+    def parse_jwt(self, token: str) -> dict:
+        """Parse the JWT"""
+        if "Bearer" in token:
+            token = token.replace("Bearer ", "")
+
+        decoded_jwt: dict = jwt.decode(token, options={"verify_signature": False})
+
+        return decoded_jwt
+
+    def decode_jwt(self, token: str, issuer) -> dict:
+        """Decode the JWT"""
+        # Get the public keys
+        # Get the JWKS client
+        jwks_client = self.get_jwks_client(issuer)
+        if "Bearer" in token:
+            token = token.replace("Bearer ", "")
+        # Fetch the signing key using the PyJWKClient
+        signing_key = jwks_client.get_signing_key_from_jwt(token)
+
+        # Decode and verify the token
+        claims = jwt.decode(
+            token,
+            signing_key.key,
+            algorithms=["RS256"],
+            # audience=user_pool_id,
+            issuer=issuer,
+            options={"verify_aud": False},  # Disable audience verification
+        )
+
+        # Optional claim checks
+        if claims["token_use"] != "id":
+            # we are currently only using ID tokens
+            raise RuntimeError("Not an id token")
+
+        return claims
+
+    def get_jwks_client(self, issuer) -> PyJWKClient:
+        """Get the JWT Client"""
+        if (
+            issuer in jwks_cache.cache
+            and (time.time() - jwks_cache.cache.get(issuer, {})["timestamp"]) < 3600
+        ):
+            # Return cached JWKS client if it’s less than an hour old
+            return jwks_cache.cache[issuer]["client"]
+        else:
+            # Create a new PyJWKClient and cache it
+            jwks_url = f"{issuer}/.well-known/jwks.json"
+            jwks_client = PyJWKClient(jwks_url)
+            jwks_cache.cache[issuer] = {"client": jwks_client, "timestamp": time.time()}
+            return jwks_client