fixing flake8 issues in sts module

cyphronix · cyphronix · commit 824871010701 · 2024-12-11T17:35:10.000-07:00
diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/app.py
@@ -194,7 +194,7 @@ def load_sra_cloudwatch_dashboard() -> dict:
 ssm_params = sra_ssm_params.SRASSMParams()
 iam = sra_iam.SRAIAM()
 dynamodb = sra_dynamodb.SRADynamoDB()
-sts = sra_sts.sra_sts()
+sts = sra_sts.SRASTS()
 repo = sra_repo.SRARepo()
 s3 = sra_s3.SRAS3()
 lambdas = sra_lambda.SRALambda()
@@ -235,6 +235,7 @@ def get_resource_parameters(event: dict) -> None:
     repo.REPO_BRANCH = repo.REPO_ZIP_URL.split(".")[1].split("/")[len(repo.REPO_ZIP_URL.split(".")[1].split("/")) - 1]  # noqa: ECE001
     repo.SOLUTIONS_DIR = f"/tmp/aws-security-reference-architecture-examples-{repo.REPO_BRANCH}/aws_sra_examples/solutions"  # noqa: S108
 
+    # TODO(liamschn): the CONFIGURATION_ROLE needs to be a resource parameter
     sts.CONFIGURATION_ROLE = "sra-execution"
     governed_regions_param = ssm_params.get_ssm_parameter(
         ssm_params.MANAGEMENT_ACCOUNT_SESSION, REGION, "/sra/regions/customer-control-tower-regions"
diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_ssm_params.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_ssm_params.py
@@ -2,7 +2,7 @@
 
 Version: 1.0
 
-'common_prerequisites' solution in the repo, https://github.com/aws-samples/aws-security-reference-architecture-examples
+SSM Params module for SRA in the repo, https://github.com/aws-samples/aws-security-reference-architecture-examples
 
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 SPDX-License-Identifier: MIT-0
diff --git a/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_sts.py b/aws_sra_examples/solutions/genai/bedrock_org/lambda/src/sra_sts.py
@@ -1,3 +1,12 @@
+"""Lambda module to use SRA STS service resources in the organization.
+
+Version: 0.1
+
+STS module for SRA in the repo, https://github.com/aws-samples/aws-security-reference-architecture-examples
+
+Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+SPDX-License-Identifier: MIT-0
+"""
 import logging
 import os
 from typing import Any
@@ -7,11 +16,13 @@
 from botocore.config import Config
 import botocore.exceptions
 
-class sra_sts:
+
+class SRASTS:
+    """Class to manage STS resources."""
+
     PROFILE = "default"
 
     UNEXPECTED = "Unexpected!"
-    # TODO(liamschn): this needs to be made into an SSM parameter
     CONFIGURATION_ROLE: str = ""
     BOTO3_CONFIG = Config(retries={"max_attempts": 10, "mode": "standard"})
     PARTITION: str = ""
@@ -22,16 +33,24 @@ class sra_sts:
     log_level: str = os.environ.get("LOG_LEVEL", "INFO")
     LOGGER.setLevel(log_level)
 
-    def __init__(self, profile: str="default") -> None:
+    def __init__(self, profile: str = "default") -> None:
+        """Initialize class object.
+
+        Args:
+            profile (str): AWS credentials profile name. Defaults to "default".
+
+        Raises:
+            ValueError: Error message
+        """
         self.PROFILE = profile
-        print(f"STS PROFILE INFO: {self.PROFILE}")
+        self.LOGGER.info(f"Initial PROFILE: {self.PROFILE}")
 
         try:
             if self.PROFILE != "default":
                 self.MANAGEMENT_ACCOUNT_SESSION = boto3.Session(profile_name=self.PROFILE)
-                print(f"STS INFO: {self.MANAGEMENT_ACCOUNT_SESSION.client('sts').get_caller_identity()}")
+                self.LOGGER.info(f"STS INFO: {self.MANAGEMENT_ACCOUNT_SESSION.client('sts').get_caller_identity()}")
             else:
-                print(f"STS PROFILE AGAIN: {self.PROFILE}")
+                self.LOGGER.info(f"Subsequent PROFILE: {self.PROFILE}")
                 self.MANAGEMENT_ACCOUNT_SESSION = boto3.Session()
 
             self.STS_CLIENT = self.MANAGEMENT_ACCOUNT_SESSION.client("sts")
@@ -48,7 +67,7 @@ def __init__(self, profile: str="default") -> None:
 
             else:
                 self.LOGGER.info(f"Error: {error}")
-                raise error
+                raise ValueError(f"Error: {error}") from None
 
         try:
             self.MANAGEMENT_ACCOUNT = self.STS_CLIENT.get_caller_identity().get("Account")
@@ -59,18 +78,19 @@ def __init__(self, profile: str="default") -> None:
                 self.LOGGER.info("Token has expired, please re-run with proper credentials set.")
             else:
                 self.LOGGER.info(f"Error: {error}")
-                raise error
+                raise ValueError(f"Error: {error}") from None
 
     def assume_role(self, account: str, role_name: str, service: str, region_name: str) -> Any:
         """Get boto3 client assumed into an account for a specified service.
 
         Args:
             account: aws account id
+            role_name: aws role name
             service: aws service
             region_name: aws region
 
         Returns:
-            client: boto3 client
+            Any: boto3 client
         """
         self.LOGGER.info(f"ASSUME ROLE CALLER ID INFO: {self.MANAGEMENT_ACCOUNT_SESSION.client('sts').get_caller_identity()}")
         self.LOGGER.info(f"ASSUME ROLE ACCOUNT (CLIENT): {account}; ROLE NAME: {role_name}; SERVICE: {service}; REGION: {region_name}")
@@ -81,32 +101,29 @@ def assume_role(self, account: str, role_name: str, service: str, region_name: s
                 RoleSessionName="SRA-AssumeCrossAccountRole",
                 DurationSeconds=900,
             )
-            assumed_client = self.MANAGEMENT_ACCOUNT_SESSION.client(
-                service, # type: ignore
+            return self.MANAGEMENT_ACCOUNT_SESSION.client(
+                service,  # type: ignore
                 region_name=region_name,
                 aws_access_key_id=sts_response["Credentials"]["AccessKeyId"],
                 aws_secret_access_key=sts_response["Credentials"]["SecretAccessKey"],
                 aws_session_token=sts_response["Credentials"]["SessionToken"],
             )
-            return assumed_client
-        else:
-            assumed_client = self.MANAGEMENT_ACCOUNT_SESSION.client(
-                service, # type: ignore
-                region_name=region_name, 
-                config=self.BOTO3_CONFIG)
-            return assumed_client
-
+        return self.MANAGEMENT_ACCOUNT_SESSION.client(
+            service,  # type: ignore
+            region_name=region_name,
+            config=self.BOTO3_CONFIG)
 
     def assume_role_resource(self, account: str, role_name: str, service: str, region_name: str) -> Any:
         """Get boto3 resource assumed into an account for a specified service.
 
         Args:
             account: aws account id
+            role_name: aws role name
             service: aws service
             region_name: aws region
 
         Returns:
-            client: boto3 client
+            Any: boto3 client
         """
         self.LOGGER.info(f"ASSUME ROLE CALLER ID INFO: {self.MANAGEMENT_ACCOUNT_SESSION.client('sts').get_caller_identity()}")
         self.LOGGER.info(f"ASSUME ROLE ACCOUNT (RESOURCE): {account}; ROLE NAME: {role_name}; SERVICE: {service}; REGION: {region_name}")
@@ -116,16 +133,23 @@ def assume_role_resource(self, account: str, role_name: str, service: str, regio
             RoleSessionName="SRA-AssumeCrossAccountRole",
             DurationSeconds=900,
         )
-        assumed_resource = self.MANAGEMENT_ACCOUNT_SESSION.resource(
-            service, # type: ignore
+        return self.MANAGEMENT_ACCOUNT_SESSION.resource(
+            service,  # type: ignore
             region_name=region_name,
             aws_access_key_id=sts_response["Credentials"]["AccessKeyId"],
             aws_secret_access_key=sts_response["Credentials"]["SecretAccessKey"],
             aws_session_token=sts_response["Credentials"]["SessionToken"],
         )
-        return assumed_resource
 
     def get_lambda_execution_role(self) -> str:
+        """Get the current lambda execution role arn.
+
+        Raises:
+            ValueError: Unexpected error getting caller identity
+
+        Returns:
+            str: lambda execution role arn
+        """
         try:
             response = self.STS_CLIENT.get_caller_identity()
             return response["Arn"]
