2021-10-05 22:18:39 -07:00
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
"""
Purpose
Shows how to use the AWS SDK for Python (Boto3) with AWS Audit Manager to create a
custom framework with all standard controls using AWS Security Hub as their data source.
"""
2021-11-19 03:43:24 -08:00
# snippet-start:[python.example_code.auditmanager.Scenario_CustomFrameworkFromSecurityHub]
2021-10-05 22:18:39 -07:00
import logging
import boto3
from botocore . exceptions import ClientError
logger = logging . getLogger ( __name__ )
2021-11-19 03:43:24 -08:00
class SecurityHub :
def __init__ ( self , auditmanager_client ) :
self . auditmanager_client = auditmanager_client
2021-10-05 22:18:39 -07:00
2021-11-19 03:43:24 -08:00
def get_sechub_controls ( self ) :
"""
Gets the list of controls that use Security Hub as their data source.
2021-10-05 22:18:39 -07:00
2021-11-19 03:43:24 -08:00
:return: The list of Security Hub controls.
"""
2023-10-18 10:35:05 -07:00
print ( " - " * 88 )
2021-11-19 03:43:24 -08:00
next_token = None
page = 1
sechub_control_list = [ ]
while True :
print ( " Page [ " + str ( page ) + " ] " )
if next_token is None :
control_list = self . auditmanager_client . list_controls (
2023-10-18 10:35:05 -07:00
controlType = " Standard " , maxResults = 100
)
2021-11-19 03:43:24 -08:00
else :
control_list = self . auditmanager_client . list_controls (
2023-10-18 10:35:05 -07:00
controlType = " Standard " , nextToken = next_token , maxResults = 100
)
print ( " Total controls found: " , len ( control_list . get ( " controlMetadataList " ) ) )
for control in control_list . get ( " controlMetadataList " ) :
2021-11-19 03:43:24 -08:00
control_details = self . auditmanager_client . get_control (
2023-10-18 10:35:05 -07:00
controlId = control . get ( " id " )
) . get ( " control " , { } )
if " AWS Security Hub " in control_details . get ( " controlSources " ) :
sechub_control_list . append ( { " id " : control_details . get ( " id " ) } )
next_token = control_list . get ( " nextToken " )
2021-11-19 03:43:24 -08:00
if not next_token :
break
page + = 1
2023-10-18 10:35:05 -07:00
print ( " Number of Security Hub controls found: " , len ( sechub_control_list ) )
2021-11-19 03:43:24 -08:00
return sechub_control_list
2021-10-05 22:18:39 -07:00
2021-11-19 03:43:24 -08:00
def create_custom_framework ( self , am_controls ) :
"""
Create a custom framework with a list of controls.
2021-10-05 22:18:39 -07:00
2021-11-19 03:43:24 -08:00
:param am_controls: The list of controls to include in the framework.
"""
try :
2023-10-18 10:35:05 -07:00
print ( " Creating custom framework... " )
2021-11-19 03:43:24 -08:00
custom_framework = self . auditmanager_client . create_assessment_framework (
2023-10-18 10:35:05 -07:00
name = " All Security Hub Controls Framework " ,
controlSets = [ { " name " : " Security-Hub " , " controls " : am_controls } ] ,
)
print (
f " Successfully created the custom framework: "
f " { custom_framework . get ( ' framework ' ) . get ( ' name ' ) } : "
f " { custom_framework . get ( ' framework ' ) . get ( ' id ' ) } "
)
print ( " - " * 88 )
2021-11-19 03:43:24 -08:00
except ClientError :
logger . exception ( " Failed to create custom framework. " )
raise
def run_demo ( ) :
2023-10-18 10:35:05 -07:00
print ( " - " * 88 )
2021-11-19 03:43:24 -08:00
print ( " Welcome to the AWS Audit Manager Security Hub demo! " )
2023-10-18 10:35:05 -07:00
print ( " - " * 88 )
2021-10-05 22:18:39 -07:00
print ( " This script creates a custom framework with all Security Hub controls. " )
2023-10-18 10:35:05 -07:00
print ( " - " * 88 )
sechub = SecurityHub ( boto3 . client ( " auditmanager " ) )
2021-11-19 03:43:24 -08:00
am_controls = sechub . get_sechub_controls ( )
sechub . create_custom_framework ( am_controls )
2023-10-18 10:35:05 -07:00
if __name__ == " __main__ " :
2021-11-19 03:43:24 -08:00
run_demo ( )
# snippet-end:[python.example_code.auditmanager.Scenario_CustomFrameworkFromSecurityHub]
