Updates to scenario flow.

Author: rlhagerm

python/example_code/controltower/README.md

@@ -52,16 +52,16 @@ Code examples that show you how to perform the essential operations within a ser
 
 Code excerpts that show you how to call individual service functions.
 
-- [DisableBaseline](controltower_wrapper.py#L365)
-- [DisableControl](controltower_wrapper.py#L240)
-- [EnableBaseline](controltower_wrapper.py#L64)
-- [EnableControl](controltower_wrapper.py#L143)
-- [GetControlOperation](controltower_wrapper.py#L186)
-- [ListBaselines](controltower_wrapper.py#L36)
-- [ListEnabledBaselines](controltower_wrapper.py#L305)
-- [ListEnabledControls](controltower_wrapper.py#L401)
-- [ListLandingZones](controltower_wrapper.py#L278)
-- [ResetEnabledBaseline](controltower_wrapper.py#L332)
+- [DisableBaseline](controltower_wrapper.py#L392)
+- [DisableControl](controltower_wrapper.py#L263)
+- [EnableBaseline](controltower_wrapper.py#L69)
+- [EnableControl](controltower_wrapper.py#L159)
+- [GetControlOperation](controltower_wrapper.py#L209)
+- [ListBaselines](controltower_wrapper.py#L39)
+- [ListEnabledBaselines](controltower_wrapper.py#L330)
+- [ListEnabledControls](controltower_wrapper.py#L431)
+- [ListLandingZones](controltower_wrapper.py#L300)
+- [ResetEnabledBaseline](controltower_wrapper.py#L358)
 
 
 <!--custom.examples.start-->

python/example_code/controltower/controltower_wrapper.py

@@ -4,6 +4,7 @@
 import logging
 import boto3
 import time
+from typing import Dict, List, Optional, Any
 
 from botocore.exceptions import ClientError
 
@@ -17,7 +18,9 @@
 class ControlTowerWrapper:
     """Encapsulates AWS Control Tower and Control Catalog functionality."""
 
-    def __init__(self, controltower_client, controlcatalog_client):
+    def __init__(
+        self, controltower_client: boto3.client, controlcatalog_client: boto3.client
+    ):
         """
         :param controltower_client: A Boto3 Amazon ControlTower client.
         :param controlcatalog_client: A Boto3 Amazon ControlCatalog client.
@@ -66,10 +69,10 @@ def list_baselines(self):
     # snippet-start:[python.example_code.controltower.EnableBaseline]
     def enable_baseline(
         self,
-        target_identifier,
-        identity_center_baseline,
-        baseline_identifier,
-        baseline_version,
+        target_identifier: str,
+        identity_center_baseline: str,
+        baseline_identifier: str,
+        baseline_version: str,
     ):
         """
         Enables a baseline for the specified target if it's not already enabled.
@@ -154,7 +157,7 @@ def list_controls(self):
     # snippet-end:[python.example_code.controltower.ListControls]
 
     # snippet-start:[python.example_code.controltower.EnableControl]
-    def enable_control(self, control_arn, target_identifier):
+    def enable_control(self, control_arn: str, target_identifier: str):
         """
         Enables a control for a specified target.
 
@@ -187,8 +190,11 @@ def enable_control(self, control_arn, target_identifier):
             ):
                 logger.info("Control is already enabled for this target")
                 return None
-            elif (err.response["Error"]["Code"] == "ResourceNotFoundException"
-                    and "not registered with AWS Control Tower" in err.response["Error"]["Message"]):
+            elif (
+                err.response["Error"]["Code"] == "ResourceNotFoundException"
+                and "not registered with AWS Control Tower"
+                in err.response["Error"]["Message"]
+            ):
                 logger.error("Control Tower must be enabled to work with controls.")
                 return None
             logger.error(
@@ -201,7 +207,7 @@ def enable_control(self, control_arn, target_identifier):
     # snippet-end:[python.example_code.controltower.EnableControl]
 
     # snippet-start:[python.example_code.controltower.GetControlOperation]
-    def get_control_operation(self, operation_id):
+    def get_control_operation(self, operation_id: str):
         """
         Gets the status of a control operation.
 
@@ -228,7 +234,7 @@ def get_control_operation(self, operation_id):
     # snippet-end:[python.example_code.controltower.GetControlOperation]
 
     # snippet-start:[python.example_code.controltower.GetBaselineOperation]
-    def get_baseline_operation(self, operation_id):
+    def get_baseline_operation(self, operation_id: str):
         """
         Gets the status of a baseline operation.
 
@@ -255,7 +261,7 @@ def get_baseline_operation(self, operation_id):
     # snippet-end:[python.example_code.controltower.GetBaselineOperation]
 
     # snippet-start:[python.example_code.controltower.DisableControl]
-    def disable_control(self, control_arn, target_identifier):
+    def disable_control(self, control_arn: str, target_identifier: str):
         """
         Disables a control for a specified target.
 
@@ -350,7 +356,7 @@ def list_enabled_baselines(self):
     # snippet-end:[python.example_code.controltower.ListEnabledBaselines]
 
     # snippet-start:[python.example_code.controltower.ResetEnabledBaseline]
-    def reset_enabled_baseline(self, enabled_baseline_identifier):
+    def reset_enabled_baseline(self, enabled_baseline_identifier: str):
         """
         Resets an enabled baseline for a specific target.
 
@@ -384,7 +390,7 @@ def reset_enabled_baseline(self, enabled_baseline_identifier):
     # snippet-end:[python.example_code.controltower.ResetEnabledBaseline]
 
     # snippet-start:[python.example_code.controltower.DisableBaseline]
-    def disable_baseline(self, enabled_baseline_identifier):
+    def disable_baseline(self, enabled_baseline_identifier: str):
         """
         Disables a baseline for a specific target and waits for the operation to complete.
 
@@ -423,7 +429,7 @@ def disable_baseline(self, enabled_baseline_identifier):
     # snippet-end:[python.example_code.controltower.DisableBaseline]
 
     # snippet-start:[python.example_code.controltower.ListEnabledControls]
-    def list_enabled_controls(self, target_identifier):
+    def list_enabled_controls(self, target_identifier: str):
         """
         Lists all enabled controls for a specific target.
 
@@ -445,8 +451,11 @@ def list_enabled_controls(self, target_identifier):
                     "Access denied. Please ensure you have the necessary permissions."
                 )
                 return enabled_controls
-            elif (err.response["Error"]["Code"] == "ResourceNotFoundException"
-                    and "not registered with AWS Control Tower" in err.response["Error"]["Message"]):
+            elif (
+                err.response["Error"]["Code"] == "ResourceNotFoundException"
+                and "not registered with AWS Control Tower"
+                in err.response["Error"]["Message"]
+            ):
                 logger.error("Control Tower must be enabled to work with controls.")
                 return enabled_controls
             else:

python/example_code/controltower/hello/hello_controltower.py

@@ -3,9 +3,10 @@
 
 # snippet-start:[python.example_code.controltower.Hello]
 import boto3
+from typing import Any, List
 
 
-def hello_controltower(controltower_client):
+def hello_controltower(controltower_client: Any) -> None:
     """
     Use the AWS SDK for Python (Boto3) to create an AWS Control Tower client
     and list all available baselines.
@@ -19,7 +20,7 @@ def hello_controltower(controltower_client):
     paginator = controltower_client.get_paginator("list_baselines")
     page_iterator = paginator.paginate()
 
-    baseline_names: [str] = []
+    baseline_names: List[str] = []
     try:
         for page in page_iterator:
             for baseline in page["baselines"]:

python/example_code/controltower/scenario_controltower.py

@@ -5,6 +5,7 @@
 import logging
 import sys
 import time
+from typing import Any, Optional
 
 import boto3
 from botocore.exceptions import ClientError
@@ -20,9 +21,12 @@
 
 # snippet-start:[python.example_code.controltower.ControlTowerScenario]
 class ControlTowerScenario:
+    IDENTITY_CENTER_BASELINE = "baseline/LN25R72TTG6IGPTQ"
     stack_name = ""
 
-    def __init__(self, controltower_wrapper, org_client):
+    def __init__(
+        self, controltower_wrapper: ControlTowerWrapper, org_client: boto3.client
+    ):
         """
         :param controltower_wrapper: An instance of the ControlTowerWrapper class.
         :param org_client: A Boto3 Organization client.
@@ -36,17 +40,17 @@ def __init__(self, controltower_wrapper, org_client):
         self.landing_zone_id = None
         self.use_landing_zone = False
 
-    def run_scenario(self):
+    def run_scenario(self) -> None:
         print("-" * 88)
         print(
             "\tWelcome to the AWS Control Tower with ControlCatalog example scenario."
         )
         print("-" * 88)
 
         print(
-            "This demo will walk you through working with AWS Control Tower for landing zones,"
+            "This demo will walk you through working with AWS Control Tower for landing zones,\n"
+            "managing baselines, and working with controls."
         )
-        print("managing baselines, and working with controls.")
 
         self.account_id = boto3.client("sts").get_caller_identity()["Account"]
 
@@ -102,7 +106,7 @@ def run_scenario(self):
             enabled_baselines = self.controltower_wrapper.list_enabled_baselines()
             for baseline in enabled_baselines:
                 # If the Identity Center baseline is enabled, the identifier must be used for other baselines.
-                if "baseline/LN25R72TTG6IGPTQ" in baseline["baselineIdentifier"]:
+                if self.IDENTITY_CENTER_BASELINE in baseline["baselineIdentifier"]:
                     identity_center_baseline = baseline
                 print(f"{baseline['baselineIdentifier']}")
 
@@ -151,6 +155,15 @@ def run_scenario(self):
                     )
                     print(f"\nDisabled baseline operation id {operation_id}.")
 
+                    # Re-enable the baseline for the next step.
+                    print("\nEnabling Control Tower Baseline.")
+                    self.controltower_wrapper.enable_baseline(
+                        self.ou_arn,
+                        ic_baseline_arn,
+                        control_tower_baseline["arn"],
+                        "4.0",
+                    )
+
         # List and Enable Controls.
         print("\nManaging Controls:")
         controls = self.controltower_wrapper.list_controls()
@@ -186,10 +199,8 @@ def run_scenario(self):
 
                 if operation_id:
                     print(f"Enabled control with operation id {operation_id}")
-                else:
-                    print("Control is already enabled for this target")
 
-            if q.ask(
+            if control_arn and q.ask(
                 f"Do you want to disable the control? (y/n) ",
                 q.is_yesno,
             ):

python/example_code/controltower/test/test_scenario_run.py

@@ -171,6 +171,19 @@ def mock_get_caller_identity():
                 self.baseline_operation_id,
                 "SUCCEEDED",
             )
+            runner.add(
+                self.scenario_data.controltower_stubber.stub_enable_baseline,
+                self.baseline_arn,
+                "4.0",
+                self.sandbox_ou_arn,
+                self.enabled_baseline_arn,
+                self.baseline_operation_id,
+            )
+            runner.add(
+                self.scenario_data.controltower_stubber.stub_get_baseline_operation,
+                self.baseline_operation_id,
+                "SUCCEEDED",
+            )
 
             # List and enable controls
             runner.add(

scenarios/basics/controltower/SPECIFICATION.md

@@ -130,11 +130,19 @@ Baseline operation status: SUCCEEDED
 Reset baseline operation id 64f9c26e-c2d4-46c1-8863-f2b6382c2b4d.
 
 Do you want to disable the Control Tower Baseline? (y/n) y
-Disabling baseline ARN: arn:aws:controltower:us-east-1:478643181688:enabledbaseline/XOVK3ATZUCD5A04QV
-Conflict disabling baseline: AWS Control Tower cannot perform a DisableBaseline operation on a target OU with enabled optional controls.. Skipping disable step.
+Disabling baseline ARN: arn:aws:controltower:us-east-1:478643181688:enabledbaseline/XOVK3ATZUCD5BUH7N
+Baseline operation status: IN_PROGRESS
+Baseline operation status: SUCCEEDED
 
-Disabled baseline operation id None.
+Disabled baseline operation id ff86f9d7-8f08-4bfa-b071-5469c0dfbe60.
 
+Enabling Control Tower Baseline.
+Baseline operation status: IN_PROGRESS
+Baseline operation status: IN_PROGRESS
+Baseline operation status: IN_PROGRESS
+Baseline operation status: IN_PROGRESS
+Baseline operation status: IN_PROGRESS
+Baseline operation status: SUCCEEDED
 
 ```