mypy fixes

aws-samples · Sep 5, 2024 · 74792d7 · 74792d7
1 parent 6d5a0bd
commit 74792d7
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 28 deletions.
diff --git a/aws_sra_examples/solutions/security_lake/security_lake_org/README.md b/aws_sra_examples/solutions/security_lake/security_lake_org/README.md
@@ -66,7 +66,7 @@ AWS SRA Security Lake solution will automate enabling Amazon Security Lake by de
 
 - The python boto3 SDK lambda layer to enable capability for Lambda to enable features of the Security Lake service.
 - This is downloaded during the deployment process and packaged into a layer that is used by the Lambda function in this solution.
-- The Security Lake API available in the current Lambda environment (as of 09/03/2024) is 1.34.145, however, enhanced functionality of the Security Lake API used in this solution requires at least 1.35.10 (see references below).
+- The Security Lake API available in the current Lambda environment (as of 09/03/2024) is 1.20.32, however, enhanced functionality of the Security Lake API used in this solution requires at least 1.35.10 (see references below).
 - Note: Future revisions to this solution will remove this layer when boto3 is updated within the Lambda environment.
 
 #### 1.8 Compliance Event Rule<!-- omit in toc -->
@@ -156,7 +156,7 @@ Choose a Deployment Method:
 In the `management account (home region)`, launch the [sra-security-lake-org-main-ssm.yaml](templates/sra-security-lake-org-main-ssm.yaml) template. This uses an approach where some of the CloudFormation parameters are populated from SSM parameters created by the [SRA Prerequisites Solution](../../common/common_prerequisites/).
 
   ```bash
-  aws cloudformation deploy --template-file $PWD/aws-sra-examples/aws_sra_examples/solutions/security-lake/security-lake/templates/sra-security-lake-org-main-ssm.yaml --stack-name sra-security-lake-org-main-ssm --capabilities CAPABILITY_NAMED_IAM --parameter-overrides pSecurityLakeWarning=<ACCEPT_OR_REJECT>
+  aws cloudformation deploy --template-file $PWD/aws_sra_examples/solutions/security-lake/security-lake-org/templates/sra-security-lake-org-main-ssm.yaml --stack-name sra-security-lake-org-main-ssm --capabilities CAPABILITY_NAMED_IAM --parameter-overrides pSecurityLakeWarning=<ACCEPT_OR_REJECT>
   ```
 
 ##### Important<!-- omit in toc -->

diff --git a/aws_sra_examples/solutions/security_lake/security_lake_org/lambda/src/app.py b/aws_sra_examples/solutions/security_lake/security_lake_org/lambda/src/app.py
@@ -44,7 +44,7 @@
 
 try:
     MANAGEMENT_ACCOUNT_SESSION = boto3.Session()
-    PARTITION: str = MANAGEMENT_ACCOUNT_SESSION.get_partition_for_region(HOME_REGION)
+    PARTITION: str = MANAGEMENT_ACCOUNT_SESSION.get_partition_for_region(HOME_REGION) # type: ignore
     CFN_CLIENT = MANAGEMENT_ACCOUNT_SESSION.client("cloudformation")
 except Exception:
     LOGGER.exception(UNEXPECTED)
@@ -99,15 +99,18 @@ def process_update_event(params: dict, regions: list, accounts: dict) -> None:
     LOGGER.info("...process_update_event")
 
     if params["action"] in ["Update"]:
-        update_security_lake(params, regions)
-        update_log_sources(params, regions, accounts)
-        if params["SET_AUDIT_ACCT_DATA_SUBSCRIBER"]:
-            update_audit_acct_data_subscriber(params, regions)
-        if params["SET_AUDIT_ACCT_QUERY_SUBSCRIBER"]:
-            update_audit_acct_query_subscriber(params, regions)
+        if params["DISABLE_SECURITY_LAKE"]:
+            disable_security_lake(params, regions, accounts)
+        else:
+            update_security_lake(params, regions)
+            update_log_sources(params, regions, accounts)
+            if params["SET_AUDIT_ACCT_DATA_SUBSCRIBER"]:
+                update_audit_acct_data_subscriber(params, regions)
+            if params["SET_AUDIT_ACCT_QUERY_SUBSCRIBER"]:
+                update_audit_acct_query_subscriber(params, regions)
 
