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)
)
]