Skip to content

Commit 8d1657c

Browse files
rlhagermrlhagerm
committed
Black formatting.
1 parent d9e15f3 commit 8d1657c

File tree

5 files changed

+214
-163
lines changed

5 files changed

+214
-163
lines changed

‎python/example_code/controltower/controltower_wrapper.py‎

Lines changed: 82 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -42,27 +42,35 @@ def list_baselines(self):
4242
:raises ClientError: If the listing operation fails.
4343
"""
4444
try:
45-
paginator = self.controltower_client.get_paginator('list_baselines')
45+
paginator = self.controltower_client.get_paginator("list_baselines")
4646
baselines = []
4747
for page in paginator.paginate():
48-
baselines.extend(page['baselines'])
48+
baselines.extend(page["baselines"])
4949
return baselines
5050

5151
except ClientError as err:
5252
if err.response["Error"]["Code"] == "AccessDeniedException":
53-
logger.error("Access denied. Please ensure you have the necessary permissions.")
53+
logger.error(
54+
"Access denied. Please ensure you have the necessary permissions."
55+
)
5456
else:
5557
logger.error(
5658
"Couldn't list baselines. Here's why: %s: %s",
5759
err.response["Error"]["Code"],
58-
err.response["Error"]["Message"]
60+
err.response["Error"]["Message"],
5961
)
6062
raise
6163

6264
# snippet-end:[python.example_code.controltower.ListBaselines]
6365

6466
# snippet-start:[python.example_code.controltower.EnableBaseline]
65-
def enable_baseline(self, target_identifier, identity_center_baseline, baseline_identifier, baseline_version):
67+
def enable_baseline(
68+
self,
69+
target_identifier,
70+
identity_center_baseline,
71+
baseline_identifier,
72+
baseline_version,
73+
):
6674
"""
6775
Enables a baseline for the specified target if it's not already enabled.
6876
@@ -81,35 +89,38 @@ def enable_baseline(self, target_identifier, identity_center_baseline, baseline_
8189
parameters=[
8290
{
8391
"key": "IdentityCenterEnabledBaselineArn",
84-
"value": identity_center_baseline
92+
"value": identity_center_baseline,
8593
}
86-
]
94+
],
8795
)
8896

89-
operation_id = response['operationIdentifier']
97+
operation_id = response["operationIdentifier"]
9098
while True:
9199
status = self.get_baseline_operation(operation_id)
92100
print(f"Baseline operation status: {status}")
93-
if status in ['SUCCEEDED', 'FAILED']:
101+
if status in ["SUCCEEDED", "FAILED"]:
94102
break
95103
time.sleep(30)
96104

97-
return response['arn']
105+
return response["arn"]
98106
except ClientError as err:
99107
if err.response["Error"]["Code"] == "ValidationException":
100108
if "already enabled" in err.response["Error"]["Message"]:
101109
print("Baseline is already enabled for this target")
102110
return None
103111
else:
104-
print("Unable to enable baseline due to validation exception: %s: %s",
105-
err.response["Error"]["Code"],
106-
err.response["Error"]["Message"])
112+
print(
113+
"Unable to enable baseline due to validation exception: %s: %s",
114+
err.response["Error"]["Code"],
115+
err.response["Error"]["Message"],
116+
)
107117
logger.error(
108118
"Couldn't enable baseline. Here's why: %s: %s",
109119
err.response["Error"]["Code"],
110-
err.response["Error"]["Message"]
120+
err.response["Error"]["Message"],
111121
)
112122
raise
123+
113124
# snippet-end:[python.example_code.controltower.EnableBaseline]
114125