-        LOGGER.info("...UPDATE_COMPLETE")
-        return
+            LOGGER.info("...UPDATE_COMPLETE")
+            return
 
     LOGGER.info("...UPDATE_NO_EVENT")
 
@@ -351,19 +354,19 @@ def process_org_configuration(
         source_version: source version
     """
     LOGGER.info(f"Checking if Organization Configuration enabled in {', '.join(regions)} region(s)")
-    org_configuration_exists, exisiting_org_configuration = security_lake.get_org_configuration(sl_client)
+    org_configuration_exists, existing_org_configuration = security_lake.get_org_configuration(sl_client)
     if set_org_configuration:
         sources = [source.strip() for source in org_configuration_sources.split(",")]
         if not org_configuration_exists:
             LOGGER.info(f"Organization Configuration not enabled in {', '.join(regions)} region(s). Creating...")
             security_lake.create_organization_configuration(sl_client, regions, sources, source_version)
             LOGGER.info("Enabled Organization Configuration")
         else:
-            security_lake.update_organization_configuration(sl_client, regions, sources, source_version, exisiting_org_configuration)
+            security_lake.update_organization_configuration(sl_client, regions, sources, source_version, existing_org_configuration)
     else:
         if org_configuration_exists:
             LOGGER.info(f"Deleting Organization Configuration in {r', '.join(regions)} region(s)...")
-            security_lake.delete_organization_configuration(sl_client, exisiting_org_configuration)
+            security_lake.delete_organization_configuration(sl_client, existing_org_configuration)
             LOGGER.info("Deleted Organization Configuration")
 
 
@@ -522,7 +525,7 @@ def add_audit_acct_query_subscriber(sl_client: SecurityLakeClient, params: dict,
 
 
 def configure_audit_acct_for_query_access(params: dict, regions: list) -> None:
-    """Configureresources for query access in Audit account.
+    """Configure resources for query access in Audit account.
 
     Args:
         params: configuration parameters
@@ -591,16 +594,16 @@ def disable_security_lake(params: dict, regions: list, accounts: dict) -> None:
             subscriber_name = params["AUDIT_ACCT_QUERY_SUBSCRIBER"] + "-" + region
             security_lake.delete_subscriber(sl_client, subscriber_name, region)
 
-        org_configuration_exists, exisiting_org_configuration = security_lake.get_org_configuration(sl_client)
+        org_configuration_exists, existing_org_configuration = security_lake.get_org_configuration(sl_client)
         if org_configuration_exists:
-            LOGGER.info(f"Deleting Organization Configuration in {region} region...")
-            security_lake.delete_organization_configuration(sl_client, exisiting_org_configuration)
+            # LOGGER.info(f"Deleting Organization Configuration in {region} region...")
+            # security_lake.delete_organization_configuration(sl_client, existing_org_configuration)
 
     all_accounts = [account["AccountId"] for account in accounts]
     for source in AWS_LOG_SOURCES:
         security_lake.delete_aws_log_source(sl_client, regions, source, all_accounts, params["SOURCE_VERSION"])
 
-    security_lake.delete_security_lake(params["CONFIGURATION_ROLE_NAME"], params["DELEGATED_ADMIN_ACCOUNT_ID"], HOME_REGION, regions)  # todo: remove
+    security_lake.delete_security_lake(params["CONFIGURATION_ROLE_NAME"], params["DELEGATED_ADMIN_ACCOUNT_ID"], HOME_REGION, regions)  # todo: remove after testing
 
 
 def orchestrator(event: dict[str, Any], context: Any) -> None:
@@ -653,7 +656,7 @@ def process_event_cloudformation(event: CloudFormationCustomResourceEvent, conte
     """
     event_info = {"Event": event}
     LOGGER.info(event_info)
-    params = get_validated_parameters(event)
+    params = get_validated_parameters({"RequestType": event["RequestType"]})
     # excluded_accounts: list = [params["DELEGATED_ADMIN_ACCOUNT_ID"]]
     accounts = common.get_active_organization_accounts()
     regions = common.get_enabled_regions(params["ENABLED_REGIONS"], params["CONTROL_TOWER_REGIONS_ONLY"])

diff --git a/aws_sra_examples/solutions/security_lake/security_lake_org/lambda/src/security_lake.py b/aws_sra_examples/solutions/security_lake/security_lake_org/lambda/src/security_lake.py
@@ -12,7 +12,7 @@
 import logging
 import os
 from time import sleep
-from typing import TYPE_CHECKING, List, Literal, Sequence
+from typing import TYPE_CHECKING, List, Literal, Sequence, Union
 
 import boto3
 import botocore
@@ -23,6 +23,7 @@
 if TYPE_CHECKING:
     from mypy_boto3_glue import GlueClient
     from mypy_boto3_lakeformation import LakeFormationClient
+    from mypy_boto3_lakeformation.type_defs import ResourceTypeDef
     from mypy_boto3_organizations import OrganizationsClient
     from mypy_boto3_ram import RAMClient
     from mypy_boto3_securitylake import SecurityLakeClient
@@ -106,7 +107,7 @@ def register_delegated_admin(admin_account_id: str, region: str, service_princip
     Raises:
         ClientError: If there is an issue interacting with the AWS API
     """
-    sl_client: SecurityLakeClient = MANAGEMENT_ACCOUNT_SESSION.client("securitylake", region, config=BOTO3_CONFIG)
+    sl_client: SecurityLakeClient = MANAGEMENT_ACCOUNT_SESSION.client("securitylake", region, config=BOTO3_CONFIG)  # type: ignore
     if not check_organization_admin_enabled(admin_account_id, service_principal):
         LOGGER.info(f"Registering delegated administrator ({admin_account_id})...")
         sl_client.register_data_lake_delegated_administrator(accountId=admin_account_id)
@@ -917,13 +918,14 @@ def set_lake_formation_permissions(lf_client: LakeFormationClient, account: str,
     """
     LOGGER.info("Setting lakeformation permissions for db")
     try:
+        resource: Union[ResourceTypeDef] = {
+                "Database": {"CatalogId": account, "Name": db_name + "_subscriber"},
+                "Table": {"CatalogId": account, "DatabaseName": db_name + "_subscriber", "Name": "rl_*"},
+            }
         lf_client.grant_permissions(
             CatalogId=account,
             Principal={"DataLakePrincipalIdentifier": f"arn:aws:iam::{account}:role/sra-security-lake-query-subscriber"},
-            Resource={
-                "Database": {"CatalogId": account, "Name": db_name + "_subscriber"},
-                "Table": {"CatalogId": account, "DatabaseName": db_name + "_subscriber", "Name": "rl_*"},
-            },
+            Resource=resource,
             Permissions=["ALL"],
             PermissionsWithGrantOption=["ALL"],
         )

diff --git a/...s/solutions/security_lake/security_lake_org/templates/sra-security-lake-org-main-ssm.yaml b/...s/solutions/security_lake/security_lake_org/templates/sra-security-lake-org-main-ssm.yaml
@@ -175,7 +175,7 @@ Parameters:
     Type: String
   pCreateResourceLink:
     AllowedValues: ['true', 'false']
-    Default: 'true'
+    Default: 'false'
     Description: Indicates whether to create a resource link for shared resources in Audit (Security Tooling) account
     Type: String
   pCreateLakeFormationSlr:
@@ -292,7 +292,7 @@ Parameters:
     Type: String
   pControlTowerRegionsOnly:
     AllowedValues: ['true', 'false']
-    Default: 'false'
+    Default: 'true'
     Description: Only enable in the customer governed regions specified in Control Tower or Common Prerequisites solution
     Type: String
   pCreateLambdaLogGroup: