2021-10-05 22:18:39 -07:00
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
2021-11-19 03:43:24 -08:00
2021-10-05 22:18:39 -07:00
"""
Purpose
Shows how to use the AWS SDK for Python (Boto3) with AWS Audit Manager to create an
assessment report that consists of only one day of evidence.
"""
2021-11-19 03:43:24 -08:00
# snippet-start:[python.example_code.auditmanager.Scenario_CreateAssessmentReport]
import dateutil . parser
import logging
2021-10-05 22:18:39 -07:00
import time
import urllib . request
2021-11-19 03:43:24 -08:00
import uuid
2021-10-05 22:18:39 -07:00
import boto3
from botocore . exceptions import ClientError
2021-11-19 03:43:24 -08:00
logger = logging . getLogger ( __name__ )
class AuditReport :
def __init__ ( self , auditmanager_client ) :
self . auditmanager_client = auditmanager_client
def get_input ( self ) :
2023-10-18 10:35:05 -07:00
print ( " - " * 40 )
2021-11-19 03:43:24 -08:00
try :
2023-10-18 10:35:05 -07:00
assessment_id = input ( " Provide assessment id [uuid]: " ) . lower ( )
2021-11-19 03:43:24 -08:00
try :
assessment_uuid = uuid . UUID ( assessment_id )
except ValueError :
logger . error ( " Assessment Id is not a valid UUID: %s " , assessment_id )
raise
2023-10-18 10:35:05 -07:00
evidence_folder = input ( " Provide evidence date [yyyy-mm-dd]: " )
2021-11-19 03:43:24 -08:00
try :
evidence_date = dateutil . parser . parse ( evidence_folder ) . date ( )
except ValueError :
logger . error ( " Invalid date : %s " , evidence_folder )
raise
try :
2023-10-18 10:35:05 -07:00
self . auditmanager_client . get_assessment (
assessmentId = str ( assessment_uuid )
)
2021-11-19 03:43:24 -08:00
except ClientError :
logger . exception ( " Couldn ' t get assessment %s . " , assessment_uuid )
raise
except ( ValueError , ClientError ) :
return None , None
else :
return assessment_uuid , evidence_date
def clear_staging ( self , assessment_uuid , evidence_date ) :
"""
Find all the evidence in the report and clear it.
"""
next_token = None
page = 1
interested_folder_id_list = [ ]
while True :
print ( f " Page [ { page } ] " )
if next_token is None :
2023-10-18 10:35:05 -07:00
folder_list = (
self . auditmanager_client . get_evidence_folders_by_assessment (
assessmentId = str ( assessment_uuid ) , maxResults = 1000
)
)
2021-11-19 03:43:24 -08:00
else :
2023-10-18 10:35:05 -07:00
folder_list = (
self . auditmanager_client . get_evidence_folders_by_assessment (
assessmentId = str ( assessment_uuid ) ,
nextToken = next_token ,
maxResults = 1000 ,
)
)
folders = folder_list . get ( " evidenceFolders " )
2021-11-19 03:43:24 -08:00
print ( f " Got { len ( folders ) } folders. " )
for folder in folders :
2023-10-18 10:35:05 -07:00
folder_id = folder . get ( " id " )
if folder . get ( " name " ) == str ( evidence_date ) :
2021-11-19 03:43:24 -08:00
interested_folder_id_list . append ( folder_id )
2023-10-18 10:35:05 -07:00
if folder . get ( " assessmentReportSelectionCount " ) == folder . get (
" totalEvidence "
) :
2021-11-19 03:43:24 -08:00
print (
f " Removing folder from report selection : { folder . get ( ' name ' ) } "
2023-10-18 10:35:05 -07:00
f " { folder_id } { folder . get ( ' controlId ' ) } "
)
2021-11-19 03:43:24 -08:00
self . auditmanager_client . disassociate_assessment_report_evidence_folder (
2023-10-18 10:35:05 -07:00
assessmentId = str ( assessment_uuid ) , evidenceFolderId = folder_id
)
elif folder . get ( " assessmentReportSelectionCount " ) > 0 :
2021-11-19 03:43:24 -08:00
# Get all evidence in the folder and
# add selected evidence in the selected_evidence_list.
2023-10-18 10:35:05 -07:00
evidence_list = (
self . auditmanager_client . get_evidence_by_evidence_folder (
assessmentId = str ( assessment_uuid ) ,
controlSetId = folder_id ,
evidenceFolderId = folder_id ,
maxResults = 1000 ,
)
)
2021-11-19 03:43:24 -08:00
selected_evidence_list = [ ]
2023-10-18 10:35:05 -07:00
for evidence in evidence_list . get ( " evidence " ) :
if evidence . get ( " assessmentReportSelection " ) == " Yes " :
selected_evidence_list . append ( evidence . get ( " id " ) )
print (
f " Removing evidence report selection : { folder . get ( ' name ' ) } "
f " { len ( selected_evidence_list ) } "
)
2021-11-19 03:43:24 -08:00
self . auditmanager_client . batch_disassociate_assessment_report_evidence (
assessmentId = str ( assessment_uuid ) ,
evidenceFolderId = folder_id ,
2023-10-18 10:35:05 -07:00
evidenceIds = selected_evidence_list ,
)
next_token = folder_list . get ( " nextToken " )
2021-11-19 03:43:24 -08:00
if not next_token :
break
page + = 1
return interested_folder_id_list
def add_folder_to_staging ( self , assessment_uuid , folder_id_list ) :
print ( f " Adding folders to report : { folder_id_list } " )
for folder in folder_id_list :
self . auditmanager_client . associate_assessment_report_evidence_folder (
2023-10-18 10:35:05 -07:00
assessmentId = str ( assessment_uuid ) , evidenceFolderId = folder
)
2021-11-19 03:43:24 -08:00
def get_report ( self , assessment_uuid ) :
report = self . auditmanager_client . create_assessment_report (
2023-10-18 10:35:05 -07:00
name = " ReportViaScript " ,
description = " testing " ,
assessmentId = str ( assessment_uuid ) ,
)
if self . _is_report_generated ( report . get ( " assessmentReport " ) . get ( " id " ) ) :
2021-11-19 03:43:24 -08:00
report_url = self . auditmanager_client . get_assessment_report_url (
2023-10-18 10:35:05 -07:00
assessmentReportId = report . get ( " assessmentReport " ) . get ( " id " ) ,
assessmentId = str ( assessment_uuid ) ,
)
print ( report_url . get ( " preSignedUrl " ) )
2021-11-19 03:43:24 -08:00
urllib . request . urlretrieve (
2023-10-18 10:35:05 -07:00
report_url . get ( " preSignedUrl " ) . get ( " link " ) ,
report_url . get ( " preSignedUrl " ) . get ( " hyperlinkName " ) ,
)
print (
f " Report saved as { report_url . get ( ' preSignedUrl ' ) . get ( ' hyperlinkName ' ) } . "
)
2021-11-19 03:43:24 -08:00
else :
print ( " Report generation did not finish in 15 minutes. " )
2023-10-18 10:35:05 -07:00
print (
" Failed to download report. Go to the console and manually download "
" the report. "
)
2021-11-19 03:43:24 -08:00
def _is_report_generated ( self , assessment_report_id ) :
max_wait_time = 0
while max_wait_time < 900 :
print ( f " Checking status of the report { assessment_report_id } " )
report_list = self . auditmanager_client . list_assessment_reports ( maxResults = 1 )
2023-10-18 10:35:05 -07:00
if (
report_list . get ( " assessmentReports " ) [ 0 ] . get ( " id " )
== assessment_report_id
and report_list . get ( " assessmentReports " ) [ 0 ] . get ( " status " ) == " COMPLETE "
) :
2021-11-19 03:43:24 -08:00
return True
2023-10-18 10:35:05 -07:00
print ( " Sleeping for 5 seconds... " )
2021-11-19 03:43:24 -08:00
time . sleep ( 5 )
max_wait_time + = 5
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 samples demo! " )
2023-10-18 10:35:05 -07:00
print ( " - " * 88 )
print (
" This script creates an assessment report for an assessment with all the "
" evidence collected on the provided date. "
)
print ( " - " * 88 )
report = AuditReport ( boto3 . client ( " auditmanager " ) )
2021-11-19 03:43:24 -08:00
assessment_uuid , evidence_date = report . get_input ( )
if assessment_uuid is not None and evidence_date is not None :
folder_id_list = report . clear_staging ( assessment_uuid , evidence_date )
report . add_folder_to_staging ( assessment_uuid , folder_id_list )
report . get_report ( assessment_uuid )
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_CreateAssessmentReport]