115126
# snippet-start:[python.example_code.controltower.ListControls]
@@ -121,20 +132,22 @@ def list_controls(self):
121132
:raises ClientError: If the listing operation fails.
122133
"""
123134
try:
124-
paginator = self.controlcatalog_client.get_paginator('list_controls')
135+
paginator = self.controlcatalog_client.get_paginator("list_controls")
125136
controls = []
126137
for page in paginator.paginate():
127-
controls.extend(page['Controls'])
138+
controls.extend(page["Controls"])
128139
return controls
129140

130141
except ClientError as err:
131142
if err.response["Error"]["Code"] == "AccessDeniedException":
132-
logger.error("Access denied. Please ensure you have the necessary permissions.")
143+
logger.error(
144+
"Access denied. Please ensure you have the necessary permissions."
145+
)
133146
else:
134147
logger.error(
135148
"Couldn't list controls. Here's why: %s: %s",
136149
err.response["Error"]["Code"],
137-
err.response["Error"]["Message"]
150+
err.response["Error"]["Message"],
138151
)
139152
raise
140153

@@ -154,30 +167,30 @@ def enable_control(self, control_arn, target_identifier):
154167
print(control_arn)
155168
print(target_identifier)
156169
response = self.controltower_client.enable_control(
157-
controlIdentifier=control_arn,
158-
targetIdentifier=target_identifier
170+
controlIdentifier=control_arn, targetIdentifier=target_identifier
159171
)
160172

161-
operation_id = response['operationIdentifier']
173+
operation_id = response["operationIdentifier"]
162174
while True:
163175
status = self.get_control_operation(operation_id)
164176
print(f"Control operation status: {status}")
165-
if status in ['SUCCEEDED', 'FAILED']:
177+
if status in ["SUCCEEDED", "FAILED"]:
166178
break
167179
time.sleep(30)
168180

169181
return operation_id
170182

171183
except ClientError as err:
172-
if (err.response["Error"]["Code"] == "ValidationException" and
173-
"already enabled" in err.response["Error"][
174-
"Message"]):
184+
if (
185+
err.response["Error"]["Code"] == "ValidationException"
186+
and "already enabled" in err.response["Error"]["Message"]
187+
):
175188
logger.info("Control is already enabled for this target")
176189
return None
177190
logger.error(
178191
"Couldn't enable control. Here's why: %s: %s",
179192
err.response["Error"]["Code"],
180-
err.response["Error"]["Message"]
193+
err.response["Error"]["Message"],
181194
)
182195
raise
183196

@@ -196,15 +209,15 @@ def get_control_operation(self, operation_id):
196209
response = self.controltower_client.get_control_operation(
197210
operationIdentifier=operation_id
198211
)
199-
return response['controlOperation']['status']
212+
return response["controlOperation"]["status"]
200213
except ClientError as err:
201214
if err.response["Error"]["Code"] == "ResourceNotFoundException":
202215
logger.error("Operation not found.")
203216
else:
204217
logger.error(
205218
"Couldn't get control operation status. Here's why: %s: %s",
206219
err.response["Error"]["Code"],
207-
err.response["Error"]["Message"]
220+
err.response["Error"]["Message"],
208221
)
209222
raise
210223

@@ -223,15 +236,15 @@ def get_baseline_operation(self, operation_id):
223236
response = self.controltower_client.get_baseline_operation(
224237
operationIdentifier=operation_id
225238
)
226-
return response['baselineOperation']['status']
239+
return response["baselineOperation"]["status"]
227240
except ClientError as err:
228241
if err.response["Error"]["Code"] == "ResourceNotFoundException":
229242
logger.error("Operation not found.")
230243
else:
231244
logger.error(
232245
"Couldn't get baseline operation status. Here's why: %s: %s",
233246
err.response["Error"]["Code"],
234-
err.response["Error"]["Message"]
247+
err.response["Error"]["Message"],
235248
)
236249
raise
237250

@@ -249,15 +262,14 @@ def disable_control(self, control_arn, target_identifier):
249262
"""
250263
try:
251264
response = self.controltower_client.disable_control(
252-
controlIdentifier=control_arn,
253-
targetIdentifier=target_identifier
265+
controlIdentifier=control_arn, targetIdentifier=target_identifier
254266
)
255267

