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