Updates for v5 baselines

Author: rlhagerm

python/example_code/controltower/controltower_wrapper.py

@@ -85,16 +85,21 @@ def enable_baseline(
         :raises ClientError: If enabling the baseline fails for reasons other than it being already enabled.
         """
         try:
-            response = self.controltower_client.enable_baseline(
-                baselineIdentifier=baseline_identifier,
-                baselineVersion=baseline_version,
-                targetIdentifier=target_identifier,
-                parameters=[
+            # Only include parameters if identity_center_baseline is not empty
+            parameters = []
+            if identity_center_baseline:
+                parameters = [
                     {
                         "key": "IdentityCenterEnabledBaselineArn",
                         "value": identity_center_baseline,
                     }
-                ],
+                ]
+            
+            response = self.controltower_client.enable_baseline(
+                baselineIdentifier=baseline_identifier,
+                baselineVersion=baseline_version,
+                targetIdentifier=target_identifier,
+                parameters=parameters,
             )
 
             operation_id = response["operationIdentifier"]
@@ -110,7 +115,6 @@ def enable_baseline(
             if err.response["Error"]["Code"] == "ValidationException":
                 if "already enabled" in err.response["Error"]["Message"]:
                     print("Baseline is already enabled for this target")
-                    return None
                 else:
                     print(
                         "Unable to enable baseline due to validation exception: %s: %s",
@@ -122,7 +126,7 @@ def enable_baseline(
                 err.response["Error"]["Code"],
                 err.response["Error"]["Message"],
             )
-            raise
+            return None
 
     # snippet-end:[python.example_code.controltower.EnableBaseline]
 

python/example_code/controltower/scenario_controltower.py

@@ -121,7 +121,7 @@ def run_scenario(self) -> None:
                     else None
                 )
                 baseline_arn = self.controltower_wrapper.enable_baseline(
-                    self.ou_arn, ic_baseline_arn, control_tower_baseline["arn"], "4.0"
+                    self.ou_arn, ic_baseline_arn, control_tower_baseline["arn"], "5.0"
                 )
                 if baseline_arn:
                     print(f"Enabled baseline ARN: {baseline_arn}")
@@ -133,9 +133,10 @@ def run_scenario(self) -> None:
                             == control_tower_baseline["arn"]
                         ):
                             baseline_arn = enabled_baseline["arn"]
-                    print("No change, the selected baseline was already enabled.")
+                    if baseline_arn:
+                        print("No change, the selected baseline was already enabled.")
 
-                if q.ask(
+                if baseline_arn and q.ask(
                     f"Do you want to reset the Control Tower Baseline? (y/n) ",
                     q.is_yesno,
                 ):
@@ -161,7 +162,7 @@ def run_scenario(self) -> None:
                         self.ou_arn,
                         ic_baseline_arn,
                         control_tower_baseline["arn"],
-                        "4.0",
+                        "5.0",
                     )
 
         # List and Enable Controls.