256-
operation_id = response['operationIdentifier']
268+
operation_id = response["operationIdentifier"]
257269
while True:
258270
status = self.get_control_operation(operation_id)
259271
print(f"Control operation status: {status}")
260-
if status in ['SUCCEEDED', 'FAILED']:
272+
if status in ["SUCCEEDED", "FAILED"]:
261273
break
262274
time.sleep(30)
263275

@@ -269,7 +281,7 @@ def disable_control(self, control_arn, target_identifier):
269281
logger.error(
270282
"Couldn't disable control. Here's why: %s: %s",
271283
err.response["Error"]["Code"],
272-
err.response["Error"]["Message"]
284+
err.response["Error"]["Message"],
273285
)
274286
raise
275287

@@ -284,22 +296,25 @@ def list_landing_zones(self):
284296
:raises ClientError: If the listing operation fails.
285297
"""
286298
try:
287-
paginator = self.controltower_client.get_paginator('list_landing_zones')
299+
paginator = self.controltower_client.get_paginator("list_landing_zones")
288300
landing_zones = []
289301
for page in paginator.paginate():
290-
landing_zones.extend(page['landingZones'])
302+
landing_zones.extend(page["landingZones"])
291303
return landing_zones
292304

293305
except ClientError as err:
294306
if err.response["Error"]["Code"] == "AccessDeniedException":
295-
logger.error("Access denied. Please ensure you have the necessary permissions.")
307+
logger.error(
308+
"Access denied. Please ensure you have the necessary permissions."
309+
)
296310
else:
297311
logger.error(
298312
"Couldn't list landing zones. Here's why: %s: %s",
299313
err.response["Error"]["Code"],
300-
err.response["Error"]["Message"]
314+
err.response["Error"]["Message"],
301315
)
302316
raise
317+
303318
# snippet-end:[python.example_code.controltower.ListLandingZones]
304319

305320
# snippet-start:[python.example_code.controltower.ListEnabledBaselines]
@@ -311,10 +326,10 @@ def list_enabled_baselines(self):
311326
:raises ClientError: If the listing operation fails.
312327
"""
313328
try:
314-
paginator = self.controltower_client.get_paginator('list_enabled_baselines')
329+
paginator = self.controltower_client.get_paginator("list_enabled_baselines")
315330
enabled_baselines = []
316331
for page in paginator.paginate():
317-
enabled_baselines.extend(page['enabledBaselines'])
332+
enabled_baselines.extend(page["enabledBaselines"])
318333
return enabled_baselines
319334

320335
except ClientError as err:
@@ -324,11 +339,12 @@ def list_enabled_baselines(self):
324339
logger.error(
325340
"Couldn't list enabled baselines. Here's why: %s: %s",
326341
err.response["Error"]["Code"],
327-
err.response["Error"]["Message"]
342+
err.response["Error"]["Message"],
328343
)
329344
raise
345+
330346
# snippet-end:[python.example_code.controltower.ListEnabledBaselines]
331-
347+
332348
# snippet-start:[python.example_code.controltower.ResetEnabledBaseline]
333349
def reset_enabled_baseline(self, enabled_baseline_identifier):
334350
"""
@@ -342,11 +358,11 @@ def reset_enabled_baseline(self, enabled_baseline_identifier):
342358
response = self.controltower_client.reset_enabled_baseline(
343359
enabledBaselineIdentifier=enabled_baseline_identifier
344360
)
345-
operation_id = response['operationIdentifier']
361+
operation_id = response["operationIdentifier"]
346362
while True:
347363
status = self.get_baseline_operation(operation_id)
348364
print(f"Baseline operation status: {status}")
349-
if status in ['SUCCEEDED', 'FAILED']:
365+
if status in ["SUCCEEDED", "FAILED"]:
350366
break
351367
time.sleep(30)
352368
return operation_id
@@ -357,11 +373,12 @@ def reset_enabled_baseline(self, enabled_baseline_identifier):
357373
logger.error(
358374
"Couldn't reset enabled baseline. Here's why: %s: %s",
359375
err.response["Error"]["Code"],
360-
err.response["Error"]["Message"]
376+
err.response["Error"]["Message"],
361377
)
362378
raise
379+
363380
# snippet-end:[python.example_code.controltower.ResetEnabledBaseline]
364-
381+
365382
# snippet-start:[python.example_code.controltower.DisableBaseline]
366383
def disable_baseline(self, enabled_baseline_identifier):
367384
"""
@@ -376,28 +393,31 @@ def disable_baseline(self, enabled_baseline_identifier):
376393
enabledBaselineIdentifier=enabled_baseline_identifier
377394
)
378395

379-
operation_id = response['operationIdentifier']
396+
operation_id = response["operationIdentifier"]
380397
while True:
381398
status = self.get_baseline_operation(operation_id)
382399
print(f"Baseline operation status: {status}")
383-
if status in ['SUCCEEDED', 'FAILED']:
400+
if status in ["SUCCEEDED", "FAILED"]:
384401
break
385402
time.sleep(30)
386403

387-
return response['operationIdentifier']
404+
return response["operationIdentifier"]
388405
except ClientError as err:
389406
if err.response["Error"]["Code"] == "ConflictException":
390-
print(f"Conflict disabling baseline: {err.response['Error']['Message']}. Skipping disable step." )
407+
print(
408+
f"Conflict disabling baseline: {err.response['Error']['Message']}. Skipping disable step."
409+
)
391410
return None
392411
else:
393412
logger.error(
394413
"Couldn't disable baseline. Here's why: %s: %s",
395414
err.response["Error"]["Code"],
396-
err.response["Error"]["Message"]
415+
err.response["Error"]["Message"],
397416
)
398417
raise
418+
399419
# snippet-end:[python.example_code.controltower.DisableBaseline]
400-
420+
401421
# snippet-start:[python.example_code.controltower.ListEnabledControls]
402422
def list_enabled_controls(self, target_identifier):
403423
"""
@@ -408,22 +428,26 @@ def list_enabled_controls(self, target_identifier):
408428
:raises ClientError: If the listing operation fails.
409429
"""
410430
try:
411-
paginator = self.controltower_client.get_paginator('list_enabled_controls')
431+
paginator = self.controltower_client.get_paginator("list_enabled_controls")
412432
enabled_controls = []
413433
for page in paginator.paginate(targetIdentifier=target_identifier):
414-
enabled_controls.extend(page['enabledControls'])
434+
enabled_controls.extend(page["enabledControls"])
415435
return enabled_controls
416436

417437
except ClientError as err:
418438
if err.response["Error"]["Code"] == "AccessDeniedException":
419-
logger.error("Access denied. Please ensure you have the necessary permissions.")
439+
logger.error(
440+
"Access denied. Please ensure you have the necessary permissions."
441+
)
420442
else:
421443
logger.error(
422444
"Couldn't list enabled controls. Here's why: %s: %s",
423445
err.response["Error"]["Code"],
424-
err.response["Error"]["Message"]
446+
err.response["Error"]["Message"],
425447
)
426448
raise
449+
427450
# snippet-end:[python.example_code.controltower.ListEnabledControls]
428451

429-
# snippet-end:[python.example_code.controltower.ControlTowerWrapper.class]
452+
453+
# snippet-end:[python.example_code.controltower.ControlTowerWrapper.class]

‎python/example_code/controltower/hello/hello_controltower.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@ def hello_controltower(controltower_client):
2222
baseline_names: [str] = []
2323
try:
2424
for page in page_iterator:
25-
for baseline in page['baselines']:
26-
baseline_names.append(baseline['name'])
25+
for baseline in page["baselines"]:
26+
baseline_names.append(baseline["name"])
2727

2828
print(f"{len(baseline_names)} baseline(s) retrieved.")
2929
for baseline_name in baseline_names:

0 commit comments

Comments
 (0)