def test_write_to_dynamodb_table(self, mock_get, mock_boto3_resource,
mock_ssm): # pylint: disable=unused-argument
"""ThreatStream - Test write action to dynamodb table"""
mock_ssm.return_value = MockSSMClient(suppress_params=True,
parameters=mock_ssm_response())
threat_stream = ThreatStream(mock_config())
threat_stream.ioc_sources = set(['test_source'])
intelligence, _, _ = threat_stream.runner({'next_url': 'next_url'})
threat_stream.write_to_dynamodb_table(intelligence)
calls = [
call('dynamodb', region_name='us-east-1'),
call().Table('prefix_threat_intel_downloader'),
call().Table().batch_writer(),
call().Table().batch_writer().__enter__(),
call().Table().batch_writer().__enter__().put_item(
Item={
'expiration_ts': 1512000062,
'source': 'test_source',
'type': 'domain',
'sub_type': 'c2_domain',
'value': 'malicious_domain2.com'
}),
call().Table().batch_writer().__exit__(None, None, None)
]
mock_boto3_resource.assert_has_calls(calls)
def test_runner(self, mock_get, mock_ssm): # pylint: disable=unused-argument
"""ThreatStream - Test connection to threatstream"""
mock_ssm.return_value = MockSSMClient(suppress_params=True,
parameters=mock_ssm_response())
threat_stream = ThreatStream(mock_config())
threat_stream.ioc_sources = set(['ioc_source'])
intelligence, next_url, continue_invoke = threat_stream.runner(None)
assert_equal(intelligence, None)
assert_equal(next_url, None)
assert_equal(continue_invoke, False)
intelligence, next_url, continue_invoke = threat_stream.runner(
{'foo': 'bar'})
assert_true(isinstance(intelligence, list))
assert_equal(len(intelligence), 1)
assert_is_not_none(next_url)
assert_equal(continue_invoke, False)
intelligence, next_url, continue_invoke = threat_stream.runner(
{'next_url': 'next_url'})
assert_true(isinstance(intelligence, list))
assert_equal(len(intelligence), 1)
assert_equal(next_url, 'next_url')
assert_equal(continue_invoke, False)
def handler(event, context):
"""Lambda handler"""
config = load_config()
config.update(parse_lambda_func_arn(context))
threat_stream = ThreatStream(config)
intelligence, next_url, continue_invoke = threat_stream.runner(event)
if intelligence:
LOGGER.info('Write %d IOCs to DynamoDB table', len(intelligence))
threat_stream.write_to_dynamodb_table(intelligence)
if context.get_remaining_time_in_millis() > END_TIME_BUFFER * 1000 and continue_invoke:
invoke_lambda_function(next_url, config)
LOGGER.debug("Time remaining (MS): %s", context.get_remaining_time_in_millis())