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