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())
def test_get_api_creds(self, mock_ssm): """ThreatStream - Test get api creds from SSM""" mock_ssm.return_value = MockSSMClient(suppress_params=True, parameters=mock_ssm_response()) threat_stream = ThreatStream(mock_config()) assert_equal(threat_stream.api_user, 'test_user') assert_equal(threat_stream.api_key, 'test_key')
def test_connect(self, mock_get, mock_ssm): # pylint: disable=unused-argument """ThreatStream - Test connection to ThreatStream.com""" 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, next_url, continue_invoke = threat_stream._connect( 'next_token') expected_intel = [{ 'value': 'malicious_domain2.com', 'itype': 'c2_domain', 'source': 'test_source', 'type': 'domain', 'expiration_ts': 1512000062 }] assert_equal(intelligence, expected_intel) assert_equal(next_url, 'next_token') assert_false(continue_invoke)
def test_process_data(self, mock_ssm): """ThreatStream - Test raw ioc data is processed correctly""" raw_data = [{ 'value': 'malicious_domain.com', 'itype': 'c2_domain', 'source': 'ioc_source', 'type': 'domain', 'expiration_ts': '2017-12-31T00:01:02.123Z', 'key1': 'value1', 'key2': 'value2' }, { 'value': 'malicious_domain2.com', 'itype': 'c2_domain', 'source': 'ioc_source2', 'type': 'domain', 'expiration_ts': '2017-11-30T00:01:02.123Z', 'key3': 'value3', 'key4': 'value4' }] mock_ssm.return_value = MockSSMClient(suppress_params=True, parameters=mock_ssm_response()) threat_stream = ThreatStream(mock_config()) threat_stream.ioc_sources = set(['ioc_source']) processed_data = threat_stream._process_data(raw_data) assert_equal(len(processed_data), 2) expected_result = { 'value': 'malicious_domain.com', 'itype': 'c2_domain', 'source': 'ioc_source', 'type': 'domain', 'expiration_ts': 1514678462 } assert_equal(processed_data[0], expected_result) expected_result = { 'value': 'malicious_domain2.com', 'itype': 'c2_domain', 'source': 'ioc_source2', 'type': 'domain', 'expiration_ts': 1512000062 } assert_equal(processed_data[1], expected_result)
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 test_get_api_creds_invalid_params(self, mock_ssm): """ThreatStream - Test get api creds from SSM with wrong parameters""" mock_ssm.return_value = MockSSMClient( suppress_params=True, parameters=mock_invalid_ssm_response()) ThreatStream(mock_config())
def test_get_api_creds_client_errors(self, mock_ssm): """ThreatStream - Test get api creds from SSM with client exception""" mock_ssm.return_value = MockSSMClient(suppress_params=False, raise_exception=True) ThreatStream(mock_config())
def test_get_api_creds_params_errors(self, mock_ssm): """ThreatStream - Test get api creds from SSM with wrong parameters""" mock_ssm.return_value = MockSSMClient(suppress_params=True) ThreatStream(mock_config())