def test_same_version_same_rules_new_object(self,
mock_table: mock.MagicMock):
"""The only thing that changes is a new S3 key - update DB entry and alert."""
match_table = analyzer_aws_lib.DynamoMatchTable(MOCK_DYNAMO_TABLE_NAME)
match_table._table.query = lambda **kwargs: {
'Items': [{
'AnalyzerVersion': 1,
'MatchedRules': {'file.yara:rule_name'},
'S3Objects': {'S3:test-bucket:test-key'}
}]
}
self._binary.s3_identifier = 'S3:{}:{}'.format(
self._binary.bucket_name, 'NEW_KEY')
needs_alert = match_table.save_matches(self._binary, 1)
self.assertTrue(needs_alert)
mock_table.assert_has_calls([
mock.call.Table().update_item(
ExpressionAttributeValues={
':s3_string_set': {'S3:test-bucket:NEW_KEY'}
},
Key={
'SHA256': 'Computed_SHA',
'AnalyzerVersion': 1
},
UpdateExpression='ADD S3Objects :s3_string_set')
])
def test_new_version_same_rules_same_objects(self,
mock_table: mock.MagicMock):
"""Same results with new Lambda version - create new DB entry but do not alert."""
match_table = analyzer_aws_lib.DynamoMatchTable(MOCK_DYNAMO_TABLE_NAME)
match_table._table.query = lambda **kwargs: {
'Items': [{
'AnalyzerVersion': 1,
'MatchedRules': {'file.yara:rule_name'},
'S3Objects': {'S3:test-bucket:test-key'}
}]
}
needs_alert = match_table.save_matches(self._binary, 2)
self.assertFalse(needs_alert)
mock_table.assert_has_calls([
mock.call.Table().put_item(
Item={
'SHA256': 'Computed_SHA',
'AnalyzerVersion': 2,
'MatchedRules': {'file.yara:rule_name'},
'MD5': 'Computed_MD5',
'S3LastModified': 'time:right_now',
'S3Metadata': {
'test-filename': 'test.txt'
},
'S3Objects': {'S3:test-bucket:test-key'}
})
])
def test_new_version_new_rules_same_objects(self,
mock_table: mock.MagicMock):
"""A previously analyzed binary matches a new YARA rule - create DB entry and alert."""
match_table = analyzer_aws_lib.DynamoMatchTable(MOCK_DYNAMO_TABLE_NAME)
match_table._table.query = lambda **kwargs: {
'Items': [{
'AnalyzerVersion': 1,
'MatchedRules': {'file.yara:rule_name'},
'S3Objects': {'S3:test-bucket:test-key'}
}]
}
self._binary.yara_matches.append(
YaraMatch('different_rule_name', 'new_file.yara', dict(), set()))
needs_alert = match_table.save_matches(self._binary, 2)
self.assertTrue(needs_alert)
mock_table.assert_has_calls([
mock.call.Table().put_item(
Item={
'SHA256': 'Computed_SHA',
'AnalyzerVersion': 2,
'MatchedRules': {
'new_file.yara:different_rule_name',
'file.yara:rule_name'
},
'MD5': 'Computed_MD5',
'S3LastModified': 'time:right_now',
'S3Metadata': {
'test-filename': 'test.txt'
},
'S3Objects': {'S3:test-bucket:test-key'}
})
])
def test_new_version_multiple_objects(self, mock_table: mock.MagicMock):
"""No alerts should fire for any of multiple binaries which have already matched."""
match_table = analyzer_aws_lib.DynamoMatchTable(MOCK_DYNAMO_TABLE_NAME)
match_table._table.query = lambda **kwargs: {
'Items': [
{
'AnalyzerVersion': 1,
'MatchedRules': {'file.yara:rule_name'},
'S3Objects': {'S3_1', 'S3_2', 'S3_3'}
}
]
}
self._binary.s3_identifier = 'S3_1'
self.assertFalse(match_table.save_matches(self._binary, 2))
match_table._table.query = lambda **kwargs: {
'Items': [
{
'AnalyzerVersion': 2,
'MatchedRules': {'file.yara:rule_name'},
'S3Objects': {'S3_1'}
},
{
'AnalyzerVersion': 1,
'MatchedRules': {'file.yara:rule_name'},
'S3Objects': {'S3_1', 'S3_2', 'S3_3'}
}
]
}
self._binary.s3_identifier = 'S3_2'
self.assertFalse(match_table.save_matches(self._binary, 2))
self._binary.s3_identifier = 'S3_3'
self.assertFalse(match_table.save_matches(self._binary, 2))
mock_table.assert_has_calls([
mock.call.Table().put_item(Item=mock.ANY),
mock.call.Table().update_item(
ExpressionAttributeValues={':s3_string_set': {'S3_2'}},
Key={'SHA256': 'Computed_SHA', 'AnalyzerVersion': 2},
UpdateExpression='ADD S3Objects :s3_string_set'
),
mock.call.Table().update_item(
ExpressionAttributeValues={':s3_string_set': {'S3_3'}},
Key={'SHA256': 'Computed_SHA', 'AnalyzerVersion': 2},
UpdateExpression='ADD S3Objects :s3_string_set'
)
])
def save_matches_and_alert(self, lambda_version, dynamo_table_name, sns_topic_arn):
"""Save match results to Dynamo and publish an alert to SNS if appropriate.
Args:
lambda_version: [int] The currently executing version of the Lambda function.
dynamo_table_name: [string] Save YARA match results to this Dynamo table.
sns_topic_arn: [string] Publish match alerts to this SNS topic ARN.
"""
table = analyzer_aws_lib.DynamoMatchTable(dynamo_table_name)
needs_alert = table.save_matches(self, lambda_version)
# Send alert if appropriate.
if needs_alert:
LOGGER.info('Publishing an SNS alert')
analyzer_aws_lib.publish_alert_to_sns(self, sns_topic_arn)
def test_new_sha(self, mock_table: mock.MagicMock):
"""A binary matches YARA rules for the first time - create DB entry and alert."""
match_table = analyzer_aws_lib.DynamoMatchTable(MOCK_DYNAMO_TABLE_NAME)
match_table._table.query = lambda **kwargs: {}
needs_alert = match_table.save_matches(self._binary, 1)
self.assertTrue(needs_alert)
mock_table.assert_has_calls([
mock.call.Table().put_item(Item={
'SHA256': 'Computed_SHA',
'AnalyzerVersion': 1,
'MatchedRules': {'file.yara:rule_name'},
'MD5': 'Computed_MD5',
'S3LastModified': 'time:right_now',
'S3Metadata': {'test-filename': 'test.txt'},
'S3Objects': {'S3:test-bucket:test-key'}
})
])
def save_matches_and_alert(
self, analyzer_version: int, dynamo_table_name: str, sns_topic_arn: str,
sns_enabled: bool = True) -> None:
"""Save match results to Dynamo and publish an alert to SNS if appropriate.
Args:
analyzer_version: The currently executing version of the Lambda function.
dynamo_table_name: Save YARA match results to this Dynamo table.
sns_topic_arn: Publish match alerts to this SNS topic ARN.
sns_enabled: If True, match alerts are sent to SNS when applicable.
"""
table = analyzer_aws_lib.DynamoMatchTable(dynamo_table_name)
needs_alert = table.save_matches(self, analyzer_version)
# Send alert if appropriate.
if needs_alert and sns_enabled:
LOGGER.info('Publishing a YARA match alert to %s', sns_topic_arn)
subject = '[BiAlert] {} matches a YARA rule'.format(
self.filepath or self.computed_sha)
analyzer_aws_lib.publish_to_sns(self, sns_topic_arn, subject)
def setUp(self):
"""Before each test, create the mock environment."""
# Create a mock Dynamo table.
self._mock_dynamo_client = boto3_mocks.MockDynamoDBClient(
MOCK_DYNAMO_TABLE_NAME, HASH_KEY, RANGE_KEY)
self._mock_dynamo_table = self._mock_dynamo_client.tables[
MOCK_DYNAMO_TABLE_NAME]
# Setup mocks.
boto3.client = mock.MagicMock(return_value=self._mock_dynamo_client)
self._binary = binary_info.BinaryInfo('Bucket', 'Key', None)
self._binary.reported_md5 = 'Original_MD5'
self._binary.observed_path = '/bin/path/run.exe'
self._binary.yara_matches = [
yara_mocks.YaraMatchMock('file.yara', 'rule_name')
]
self._binary.computed_sha = 'Computed_SHA'
self._binary.computed_md5 = 'Computed_MD5'
self._match_table = analyzer_aws_lib.DynamoMatchTable(
MOCK_DYNAMO_TABLE_NAME)
def test_old_version(self, mock_logger: mock.MagicMock,
mock_table: mock.MagicMock):
"""Analyze with an older version of the Lambda function - update DB but do not alert."""
match_table = analyzer_aws_lib.DynamoMatchTable(MOCK_DYNAMO_TABLE_NAME)
match_table._table.query = lambda **kwargs: {
'Items': [{
'AnalyzerVersion': 1,
'MatchedRules': {'file.yara:rule_name'},
'S3Objects': {'S3:test-bucket:test-key'}
}]
}
self._binary.yara_matches.append(
YaraMatch('different_rule_name', 'new_file.yara', dict(), set()))
needs_alert = match_table.save_matches(self._binary, 0)
self.assertFalse(needs_alert) # Don't alert even if there was a change
mock_logger.assert_has_calls([
mock.call.warning(
'Current Lambda version %d is < version %d from previous analysis',
0, 1)
])
mock_table.assert_has_calls(
[mock.call.Table().put_item(Item=mock.ANY)])
