@patch('stream_alert.threat_intel_downloader.main.load_config',
Mock(return_value=CONFIG))
@patch('boto3.client')
@patch('stream_alert.threat_intel_downloader.threat_stream.requests.get',
side_effect=mock_requests_get)
def test_handler_next_token(mock_get, mock_ssm):
"""Threat Intel Downloader - Test handler with next token passed in"""
mock_ssm.return_value = MockSSMClient(suppress_params=True,
parameters=mock_ssm_response())
handler({'next_url': 'next_token'}, get_mock_context())
mock_get.assert_called()
@patch('boto3.client', Mock(return_value=MockLambdaClient()))
def test_invoke_lambda_function():
"""Threat Intel Downloader - Test invoke_lambda_function"""
config = {
'region': 'us-east-1',
'account_id': '123456789012',
'function_name': 'prefix_threat_intel_downloader',
'qualifier': 'development'
}
invoke_lambda_function('next_token', config)
@patch('boto3.client', Mock(return_value=MockLambdaClient()))
@raises(ThreatStreamLambdaInvokeError)
def test_invoke_lambda_function_error():
"""Threat Intel Downloader - Test invoke_lambda_function with error"""
class TestThreatStream(object):
"""Test class to test ThreatStream functionalities"""
@patch('stream_alert.threat_intel_downloader.main.load_config',
Mock(return_value=load_config('tests/unit/conf/')))
def setup(self):
"""Setup TestThreatStream"""
context = get_mock_context(100000)
self.threatstream = ThreatStream(context.invoked_function_arn,
context.get_remaining_time_in_millis)
@staticmethod
def _get_fake_intel(value, source):
return {
'value': value,
'itype': 'c2_domain',
'source': source,
'type': 'domain',
'expiration_ts': '2017-11-30T00:01:02.123Z',
'key1': 'value1',
'key2': 'value2'
}
@staticmethod
def _get_http_response(next_url=None):
return {
'key1':
'value1',
'objects': [
TestThreatStream._get_fake_intel('malicious_domain.com',
'ioc_source'),
TestThreatStream._get_fake_intel('malicious_domain2.com',
'test_source')
],
'meta': {
'next': next_url,
'offset': 100
}
}
@patch('stream_alert.threat_intel_downloader.main.load_config',
Mock(return_value=load_config('tests/unit/conf/')))
def test_load_config(self):
"""ThreatStream - Load Config"""
arn = 'arn:aws:lambda:region:123456789012:function:name:development'
expected_config = {
'account_id':
'123456789012',
'function_name':
'name',
'qualifier':
'development',
'region':
'region',
'enabled':
True,
'excluded_sub_types':
['bot_ip', 'brute_ip', 'scan_ip', 'spam_ip', 'tor_ip'],
'handler':
'main.handler',
'ioc_filters': ['crowdstrike', '@airbnb.com'],
'ioc_keys': ['expiration_ts', 'itype', 'source', 'type', 'value'],
'ioc_types': ['domain', 'ip', 'md5'],
'memory':
'128',
'source_bucket':
'unit-testing.streamalert.source',
'source_current_hash':
'<auto_generated>',
'source_object_key':
'<auto_generated>',
'timeout':
'60'
}
assert_equal(self.threatstream._load_config(arn), expected_config)
def test_process_data(self):
"""ThreatStream - Process Raw IOC Data"""
raw_data = [
self._get_fake_intel('malicious_domain.com', 'ioc_source'),
self._get_fake_intel('malicious_domain2.com', 'ioc_source2'),
# this will get filtered out
self._get_fake_intel('malicious_domain3.com', 'bad_source_ioc'),
]
self.threatstream._config['ioc_filters'] = {'ioc_source'}
processed_data = self.threatstream._process_data(raw_data)
expected_result = [{
'value': 'malicious_domain.com',
'itype': 'c2_domain',
'source': 'ioc_source',
'type': 'domain',
'expiration_ts': 1512000062
}, {
'value': 'malicious_domain2.com',
'itype': 'c2_domain',
'source': 'ioc_source2',
'type': 'domain',
'expiration_ts': 1512000062
}]
assert_equal(processed_data, expected_result)
@patch('boto3.client')
def test_load_api_creds(self, mock_ssm):
"""ThreatStream - Load API creds from SSM"""
params = {
ThreatStream.CRED_PARAMETER_NAME:
'{"api_user": "******", "api_key": "test_key"}'
}
mock_ssm.return_value = MockSSMClient(suppress_params=True,
parameters=params)
self.threatstream._load_api_creds()
assert_equal(self.threatstream.api_user, 'test_user')
assert_equal(self.threatstream.api_key, 'test_key')
@patch('boto3.client')
def test_load_api_creds_cached(self, mock_ssm):
"""ThreatStream - Load API creds from SSM, Cached"""
params = {
ThreatStream.CRED_PARAMETER_NAME:
'{"api_user": "******", "api_key": "test_key"}'
}
mock_ssm.return_value = MockSSMClient(suppress_params=True,
parameters=params)
self.threatstream._load_api_creds()
assert_equal(self.threatstream.api_user, 'test_user')
assert_equal(self.threatstream.api_key, 'test_key')
self.threatstream._load_api_creds()
mock_ssm.assert_called_once(
) # Make sure the client was not loaded again
@patch('boto3.client')
@raises(ClientError)
def test_load_api_creds_client_errors(self, mock_ssm):
"""ThreatStream - Load API creds from SSM, ClientError"""
mock_ssm.return_value = MockSSMClient(suppress_params=True,
parameters={})
self.threatstream._load_api_creds()
@patch('boto3.client')
@raises(ThreatStreamCredsError)
def test_load_api_creds_empty_response(self, mock_ssm):
"""ThreatStream - Load API creds from SSM, Empty Response"""
mock_ssm.return_value.get_parameter.return_value = None
self.threatstream._load_api_creds()
assert_equal(self.threatstream.api_user, 'test_user')
assert_equal(self.threatstream.api_key, 'test_key')
@patch('boto3.client')
@raises(ThreatStreamCredsError)
def test_load_api_creds_invalid_json(self, mock_ssm):
"""ThreatStream - Load API creds from SSM with invalid JSON"""
params = {'threat_intel_downloader_api_creds': 'invalid_value'}
mock_ssm.return_value = MockSSMClient(suppress_params=True,
parameters=params)
self.threatstream._load_api_creds()
@patch('boto3.client')
@raises(ThreatStreamCredsError)
def test_load_api_creds_no_api_key(self, mock_ssm):
"""ThreatStream - Load API creds from SSM, No API Key"""
params = {
'threat_intel_downloader_api_creds':
'{"api_user": "******", "api_key": ""}'
}
mock_ssm.return_value = MockSSMClient(suppress_params=True,
parameters=params)
self.threatstream._load_api_creds()
@patch('stream_alert.threat_intel_downloader.main.datetime')
def test_epoch_now(self, date_mock):
"""ThreatStream - Epoch, Now"""
fake_date_now = datetime(year=2017, month=9, day=1)
date_mock.utcnow.return_value = fake_date_now
date_mock.utcfromtimestamp = datetime.utcfromtimestamp
expected_value = datetime(year=2017, month=11, day=30)
value = self.threatstream._epoch_time(None)
assert_equal(datetime.utcfromtimestamp(value), expected_value)
def test_epoch_from_time(self):
"""ThreatStream - Epoch, From Timestamp"""
expected_value = datetime(year=2017, month=11, day=30)
value = self.threatstream._epoch_time('2017-11-30T00:00:00.000Z')
assert_equal(datetime.utcfromtimestamp(value), expected_value)
@raises(ValueError)
def test_epoch_from_bad_time(self):
"""ThreatStream - Epoch, Error"""
self.threatstream._epoch_time('20171130T00:00:00.000Z')
def test_excluded_sub_types(self):
"""ThreatStream - Excluded Sub Types Property"""
expected_value = ['bot_ip', 'brute_ip', 'scan_ip', 'spam_ip', 'tor_ip']
assert_equal(self.threatstream.excluded_sub_types, expected_value)
def test_ioc_keys(self):
"""ThreatStream - IOC Keys Property"""
expected_value = ['expiration_ts', 'itype', 'source', 'type', 'value']
assert_equal(self.threatstream.ioc_keys, expected_value)
def test_ioc_sources(self):
"""ThreatStream - IOC Sources Property"""
expected_value = ['crowdstrike', '@airbnb.com']
assert_equal(self.threatstream.ioc_sources, expected_value)
def test_ioc_types(self):
"""ThreatStream - IOC Types Property"""
expected_value = ['domain', 'ip', 'md5']
assert_equal(self.threatstream.ioc_types, expected_value)
def test_threshold(self):
"""ThreatStream - Threshold Property"""
assert_equal(self.threatstream.threshold, 499000)
@patch('stream_alert.threat_intel_downloader.main.ThreatStream._finalize')
@patch('stream_alert.threat_intel_downloader.main.requests.get')
def test_connect(self, get_mock, finalize_mock):
"""ThreatStream - Connection to ThreatStream.com"""
get_mock.return_value.json.return_value = self._get_http_response()
get_mock.return_value.status_code = 200
self.threatstream._config['ioc_filters'] = {'test_source'}
self.threatstream._connect('previous_url')
expected_intel = [{
'value': 'malicious_domain2.com',
'itype': 'c2_domain',
'source': 'test_source',
'type': 'domain',
'expiration_ts': 1512000062
}]
finalize_mock.assert_called_with(expected_intel, None)
@patch('stream_alert.threat_intel_downloader.main.ThreatStream._finalize')
@patch('stream_alert.threat_intel_downloader.main.requests.get')
def test_connect_with_next(self, get_mock, finalize_mock):
"""ThreatStream - Connection to ThreatStream.com, with Continuation"""
next_url = 'this_url'
get_mock.return_value.json.return_value = self._get_http_response(
next_url)
get_mock.return_value.status_code = 200
self.threatstream._config['ioc_filters'] = {'test_source'}
self.threatstream._connect('previous_url')
expected_intel = [{
'value': 'malicious_domain2.com',
'itype': 'c2_domain',
'source': 'test_source',
'type': 'domain',
'expiration_ts': 1512000062
}]
finalize_mock.assert_called_with(expected_intel, next_url)
@raises(ThreatStreamRequestsError)
@patch('stream_alert.threat_intel_downloader.main.requests.get')
def test_connect_with_unauthed(self, get_mock):
"""ThreatStream - Connection to ThreatStream.com, Unauthorized Error"""
get_mock.return_value.json.return_value = self._get_http_response()
get_mock.return_value.status_code = 401
self.threatstream._connect('previous_url')
@raises(ThreatStreamRequestsError)
@patch('stream_alert.threat_intel_downloader.main.requests.get')
def test_connect_with_retry_error(self, get_mock):
"""ThreatStream - Connection to ThreatStream.com, Retry Error"""
get_mock.return_value.status_code = 500
self.threatstream._connect('previous_url')
@raises(ThreatStreamRequestsError)
@patch('stream_alert.threat_intel_downloader.main.requests.get')
def test_connect_with_unknown_error(self, get_mock):
"""ThreatStream - Connection to ThreatStream.com, Unknown Error"""
get_mock.return_value.status_code = 404
self.threatstream._connect('previous_url')
@patch(
'stream_alert.threat_intel_downloader.main.ThreatStream._load_api_creds'
)
@patch('stream_alert.threat_intel_downloader.main.ThreatStream._connect')
def test_runner(self, connect_mock, _):
"""ThreatStream - Runner"""
expected_url = (
'/api/v2/intelligence/?username=user&api_key=key&limit=1000&q='
'(status="active")+AND+(type="domain"+OR+type="ip"+OR+type="md5")+'
'AND+NOT+(itype="bot_ip"+OR+itype="brute_ip"+OR+itype="scan_ip"+'
'OR+itype="spam_ip"+OR+itype="tor_ip")')
self.threatstream.api_key = 'key'
self.threatstream.api_user = '******'
self.threatstream.runner({'none': 'test'})
connect_mock.assert_called_with(expected_url)
@patch(
'stream_alert.threat_intel_downloader.main.ThreatStream._write_to_dynamodb_table'
)
@patch(
'stream_alert.threat_intel_downloader.main.ThreatStream._invoke_lambda_function'
)
def test_finalize(self, invoke_mock, write_mock):
"""ThreatStream - Finalize with Intel"""
intel = ['foo', 'bar']
self.threatstream._finalize(intel, None)
write_mock.assert_called_with(intel)
invoke_mock.assert_not_called()
@patch(
'stream_alert.threat_intel_downloader.main.ThreatStream._write_to_dynamodb_table'
)
@patch(
'stream_alert.threat_intel_downloader.main.ThreatStream._invoke_lambda_function'
)
def test_finalize_next_url(self, invoke_mock, write_mock):
"""ThreatStream - Finalize with Next URL"""
intel = ['foo', 'bar']
self.threatstream._finalize(intel, 'next')
write_mock.assert_called_with(intel)
invoke_mock.assert_called_with('next')
@patch('boto3.resource')
def test_write_to_dynamodb_table(self, boto_mock):
"""ThreatStream - Write Intel to DynamoDB Table"""
intel = [self._get_fake_intel('malicious_domain.com', 'test_source')]
expected_intel = {
'expiration_ts': '2017-11-30T00:01:02.123Z',
'source': 'test_source',
'ioc_type': 'domain',
'sub_type': 'c2_domain',
'ioc_value': 'malicious_domain.com'
}
self.threatstream._write_to_dynamodb_table(intel)
batch_writer = boto_mock.return_value.Table.return_value.batch_writer.return_value
batch_writer.__enter__.return_value.put_item.assert_called_with(
Item=expected_intel)
@patch('boto3.resource')
@raises(ClientError)
def test_write_to_dynamodb_table_error(self, boto_mock):
"""ThreatStream - Write Intel to DynamoDB Table, Error"""
intel = [self._get_fake_intel('malicious_domain.com', 'test_source')]
err = ClientError({'Error': {'Code': 404}}, 'PutItem')
batch_writer = boto_mock.return_value.Table.return_value.batch_writer.return_value
batch_writer.__enter__.return_value.put_item.side_effect = err
self.threatstream._write_to_dynamodb_table(intel)
@patch('boto3.client')
def test_invoke_lambda_function(self, boto_mock):
"""ThreatStream - Invoke Lambda Function"""
boto_mock.return_value = MockLambdaClient()
self.threatstream._invoke_lambda_function('next_token')
boto_mock.assert_called_once()
@patch('boto3.client', Mock(return_value=MockLambdaClient()))
@raises(ThreatStreamLambdaInvokeError)
def test_invoke_lambda_function_error(self):
"""ThreatStream - Invoke Lambda Function, Error"""
MockLambdaClient._raise_exception = True
self.threatstream._invoke_lambda_function('next_token')
def test_invoke_lambda_function(self, boto_mock):
"""ThreatStream - Invoke Lambda Function"""
boto_mock.return_value = MockLambdaClient()
self.threatstream._invoke_lambda_function('next_token')
boto_mock.assert_called_once()
class TestAppIntegration(object):
"""Test class for the AppIntegration"""
# Remove all abstractmethods so we can instantiate AppIntegration for testing
# Also patch some abstractproperty attributes
@patch.object(AppIntegration, '__abstractmethods__', frozenset())
@patch('app_integrations.apps.app_base.Batcher', Mock())
def setup(self):
"""Setup before each method"""
self._app = AppIntegration(
AppConfig(get_valid_config_dict('duo_admin'), None))
@patch('logging.Logger.debug')
def test_no_sleep(self, log_mock):
"""App Integration - App Base, No Sleep on First Poll"""
self._app._sleep()
log_mock.assert_called_with('Skipping sleep for first poll')
@patch('time.sleep')
@patch('app_integrations.apps.app_base.AppIntegration._sleep_seconds',
Mock(return_value=1))
def test_sleep(self, time_mock):
"""App Integration - App Base, Sleep"""
self._app._poll_count = 1
self._app._sleep()
time_mock.assert_called_with(1)
def test_validate_auth_(self):
"""App Integration - Validate Authentication Info"""
assert_is_none(self._app._validate_auth())
@raises(AppIntegrationConfigError)
def test_validate_auth_empty(self):
"""App Integration - Validate Authentication Info, No Config Exception"""
self._app._config.clear()
self._app._validate_auth()
@raises(AppIntegrationConfigError)
def test_validate_auth_no_auth(self):
"""App Integration - Validate Authentication Info, No Auth Exception"""
del self._app._config['auth']
self._app._validate_auth()
@raises(AppIntegrationConfigError)
def test_validate_auth_missing_auth(self):
"""App Integration - Validate Authentication Info, Missing Auth Key Exception"""
with patch.object(AppIntegration,
'required_auth_info') as auth_keys_mock:
auth_keys_mock.return_value = {'new_auth_key'}
self._app._validate_auth()
def test_check_http_response_good(self):
"""App Integration - Check HTTP Response, Success"""
response = Mock(status_code=200)
assert_true(self._app._check_http_response(response))
@patch('logging.Logger.error')
def test_check_http_response_bad(self, log_mock):
"""App Integration - Check HTTP Response, Failure"""
response = Mock(status_code=404, content='hey')
# Check to make sure this resulted in a return of False
assert_false(self._app._check_http_response(response))
# Make sure the logger was called with the proper info
log_mock.assert_called_with(
'HTTP request failed for service \'%s\': [%d] %s', 'type', 404,
'hey')
def test_initialize(self):
"""App Integration - Initialize, Valid"""
assert_true(self._app._initialize())
@patch('logging.Logger.error')
def test_initialize_running(self, log_mock):
"""App Integration - Initialize, Already Running"""
self._app._config['current_state'] = 'running'
assert_false(self._app._initialize())
log_mock.assert_called_with('App already running for service \'%s\'.',
'type')
@patch('logging.Logger.error')
def test_initialize_partial(self, log_mock):
"""App Integration - Initialize, Partial Execution"""
self._app._config['current_state'] = 'partial'
assert_false(self._app._initialize())
log_mock.assert_called_with(
'App in partial execution state for service \'%s\'.', 'type')
def test_finalize(self):
"""App Integration - Finalize, Valid"""
test_new_time = 50000000
self._app._last_timestamp = test_new_time
self._app._finalize()
assert_equal(self._app._config.last_timestamp, test_new_time)
@patch('boto3.client', Mock(return_value=MockLambdaClient()))
@patch('app_integrations.config.AppConfig.mark_success')
def test_finalize_more_logs(self, config_mock):
"""App Integration - Finalize, More Logs"""
self._app._more_to_poll = True
self._app._finalize()
config_mock.assert_not_called()
@raises(ClientError)
@patch('boto3.client', Mock(return_value=MockLambdaClient()))
def test_finalize_more_logs_error(self):
"""App Integration - Finalize, More Logs"""
MockLambdaClient._raise_exception = True
self._app._more_to_poll = True
self._app._finalize()
@patch('logging.Logger.error')
def test_finalize_zero_time(self, log_mock):
"""App Integration - Finalize, Zero Time Error"""
self._app._finalize()
log_mock.assert_called_with(
'Ending last timestamp is 0. This should not happen and '
'is likely due to the subclass not setting this value.')
@patch('logging.Logger.error')
def test_finalize_same_time(self, log_mock):
"""App Integration - Finalize, Same Time Error"""
self._app._last_timestamp = self._app._config.start_last_timestamp
self._app._finalize()
log_mock.assert_called_with(
'Ending last timestamp is the same as '
'the beginning last timestamp. This could occur if '
'there were no logs collected for this execution.')
@patch('logging.Logger.info')
def test_gather_success(self, log_mock):
"""App Integration - Gather, Success"""
with patch.object(AppIntegration,
'_gather_logs') as subclass_gather_mock:
subclass_gather_mock.return_value = ['log01', 'log02', 'log03']
self._app._gather()
log_mock.assert_called()
assert_equal(log_mock.call_args_list[-1][0][0],
'Gather process for \'%s\' executed in %f seconds.')
@patch('logging.Logger.error')
def test_gather_no_logs(self, log_mock):
"""App Integration - Gather, No Logs"""
with patch.object(AppIntegration,
'_gather_logs') as subclass_gather_mock:
subclass_gather_mock.return_value = []
self._app._gather()
log_mock.assert_called_with(
'Gather process for service \'%s\' was not able '
'to poll any logs on poll #%d', 'type', 1)
@patch('app_integrations.apps.app_base.AppIntegration._finalize')
@patch('app_integrations.apps.app_base.AppIntegration._sleep_seconds',
Mock(return_value=1))
@patch('app_integrations.config.AppConfig.remaining_ms',
Mock(return_value=5000))
def test_gather_entry(self, finalize_mock):
"""App Integration - Gather, Entry Point"""
self._app.gather()
finalize_mock.assert_called()
@patch('app_integrations.apps.app_base.AppIntegration._gather')
@patch('app_integrations.apps.app_base.AppIntegration._sleep_seconds',
Mock(return_value=1))
@patch('app_integrations.config.AppConfig.remaining_ms',
Mock(side_effect=[8000, 8000, 2000, 2000]))
def test_gather_multiple(self, gather_mock):
"""App Integration - Gather, Entry Point, Multiple Calls"""
# 3 == number of 'seconds' this ran for. This is compared against the remaining_ms mock
gather_mock.side_effect = [3, 3]
self._app._more_to_poll = True
self._app.gather()
assert_equal(gather_mock.call_count, 2)
@patch('app_integrations.apps.app_base.AppIntegration._finalize')
def test_gather_running(self, finalize_mock):
"""App Integration - Gather, Entry Point, Already Running"""
self._app._config['current_state'] = 'running'
self._app.gather()
finalize_mock.assert_not_called()
@patch('requests.get')
def test_make_request_bad_response(self, requests_mock):
"""App Integration - Make Request, Bad Response"""
failed_message = 'something went wrong'
requests_mock.return_value = Mock(
status_code=404,
content=failed_message,
json=Mock(return_value={'message': failed_message}))
result, response = self._app._make_get_request('hostname', None, None)
assert_false(result)
assert_equal(response['message'], failed_message)
# The .json should be called on the response once, to return the response.
assert_equal(requests_mock.return_value.json.call_count, 1)
@patch('requests.get')
def test_make_request_timeout(self, requests_mock):
"""App Integration - Make Request, Timeout"""
requests_mock.side_effect = ConnectTimeout(None, response='too slow')
result, response = self._app._make_get_request('hostname', None, None)
assert_false(result)
assert_is_none(response)