def test_scan_dynamodb_table(self): """ Test that a scan of the DynamoDB table for stuffed animals returns all the items. """ table: Table = self.dynamodb_resource.Table('stuffed-animals') response = table.scan() items: List[dict] = response.get('Items') self.assertEqual(2, len(items)) lily = items[0] self.assertEqual('Bear', lily.get('species')) self.assertEqual( 'A pillow pet who likes to hang out with Dotty, hibernate, and watch Packers games.', lily.get('description')) self.assertEqual('Lily', lily.get('name')) self.assertEqual(Decimal(2), lily.get('id')) self.assertEqual('New York, NY', lily.get('location')) dotty = items[1] self.assertEqual('Horse', dotty.get('species')) self.assertEqual( 'Small spotted horse who loves to cuddle. She also has beautiful flappy ears.', dotty.get('description')) self.assertEqual('Dotty', dotty.get('name')) self.assertEqual(Decimal(1), dotty.get('id')) self.assertEqual('Sleeping under a blanket', dotty.get('favorite_activity'))
def test_populate_na_requirement(self, get_cs_requirements, s3_stubber, special_account): """Test populate with requirement that doesn't apply to account""" cs_data = create_cloudsploit_data({}, [ { 'plugin': 'pluginNameOne', 'category': 'S3', 'title': 'Plugin Name One', 'resource': 'arn:aws:s3:::aaa-aaa-aaa-aaa', 'region': 'global', 'status': 'FAIL', 'statusNumber': 2, 'message': 'sample text' }, ]) get_cs_requirements.return_value = [{ 'requirementId': '10', 'weight': Decimal(100), 'description': 'sample text', 'source': 'cloudsploit', 'severity': 'medium', 'onlyAppliesTo': ['other-than-special'], 'cloudsploit': { 'finding': 'Plugin Name One' } }] s3_stubber.add_response('get_object', create_cloudsploit_s3_response(cs_data), { 'Bucket': os.getenv('CLOUDSPLOIT_RESULT_BUCKET'), 'Key': os.getenv('CLOUDSPLOIT_PREFIX') + '/' + special_account['accountId'] + '/latest.json', }) scan_id = scans_table.create_new_scan_id() cloudsploit_populate({'scanId': scan_id, 'accountId': special_account['accountId']}, {}) # check NCR record is created ncr_records_in_db = ncr_table.query_all(KeyConditionExpression=Key('scanId').eq(scan_id)) assert ncr_records_in_db == [] # check score record is created score_records_in_db = scores_table.query_all(KeyConditionExpression=Key('scanId').eq(scan_id)) for record in score_records_in_db: assert record.pop('ttl') assert score_records_in_db == [ scores_table.new_score( scan_id, special_account['accountId'], { 'requirementId': '10', 'severity': 'medium', 'weight': Decimal(100) }, scores_table.NOT_APPLICABLE, scores_table.NOT_APPLICABLE, ) ]
def __setattr__(self, key, value): if key.startswith('_'): self.__dict__[key] = value else: if isinstance(value, float): value = Decimal(str(value)) self._data[key] = value
def insert_product(product): return table.put_item( Item={ 'vendor': product['vendor'], 'productName': product['name'], 'colors': product['colors'], 'price': Decimal(product['price']), 'category': product['category'] })
def test_add_not_applicable_score(self, get_cs_requirements, special_account): """Tests adding N/A scores to the scores-table for requirements not applying to the account""" scan_id = ScansTable.create_new_scan_id() event = { 'scanId': scan_id, 'accountId': special_account['accountId'], 'error': { 'Cause': json.dumps({'message': 'Test Error'}) }, } get_cs_requirements.return_value = [{ 'requirementId': '10', 'weight': Decimal(100), 'description': 'sample text', 'source': 'cloudsploit', 'severity': 'medium', 'onlyAppliesTo': ['other-special'], 'cloudsploit': { 'finding': 'Plugin Name One' } }] cloudsploit_error(event, self.context) expected_results = { 'accountId': special_account['accountId'], 'score': { 'medium': { 'numFailing': scores_table.NOT_APPLICABLE, 'numResources': scores_table.NOT_APPLICABLE, 'weight': Decimal('100'), } }, 'requirementId': '10', 'scanId': scan_id, 'accntId_rqrmntId': f'{special_account["accountId"]}#10'} result = scores_table.get_item(Key={'scanId': scan_id, 'accntId_rqrmntId': f'{special_account["accountId"]}#10'})['Item'] assert result.pop('ttl') assert result == expected_results
def test_add_dnc_score(self, get_cs_requirements, regular_account): """Tests adding scores for account/requirement""" scan_id = ScansTable.create_new_scan_id() event = { 'scanId': scan_id, 'accountId': regular_account['accountId'], 'error': { 'Cause': json.dumps({'message': 'Test Error'}) }, } get_cs_requirements.return_value = [{ 'requirementId': '10', 'weight': Decimal(100), 'description': 'sample text', 'source': 'cloudsploit', 'severity': 'medium', 'cloudsploit': { 'finding': 'Plugin Name One' } }] cloudsploit_error(event, self.context) expected_results = { 'accountId': regular_account['accountId'], 'score': { 'medium': { 'numFailing': scores_table.DATA_NOT_COLLECTED, 'numResources': scores_table.DATA_NOT_COLLECTED, 'weight': Decimal('100'), } }, 'requirementId': '10', 'scanId': scan_id, 'accntId_rqrmntId': f'{regular_account["accountId"]}#10'} result = scores_table.get_item(Key={'scanId': scan_id, 'accntId_rqrmntId': f'{regular_account["accountId"]}#10'})['Item'] assert result.pop('ttl') assert result == expected_results
def test_query_dynamodb_table(self): """ Test that a query of the DynamoDB table for a stuffed animal by its ID returns the expected record. """ table: Table = self.dynamodb_resource.Table('stuffed-animals') response = table.query(KeyConditionExpression=Key('id').eq(2)) items: List[dict] = response.get('Items') self.assertEqual(1, len(items)) lily = items[0] self.assertEqual('Bear', lily.get('species')) self.assertEqual( 'A pillow pet who likes to hang out with Dotty, hibernate, and watch Packers games.', lily.get('description')) self.assertEqual('Lily', lily.get('name')) self.assertEqual(Decimal(2), lily.get('id')) self.assertEqual('New York, NY', lily.get('location'))
def test_adds_error_to_scan(self, get_cs_requirements, regular_account): scan_id = ScansTable.create_new_scan_id() event = { 'scanId': scan_id, 'accountId': regular_account['accountId'], 'error': { 'Cause': json.dumps({'message': 'Test Error'}) }, } get_cs_requirements.return_value = [{ 'requirementId': '10', 'weight': Decimal(100), 'description': 'sample text', 'source': 'cloudsploit', 'severity': 'medium', 'cloudsploit': { 'finding': 'Plugin Name One' } }] cloudsploit_error(event, self.context) expected_results = { 'scan': ScansTable.SCAN, 'scanId': scan_id, 'errors': [{ 'error': { 'Cause': { 'message': 'Test Error', 'trace': None, }, }, 'functionName': 'function-name', }] } assert scans_table.get_item(Key={'scan': ScansTable.SCAN, 'scanId': scan_id})['Item'] == expected_results
def test_s3importer_single_resource_score(self, get_item, s3_stubber): s3_key = 'sample key' requirement_id = '100' account_ids = ('1', ) event = { 'scanId': scans_table.create_new_scan_id(), 'requirementId': requirement_id, 'accountIds': [account_ids[0]], } s3_import_requirement = create_s3import_requirement( requirement_id, 'critical', s3_key) s3_data = create_s3_data(account_ids, 2, requirement_id) get_item.return_value = {'Item': s3_import_requirement} s3_stubber.add_response('get_object', create_s3_response(s3_data), { 'Bucket': ANY, 'Key': s3_key }) s3import_handler(event, {}) score_records_in_db = scores_table.query( KeyConditionExpression=Key('scanId').eq(event['scanId']))['Items'] for record in score_records_in_db: assert record.pop('ttl') account_id = account_ids[0] scan_id = event['scanId'] assert score_records_in_db == [ scores_table.new_score( scan_id, account_id, s3_import_requirement, Decimal(2), ) ]
EVENT = { 'scanId': SCAN_ID, 'requirementId': '1', 'accountIds': ['111', '222', '333'], 'error': 'test-error' } DESIRED_OUTPUT = [ scores_table.new_score( SCAN_ID, '111', { 'requirementId': '1', 'severity': 'high', 'weight': Decimal(1000), }, scores_table.DATA_NOT_COLLECTED, scores_table.DATA_NOT_COLLECTED, ), scores_table.new_score( SCAN_ID, '222', { 'requirementId': '1', 'severity': 'high', 'weight': Decimal(1000), }, scores_table.DATA_NOT_COLLECTED, scores_table.DATA_NOT_COLLECTED, ),
def object_hook(obj): for k, v in obj.items(): if isinstance(v, float): obj[k] = Decimal(str(v)) return obj
def score_calc_handler(event, context): """ :param event: { scanId: string, accountIds: list of accountIds, } :param context: dict :return: None """ scan_id = event['openScan']['scanId'] account_ids = event['load']['accountIds'] date = scan_id[0:10] all_scores_to_put = [] all_account_scores = [] for account_id in account_ids: account_name = accounts_table.get_account(account_id).get('account_name') scores_to_put = { record['requirementId']: record for record in scores_table.query_all( KeyConditionExpression=Key('scanId').eq(scan_id) & Key('accntId_rqrmntId').begins_with(account_id) ) } existing_ncr_records = ncr_table.query_all( KeyConditionExpression=Key('scanId').eq(scan_id) & Key('accntId_rsrceId_rqrmntId').begins_with(account_id), ) grouped_ncr_data = defaultdict(list) for ncr_object in existing_ncr_records: grouped_ncr_data[ncr_object['requirementId']].append(ncr_object) for requirement_object in all_requirements: severity = requirement_object['severity'] record_to_edit = scores_to_put.get(requirement_object['requirementId'], False) if record_to_edit is False: continue # data not collected for this account for this scan for this requirement, moving on score_object = record_to_edit['score'][severity] # check if score is DNC if so we skip counting failing resources if scores_table.DATA_NOT_COLLECTED in score_object.values(): continue if score_object['numFailing'] is None: score_object['numFailing'] = Decimal(0) matching_ncrs = grouped_ncr_data.get(requirement_object['requirementId'], []) for ncr_record in matching_ncrs: is_excluded = ncr_record.get('exclusionApplied', False) # TODO handle hidden ncr's also (decrement numResources) if is_excluded: continue else: score_object['numFailing'] += 1 all_scores_to_put.append(record_to_edit) account_score = { 'accountId': account_id, 'accountName': account_name, 'date': date, 'scanId': scan_id, 'score': scores_table.weighted_score_aggregate_calc(scores_to_put.values()) } all_account_scores.append(account_score) scores_table.batch_put_records(all_scores_to_put) account_scores_table.batch_put_records(all_account_scores)
'exclusionPendingExpiration': 'sampletext', }, 'allowedActions': { 'updateExclusionReason': True, 'requestExclusion': True, 'addJustification': True, 'requestExclusionExtension': True, 'remediate': True, } }] ACCOUNT_SCORES_DATA = [{ 'accountId': '111111111111', 'date': '2020-03-03', 'score': { 'weightedScore': Decimal(str(0.8)), 'criticalCount': 2 } }, { 'accountId': '222222222222', 'date': '2020-03-03', 'score': { 'weightedScore': Decimal(str(.6)), 'criticalCount': 1 } }, { 'accountId': '111111111111', 'date': '2020-03-01', 'score': { 'weightedScore': Decimal(str(.3)), }
def test_populate(self, get_cs_requirements, s3_stubber, regular_account): cs_data = create_cloudsploit_data({}, [ { 'plugin': 'pluginNameOne', 'category': 'S3', 'title': 'Plugin Name One', 'resource': 'arn:aws:s3:::aaa-aaa-aaa-aaa', 'region': 'global', 'status': 'FAIL', 'statusNumber': 2, 'message': 'sample text' }, { 'plugin': 'pluginNameOne', 'category': 'S3', 'title': 'Plugin Name Two', 'resource': 'arn:aws:s3:::aaa-aaa-aaa-aaa', 'region': 'global', 'status': 'FAIL', 'statusNumber': 2, 'message': 'sample text' }, { 'plugin': 'pluginNameOne', 'category': 'S3', 'title': 'Plugin Name Two', 'resource': 'arn:aws:s3:::aaa-aaa-aaa-aaa', 'region': 'global', 'status': 'FAIL', 'statusNumber': 2, 'message': 'sample text2' }, ]) get_cs_requirements.return_value = [{ 'requirementId': '10', 'weight': Decimal(100), 'description': 'sample text', 'source': 'cloudsploit', 'severity': 'medium', 'service': 'organizations', 'component': 'account', 'cloudsploit': { 'finding': [ 'Plugin Name One', 'Plugin Name Two', ] } }] s3_stubber.add_response('get_object', create_cloudsploit_s3_response(cs_data), { 'Bucket': os.getenv('CLOUDSPLOIT_RESULT_BUCKET'), 'Key': os.getenv('CLOUDSPLOIT_PREFIX') + '/' + regular_account['accountId'] + '/latest.json', }) scan_id = scans_table.create_new_scan_id() cloudsploit_populate({'scanId': scan_id, 'accountId': regular_account['accountId']}, {}) # check NCR record is created ncr_records_in_db = ncr_table.query_all(KeyConditionExpression=Key('scanId').eq(scan_id)) for record in ncr_records_in_db: assert record.pop('ttl') assert ncr_records_in_db == [ ncr_table.new_ncr_record({ 'accountId': regular_account['accountId'], 'accountName': regular_account['account_name'], 'cloudsploitStatus': 'FAIL', 'requirementId': '10', 'resourceId': 'aaa-aaa-aaa-aaa', 'resourceType': 'organizations-account', 'region': 'global', 'reason': 'sample text\nsample text2', }, scan_id) ] # check score record is created score_records_in_db = scores_table.query_all(KeyConditionExpression=Key('scanId').eq(scan_id)) for record in score_records_in_db: assert record.pop('ttl') assert score_records_in_db == [ scores_table.new_score( scan_id, regular_account['accountId'], { 'requirementId': '10', 'severity': 'medium', 'weight': Decimal(100) }, Decimal(3) ) ]