removed prints, updated contact logic

aws-samples · Feb 1, 2024 · 7318195 · 7318195
1 parent 5451ab1
commit 7318195
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 19 deletions.
diff --git a/aws_sra_examples/solutions/shield_advanced/shield_advanced/lambda/src/app.py b/aws_sra_examples/solutions/shield_advanced/shield_advanced/lambda/src/app.py
@@ -9,7 +9,6 @@
 """
 from __future__ import annotations
 
-import json
 import logging
 import os
 import re
@@ -521,15 +520,12 @@ def setup_shield_global(params: dict, accounts: list) -> None:
     if params["SHIELD_ACCOUNTS_TO_PROTECT"] == "ALL":
         LOGGER.info("Protect all accounts")
     else:
-        LOGGER.info("")
         accounts = []
-        print(f'SHIELD_ACCOUNTS_TO_PROTECT {params["SHIELD_ACCOUNTS_TO_PROTECT"]}')
         for account in params["SHIELD_ACCOUNTS_TO_PROTECT"].split(","):
-            print(f"Adding AccountId: {account} to accounts")
             accounts.append({"AccountId": account})
     for account in accounts:
         account_id = account["AccountId"]
-        print(account_id)
+        LOGGER.info(f"Configuring account {account_id}")
         shield.check_if_key_in_object(account_id, shield.RESOURCES_BY_ACCOUNT, "dict")
 
         account_session: boto3.Session = common.assume_role(params["CONFIGURATION_ROLE_NAME"], "sra-configure-shield", account_id)
@@ -554,7 +550,7 @@ def teardown_shield(account_session: boto3.Session, account_id: str, params: dic
     LOGGER.info(f"Teardown shield in for account {account_id} in ")
     shield.build_resources_by_account(account_session, params, account_id)
     shield_client = account_session.client("shield")
-    shield.disable_proactive_engagement(shield_client)  #
+    shield.disable_proactive_engagement(shield_client)
 
     while len(shield.RESOURCES_BY_ACCOUNT[account_id]["buckets"]) > 0:
         bucket = shield.RESOURCES_BY_ACCOUNT[account_id]["buckets"].pop()
@@ -563,11 +559,8 @@ def teardown_shield(account_session: boto3.Session, account_id: str, params: dic
     while len(shield.RESOURCES_BY_ACCOUNT[account_id]["resources_to_protect"]) > 0:
         resource = shield.RESOURCES_BY_ACCOUNT[account_id]["resources_to_protect"].pop()
         if resource not in resources_processed:
-            # if "::" in resource or region in resource and resource:
             shield.delete_protection(shield_client, resource)
             resources_processed.append(resource)
-            # else:
-            #     shield.RESOURCES_BY_ACCOUNT[account_id]["resources_to_protect"].append(resource)
     shield.delete_protection_group(shield_client, params, account_id)
     shield.update_emergency_contacts(shield_client, params, True)
 

diff --git a/aws_sra_examples/solutions/shield_advanced/shield_advanced/lambda/src/shield.py b/aws_sra_examples/solutions/shield_advanced/shield_advanced/lambda/src/shield.py
@@ -15,7 +15,6 @@
 from typing import TYPE_CHECKING, Any, Literal, Sequence
 
 import boto3
-import common
 from botocore.exceptions import ClientError
 
 if TYPE_CHECKING:
@@ -51,9 +50,6 @@
 
 
 UNEXPECTED = "Unexpected!"
-# shield_THROTTLE_PERIOD: float = 0.2
-ENABLE_RETRY_ATTEMPTS: int = 10
-ENABLE_RETRY_SLEEP_INTERVAL: int = 10
 RESOURCES_BY_ACCOUNT: dict = {}
 
 try:
@@ -644,7 +640,6 @@ def create_protection_group(shield_client: ShieldClient, params: dict, account_i
         account_id: AWS account id
     """
     for i in range(0, 5):
-        print(i)
         pg_id: str = params[f"PROTECTION_GROUP_{i}_ID"]
         pg_account_id: str = params[f"PROTECTION_GROUP_{i}_ACCOUNT_ID"]
         pg_aggregation: Literal["SUM", "MEAN", "MAX"] = params[f"PROTECTION_GROUP_{i}_AGGREGATION"]
@@ -657,10 +652,7 @@ def create_protection_group(shield_client: ShieldClient, params: dict, account_i
             "APPLICATION_LOAD_BALANCER",
             "GLOBAL_ACCELERATOR",
         ] = params[f"PROTECTION_GROUP_{i}_RESOURCE_TYPE"]
-        print(f" pg_account_id {pg_account_id}")
         pg_members: list = params[f"PROTECTION_GROUP_{i}_MEMBERS"]
-        print(f"pg_members{pg_members}")
-        print(f"i {i}")
         if pg_id != "" and pg_account_id == account_id:
             if check_if_protection_group_exists(shield_client, pg_id):
                 LOGGER.info(f"Protection_Group_{i} already exists in {account_id}")
@@ -695,7 +687,10 @@ def check_emergency_contacts(shield_client: ShieldClient) -> bool:
         emergency_contacts_response: DescribeEmergencyContactSettingsResponseTypeDef = shield_client.describe_emergency_contact_settings()
         api_call_details = {"API_Call": "shield:DescribeEmergencyContactSettings", "API_Response": emergency_contacts_response}
         LOGGER.info(api_call_details)
-        return True
+        if len(emergency_contacts_response) > 0:
+            return True
+        else:
+            return False
     except shield_client.exceptions.ResourceNotFoundException:
         return False
 
@@ -707,7 +702,6 @@ def enable_proactive_engagement(shield_client: ShieldClient, params: dict) -> No
         shield_client: shield client
         params: environment variables
     """
-    print(f"Before IF SHIELD_ENABLE_PROACTIVE_ENGAGEMENT is set to {params['SHIELD_ENABLE_PROACTIVE_ENGAGEMENT']}")
     if params["SHIELD_ENABLE_PROACTIVE_ENGAGEMENT"] == "true":
         if check_proactive_engagement_enabled(shield_client, params):
             update_emergency_contacts(shield_client, params)