def boto_volume_for_test(test, cluster_id): """ Create an in-memory boto3 Volume, avoiding any AWS API calls. """ # Create a session directly rather than allow lazy loading of a default # session. region_name = u"some-test-region-1" s = Boto3Session( botocore_session=botocore_get_session(), region_name=region_name, ) ec2 = s.resource("ec2", region_name=region_name) stubber = Stubber(ec2.meta.client) # From this point, any attempt to interact with AWS API should fail with # botocore.exceptions.StubResponseError stubber.activate() volume_id = u"vol-{}".format(random_name(test)) v = ec2.Volume(id=volume_id) tags = [] if cluster_id is not None: tags.append( dict( Key=CLUSTER_ID_LABEL, Value=cluster_id, ), ) # Pre-populate the metadata to prevent any attempt to load the metadata by # API calls. v.meta.data = dict( Tags=tags ) return v
def create_client_sts_stub(service, *args, **kwargs): client = _original_create_client(service, *args, **kwargs) stub = Stubber(client) response = self.create_assume_role_response(expected_creds) self.actual_client_region = client.meta.region_name stub.add_response('assume_role', response) stub.activate() return client
def test_delete_tags(self): stubber = Stubber(self.instance_resource.meta.client) stubber.add_response('delete_tags', {}) stubber.activate() response = self.instance_resource.delete_tags(Tags=[{'Key': 'foo'}]) stubber.assert_no_pending_responses() self.assertEqual(response, {}) stubber.deactivate()
def set_cloudformation_stubber_for_client(self, redshift_client): stubber = Stubber(redshift_client) with open(self.resource_path+'/DescribeSourceClusterResponse.json') as describe_source_response: describe_source_cluster_response = json.load(describe_source_response) expected_source_params = {'ClusterIdentifier': 'rscopyunloadtest3-redshiftclustersource-1so4t2ip0ei3a'} stubber.add_response('describe_clusters', describe_source_cluster_response, expected_source_params) with open(self.resource_path+'/DescribeTargetClusterResponse.json') as describe_target_response: describe_target_cluster_response = json.load(describe_target_response) expected_target_params = {'ClusterIdentifier': 'rscopyunloadtest3-redshiftclustertarget-oaw35zvu02h'} stubber.add_response('describe_clusters', describe_target_cluster_response, expected_target_params) stubber.activate()
def test_multipart_download_with_multiple_parts_and_extra_args(self): client = Session().create_client('s3') stubber = Stubber(client) response_body = b'foobarbaz' response = {'Body': six.BytesIO(response_body)} expected_params = { 'Range': mock.ANY, 'Bucket': mock.ANY, 'Key': mock.ANY, 'RequestPayer': 'requester'} stubber.add_response('get_object', response, expected_params) stubber.activate() downloader = MultipartDownloader( client, TransferConfig(), InMemoryOSLayer({}), SequentialExecutor) downloader.download_file( 'bucket', 'key', 'filename', len(response_body), {'RequestPayer': 'requester'}) stubber.assert_no_pending_responses()
def create_session(self, profile=None): session = Session(profile=profile) # We have to set bogus credentials here or otherwise we'll trigger # an early credential chain resolution. sts = session.create_client( 'sts', aws_access_key_id='spam', aws_secret_access_key='eggs', ) stubber = Stubber(sts) stubber.activate() assume_role_provider = AssumeRoleProvider( load_config=lambda: session.full_config, client_creator=lambda *args, **kwargs: sts, cache={}, profile_name=profile, credential_sourcer=CanonicalNameCredentialSourcer([ self.env_provider, self.container_provider, self.metadata_provider ]) ) component_name = 'credential_provider' resolver = session.get_component(component_name) available_methods = [p.METHOD for p in resolver.providers] replacements = { 'env': self.env_provider, 'iam-role': self.metadata_provider, 'container-role': self.container_provider, 'assume-role': assume_role_provider } for name, provider in replacements.items(): try: index = available_methods.index(name) except ValueError: # The provider isn't in the session continue resolver.providers[index] = provider session.register_component( 'credential_provider', resolver ) return session, stubber
def _boto3_stubber(service, mocked_requests): client = boto3.client(service, region) stubber = Stubber(client) # Save a ref to the stubber so that we can deactivate it at the end of the test. created_stubbers.append(stubber) # Attach mocked requests to the Stubber and activate it. if not isinstance(mocked_requests, list): mocked_requests = [mocked_requests] for mocked_request in mocked_requests: stubber.add_response( mocked_request.method, mocked_request.response, expected_params=mocked_request.expected_params ) stubber.activate() # Add stubber to the collection of mocked clients. This allows to mock multiple clients. # Mocking twice the same client will replace the previous one. mocked_clients[service] = client return client
class TestRDSPagination(BaseSessionTest): def setUp(self): super(TestRDSPagination, self).setUp() self.region = 'us-west-2' self.client = self.session.create_client( 'rds', self.region) self.stubber = Stubber(self.client) def test_can_specify_zero_marker(self): service_response = { 'LogFileData': 'foo', 'Marker': '2', 'AdditionalDataPending': True } expected_params = { 'DBInstanceIdentifier': 'foo', 'LogFileName': 'bar', 'NumberOfLines': 2, 'Marker': '0' } function_name = 'download_db_log_file_portion' # The stubber will assert that the function is called with the expected # parameters. self.stubber.add_response( function_name, service_response, expected_params) self.stubber.activate() try: paginator = self.client.get_paginator(function_name) result = paginator.paginate( DBInstanceIdentifier='foo', LogFileName='bar', NumberOfLines=2, PaginationConfig={ 'StartingToken': '0', 'MaxItems': 3 }).build_full_result() self.assertEqual(result['LogFileData'], 'foo') self.assertIn('NextToken', result) except StubAssertionError as e: self.fail(str(e))
def test_provide_copy_source_client(self): source_client = self.session.create_client( 's3', 'eu-central-1', aws_access_key_id='foo', aws_secret_access_key='bar') source_stubber = Stubber(source_client) source_stubber.activate() self.addCleanup(source_stubber.deactivate) self.add_head_object_response(stubber=source_stubber) self.add_successful_copy_responses() call_kwargs = self.create_call_kwargs() call_kwargs['source_client'] = source_client future = self.manager.copy(**call_kwargs) future.result() # Make sure that all of the responses were properly # used for both clients. source_stubber.assert_no_pending_responses() self.stubber.assert_no_pending_responses()
class TestRDS(unittest.TestCase): def setUp(self): self.session = botocore.session.get_session() self.client = self.session.create_client('rds', 'us-west-2') self.stubber = Stubber(self.client) self.stubber.activate() def test_generate_db_auth_token(self): hostname = 'host.us-east-1.rds.amazonaws.com' port = 3306 username = '******' auth_token = self.client.generate_db_auth_token( DBHostname=hostname, Port=port, DBUsername=username) endpoint_url = 'host.us-east-1.rds.amazonaws.com:3306' self.assertIn(endpoint_url, auth_token) self.assertIn('Action=connect', auth_token) # Asserts that there is no scheme in the url self.assertTrue(auth_token.startswith(hostname))
class TestMturk(BaseSessionTest): def setUp(self): super(TestMturk, self).setUp() self.region = 'us-west-2' self.client = self.session.create_client( 'mturk', self.region) self.stubber = Stubber(self.client) self.stubber.activate() def tearDown(self): self.stubber.deactivate() def test_list_hits_aliased(self): self.stubber.add_response('list_hits_for_qualification_type', {}) self.stubber.add_response('list_hits_for_qualification_type', {}) params = {'QualificationTypeId': 'foo'} self.client.list_hi_ts_for_qualification_type(**params) self.client.list_hits_for_qualification_type(**params) self.stubber.assert_no_pending_responses()
class TestS3ObjectSummary(unittest.TestCase): def setUp(self): self.session = boto3.session.Session( aws_access_key_id='foo', aws_secret_access_key='bar', region_name='us-west-2') self.s3 = self.session.resource('s3') self.obj_summary = self.s3.ObjectSummary('my_bucket', 'my_key') self.obj_summary_size = 12 self.stubber = Stubber(self.s3.meta.client) self.stubber.activate() self.stubber.add_response( method='head_object', service_response={ 'ContentLength': self.obj_summary_size, 'ETag': 'my-etag', 'ContentType': 'binary' }, expected_params={ 'Bucket': 'my_bucket', 'Key': 'my_key' } ) def tearDown(self): self.stubber.deactivate() def test_has_load(self): self.assertTrue(hasattr(self.obj_summary, 'load'), 'load() was not injected onto ObjectSummary resource.') def test_autoloads_correctly(self): # In HeadObject the parameter returned is ContentLength, this # should get mapped to Size of ListObject since the resource uses # the shape returned to by ListObjects. self.assertEqual(self.obj_summary.size, self.obj_summary_size) def test_cannot_access_other_non_related_parameters(self): # Even though an HeadObject was used to load this, it should # only expose the attributes from its shape defined in ListObjects. self.assertFalse(hasattr(self.obj_summary, 'content_length'))
class TestSagemaker(BaseSessionTest): def setUp(self): super(TestSagemaker, self).setUp() self.region = 'us-west-2' self.client = self.session.create_client( 'sagemaker', self.region) self.stubber = Stubber(self.client) self.stubber.activate() self.hook_calls = [] def _hook(self, **kwargs): self.hook_calls.append(kwargs['event_name']) def tearDown(self): self.stubber.deactivate() def test_event_with_old_prefix(self): self.client.meta.events.register( 'provide-client-params.sagemaker.ListEndpoints', self._hook ) self.stubber.add_response('list_endpoints', {'Endpoints': []}) self.client.list_endpoints() self.assertEqual(self.hook_calls, [ 'provide-client-params.sagemaker.ListEndpoints' ]) def test_event_with_new_prefix(self): self.client.meta.events.register( 'provide-client-params.api.sagemaker.ListEndpoints', self._hook ) self.stubber.add_response('list_endpoints', {'Endpoints': []}) self.client.list_endpoints() self.assertEqual(self.hook_calls, [ 'provide-client-params.sagemaker.ListEndpoints' ])
def _can_use_parameter_in_client_call(session, case, use_alias=True): client = session.create_client( case['service'], region_name='us-east-1', aws_access_key_id='foo', aws_secret_access_key='bar') stubber = Stubber(client) stubber.activate() operation = case['operation'] params = case.get('extra_args', {}) params = params.copy() param_name = case['original_name'] if use_alias: param_name = case['new_name'] params[param_name] = case['parameter_value'] stubbed_response = case.get('stubbed_response', {}) stubber.add_response(operation, stubbed_response) try: getattr(client, operation)(**params) except ParamValidationError as e: raise AssertionError( 'Expecting %s to be valid parameter for %s.%s but received ' '%s.' % ( case['new_name'], case['service'], case['operation'], e) )
class StubbedClientTest(unittest.TestCase): def setUp(self): self.session = botocore.session.get_session() self.region = 'us-west-2' self.client = self.session.create_client( 's3', self.region, aws_access_key_id='foo', aws_secret_access_key='bar') self.stubber = Stubber(self.client) self.stubber.activate() def tearDown(self): self.stubber.deactivate() def reset_stubber_with_new_client(self, override_client_kwargs): client_kwargs = { 'service_name': 's3', 'region_name': self.region, 'aws_access_key_id': 'foo', 'aws_secret_access_key': 'bar' } client_kwargs.update(override_client_kwargs) self.client = self.session.create_client(**client_kwargs) self.stubber = Stubber(self.client) self.stubber.activate()
class TestSTSPresignedUrl(BaseSessionTest): def setUp(self): super(TestSTSPresignedUrl, self).setUp() self.client = self.session.create_client('sts', 'us-west-2') # Makes sure that no requests will go through self.stubber = Stubber(self.client) self.stubber.activate() def test_presigned_url_contains_no_content_type(self): timestamp = datetime(2017, 3, 22, 0, 0) with mock.patch('botocore.auth.datetime') as _datetime: _datetime.datetime.utcnow.return_value = timestamp url = self.client.generate_presigned_url('get_caller_identity', {}) # There should be no 'content-type' in x-amz-signedheaders expected_url = ( 'https://sts.amazonaws.com/?Action=GetCallerIdentity&' 'Version=2011-06-15&X-Amz-Algorithm=AWS4-HMAC-SHA256&' 'X-Amz-Credential=access_key%2F20170322%2Fus-east-1%2Fsts%2F' 'aws4_request&X-Amz-Date=20170322T000000Z&X-Amz-Expires=3600&' 'X-Amz-SignedHeaders=host&X-Amz-Signature=767845d2ee858069a598d5f' '8b497b75c7d57356885b1b3dba46dbbc0fc62bf5a' ) assert_url_equal(url, expected_url)
class TestADICommandS3(TransactionTestCase): fixtures = ('base/addon_3615', 'base/featured', 'base/appversion.json') date = '2014-07-10' stats_source = 's3' def add_response(self, stat): stat_path = os.path.join(hive_folder, 'src', '%s.hive' % stat) data = get_stats_data(stat_path) response = { 'Body': data, } expected_params = { 'Bucket': 'test-bucket', 'Key': os.path.join('amo_stats', stat, self.date, '000000_0'), 'Range': ANY } self.stubber.add_response('get_object', response, expected_params) def setUp(self): self.client = boto3.client('s3') self.stubber = Stubber(self.client) self.stubber.activate() def tearDown(self): self.stubber.deactivate() @override_settings(AWS_STATS_S3_BUCKET='test-bucket') @mock.patch('olympia.stats.management.commands.boto3') def test_update_counts_from_s3(self, mock_boto3): stats = ['app', 'locale', 'os', 'status', 'version'] for x in range(2): for stat in stats: self.add_response('update_counts_by_%s' % stat) mock_boto3.client.return_value = self.client management.call_command('update_counts_from_file', date=self.date, stats_source=self.stats_source) assert UpdateCount.objects.all().count() == 1 update_count = UpdateCount.objects.last() # should be identical to `statuses.userEnabled` assert update_count.count == 4 assert update_count.date == date(2014, 7, 10) assert update_count.versions == {u'3.8': 2, u'3.7': 3} assert update_count.statuses == {u'userDisabled': 1, u'userEnabled': 4} application = u'{ec8030f7-c20a-464f-9b0e-13a3a9e97384}' assert update_count.applications[application] == {u'3.6': 18} assert update_count.oses == {u'WINNT': 5} assert update_count.locales == {u'en-us': 1, u'en-US': 4} @override_settings(AWS_STATS_S3_BUCKET='test-bucket') @mock.patch('olympia.stats.management.commands.boto3') def test_download_counts_from_s3(self, mock_boto3): for x in range(2): self.add_response('download_counts') mock_boto3.client.return_value = self.client management.call_command('download_counts_from_file', date=self.date, stats_source=self.stats_source) assert DownloadCount.objects.all().count() == 2 download_count = DownloadCount.objects.get(addon_id=3615) assert download_count.count == 3 assert download_count.date == date(2014, 7, 10) assert download_count.sources == {u'search': 2, u'cb-dl-bob': 1}
def stubber(client): stubber = Stubber(client) stubber.activate() yield stubber stubber.deactivate()
class TestAccessCounts(TestCase): """Tests S3 Select""" def setUp(self): self.env_patcher = patch.dict( os.environ, { 'AWS_ACCESS_KEY_ID': 'test_key', 'AWS_SECRET_ACCESS_KEY': 'test_secret', 'AWS_DEFAULT_REGION': 'ng-north-1', 'ATHENA_DATABASE': 'athena-db', 'CLOUDTRAIL_BUCKET': 'cloudtrail-bucket', 'QUERY_RESULT_BUCKET': 'results-bucket', 'ACCESS_COUNTS_OUTPUT_DIR': 'AccessCounts', }) self.env_patcher.start() import index self.s3_stubber = Stubber(index.s3) self.s3_stubber.activate() self.athena_stubber = Stubber(index.athena) self.athena_stubber.activate() def tearDown(self): self.athena_stubber.deactivate() self.s3_stubber.deactivate() self.env_patcher.stop() def _start_query(self, query, execution_id): self.athena_stubber.add_response( method='start_query_execution', expected_params={ 'QueryExecutionContext': { 'Database': 'athena-db' }, 'QueryString': query, 'ResultConfiguration': { 'OutputLocation': 's3://results-bucket/AthenaQueryResults/' } }, service_response={'QueryExecutionId': execution_id}, ) def _end_query(self, execution_id=None): self.athena_stubber.add_response( method='get_query_execution', expected_params={'QueryExecutionId': execution_id} if execution_id is not None else None, service_response={ 'QueryExecution': { 'Status': { 'State': 'SUCCEEDED' } } }) def _run_queries(self, queries): for idx, query in enumerate(queries): self._start_query(query, str(idx)) for _ in queries: self._end_query() def test_access_counts(self): import index now = datetime.fromtimestamp(1234567890, timezone.utc) end_ts = now - timedelta(minutes=15) start_ts = now - timedelta(days=1) self.s3_stubber.add_response( method='get_object', expected_params={ 'Bucket': 'results-bucket', 'Key': 'ObjectAccessLog.last_updated_ts.txt', }, service_response={ 'Body': BytesIO(str(start_ts.timestamp()).encode()), }) self.s3_stubber.add_response(method='list_objects_v2', expected_params={ 'Bucket': 'results-bucket', 'Prefix': 'AthenaQueryResults', 'MaxKeys': 1000, }, service_response={}) self.s3_stubber.add_response(method='list_objects_v2', expected_params={ 'Bucket': 'cloudtrail-bucket', 'Prefix': 'AWSLogs/', 'Delimiter': '/', }, service_response={ 'CommonPrefixes': [{ 'Prefix': 'AWSLogs/123456/' }] }) self.s3_stubber.add_response( method='list_objects_v2', expected_params={ 'Bucket': 'cloudtrail-bucket', 'Prefix': 'AWSLogs/123456/CloudTrail/', 'Delimiter': '/', }, service_response={ 'CommonPrefixes': [{ 'Prefix': 'AWSLogs/123456/CloudTrail/ng-north-1/' }] }) self._run_queries([ index.DROP_CLOUDTRAIL, index.DROP_OBJECT_ACCESS_LOG, index.DROP_PACKAGE_HASHES ]) self._run_queries([ index.CREATE_CLOUDTRAIL, index.CREATE_OBJECT_ACCESS_LOG, index.CREATE_PACKAGE_HASHES ]) self._run_queries([ index.REPAIR_OBJECT_ACCESS_LOG, index.ADD_CLOUDTRAIL_PARTITION.format(account='123456', region='ng-north-1', year=2009, month=2, day=12), index.ADD_CLOUDTRAIL_PARTITION.format(account='123456', region='ng-north-1', year=2009, month=2, day=13), ]) self.s3_stubber.add_response(method='delete_object', expected_params={ 'Bucket': 'results-bucket', 'Key': 'ObjectAccessLog.last_updated_ts.txt', }, service_response={}) self._run_queries([ index.INSERT_INTO_OBJECT_ACCESS_LOG.format( start_ts=start_ts.timestamp(), end_ts=end_ts.timestamp()) ]) self.s3_stubber.add_response(method='put_object', expected_params={ 'Bucket': 'results-bucket', 'Key': 'ObjectAccessLog.last_updated_ts.txt', 'ContentType': 'text/plain', 'Body': str(end_ts.timestamp()), }, service_response={}) self._run_queries([ index.OBJECT_ACCESS_COUNTS, index.PACKAGE_ACCESS_COUNTS, index.PACKAGE_VERSION_ACCESS_COUNTS, index.BUCKET_ACCESS_COUNTS, index.EXTS_ACCESS_COUNTS ]) for idx, name in enumerate( ['Objects', 'Packages', 'PackageVersions', 'Bucket', 'Exts']): self.s3_stubber.add_response( method='head_object', expected_params={ 'Bucket': 'results-bucket', 'Key': f'AthenaQueryResults/{idx}.csv', }, service_response={'ContentLength': 123}) self.s3_stubber.add_response( method='copy_object', expected_params={ 'CopySource': { 'Bucket': 'results-bucket', 'Key': f'AthenaQueryResults/{idx}.csv', }, 'Bucket': 'results-bucket', 'Key': f'AccessCounts/{name}.csv', }, service_response={}) with patch('index.now', return_value=now), \ patch('time.sleep', return_value=None): index.handler(None, None)
class TestAutoscalingPagination(BaseSessionTest): def setUp(self): super(TestAutoscalingPagination, self).setUp() self.region = 'us-west-2' self.client = self.session.create_client( 'autoscaling', self.region, aws_secret_access_key='foo', aws_access_key_id='bar', aws_session_token='baz' ) self.stubber = Stubber(self.client) self.stubber.activate() def _setup_scaling_pagination( self, page_size=200, max_items=100, total_items=600 ): """ Add to the stubber to test paginating describe_scaling_activities. WARNING: This only handles cases where max_items cleanly divides page_size. """ requests_per_page = page_size / max_items if requests_per_page != ceil(requests_per_page): raise NotImplementedError( "This only handles setup where max_items is less than " "page_size and where max_items evenly divides page_size." ) requests_per_page = int(requests_per_page) num_pages = int(ceil(total_items / page_size)) previous_next_token = None for i in range(num_pages): page = self.create_describe_scaling_response(page_size=page_size) # Don't create a next_token for the final page if i + 1 == num_pages: next_token = None else: next_token = random_chars(10) expected_args = {} if previous_next_token: expected_args['StartingToken'] = previous_next_token # The same page may be accessed multiple times because we are # truncating it at max_items for _ in range(requests_per_page - 1): # The page is copied because the paginator will modify the # response object, causing issues when using the stubber. self.stubber.add_response( 'describe_scaling_activities', page.copy() ) if next_token is not None: page['NextToken'] = next_token # Copying the page here isn't necessary because it is about to # be blown away anyway. self.stubber.add_response( 'describe_scaling_activities', page ) previous_next_token = next_token def create_describe_scaling_response(self, page_size=200): """Create a valid describe_scaling_activities response.""" page = [] date = datetime.now() for _ in range(page_size): page.append({ 'AutoScalingGroupName': 'test', 'ActivityId': random_chars(10), 'Cause': 'test', 'StartTime': date, 'StatusCode': '200', }) return {'Activities': page} def test_repeated_build_full_results(self): # This ensures that we can cleanly paginate using build_full_results. max_items = 100 total_items = 600 self._setup_scaling_pagination( max_items=max_items, total_items=total_items, page_size=200 ) paginator = self.client.get_paginator('describe_scaling_activities') conf = {'MaxItems': max_items} pagination_tokens = [] result = paginator.paginate(PaginationConfig=conf).build_full_result() all_results = result['Activities'] while 'NextToken' in result: starting_token = result['NextToken'] # We should never get a duplicate pagination token. self.assertNotIn(starting_token, pagination_tokens) pagination_tokens.append(starting_token) conf['StartingToken'] = starting_token pages = paginator.paginate(PaginationConfig=conf) result = pages.build_full_result() all_results.extend(result['Activities']) self.assertEqual(len(all_results), total_items)
def stub(client): stubber = Stubber(client) stubber.add_response('change_resource_record_sets', response, expected) stubber.activate()
class TestAutoscalingPagination(BaseSessionTest): def setUp(self): super(TestAutoscalingPagination, self).setUp() self.region = 'us-west-2' self.client = self.session.create_client( 'autoscaling', self.region, aws_secret_access_key='foo', aws_access_key_id='bar', aws_session_token='baz' ) self.stubber = Stubber(self.client) self.stubber.activate() def _setup_scaling_pagination(self, page_size=200, max_items=100, total_items=600): """ Add to the stubber to test paginating describe_scaling_activities. WARNING: This only handles cases where max_items cleanly divides page_size. """ requests_per_page = page_size / max_items if requests_per_page != ceil(requests_per_page): raise NotImplementedError( "This only handles setup where max_items is less than " "page_size and where max_items evenly divides page_size." ) requests_per_page = int(requests_per_page) num_pages = int(ceil(total_items / page_size)) previous_next_token = None for i in range(num_pages): page = self.create_describe_scaling_response(page_size=page_size) # Don't create a next_token for the final page if i + 1 == num_pages: next_token = None else: next_token = random_chars(10) expected_args = {} if previous_next_token: expected_args['StartingToken'] = previous_next_token # The same page may be accessed multiple times because we are # truncating it at max_items for _ in range(requests_per_page - 1): # The page is copied because the paginator will modify the # response object, causing issues when using the stubber. self.stubber.add_response( 'describe_scaling_activities', page.copy() ) if next_token is not None: page['NextToken'] = next_token # Copying the page here isn't necessary because it is about to # be blown away anyway. self.stubber.add_response( 'describe_scaling_activities', page ) previous_next_token = next_token def create_describe_scaling_response(self, page_size=200): """Create a valid describe_scaling_activities response.""" page = [] date = datetime.now() for _ in range(page_size): page.append({ 'AutoScalingGroupName': 'test', 'ActivityId': random_chars(10), 'Cause': 'test', 'StartTime': date, 'StatusCode': '200', }) return {'Activities': page} def test_repeated_build_full_results(self): # This ensures that we can cleanly paginate using build_full_results. max_items = 100 total_items = 600 self._setup_scaling_pagination( max_items=max_items, total_items=total_items, page_size=200 ) paginator = self.client.get_paginator('describe_scaling_activities') conf = {'MaxItems': max_items} pagination_tokens = [] result = paginator.paginate(PaginationConfig=conf).build_full_result() all_results = result['Activities'] while 'NextToken' in result: starting_token = result['NextToken'] # We should never get a duplicate pagination token. self.assertNotIn(starting_token, pagination_tokens) pagination_tokens.append(starting_token) conf['StartingToken'] = starting_token pages = paginator.paginate(PaginationConfig=conf) result = pages.build_full_result() all_results.extend(result['Activities']) self.assertEqual(len(all_results), total_items)
def kinesisanalyticsv2_stubber(): kinesisanalyticsv2_stubber = Stubber(kinesisanalyticsv2) kinesisanalyticsv2_stubber.activate() yield kinesisanalyticsv2_stubber kinesisanalyticsv2_stubber.deactivate()
def ram_stubber(): ram_stubber = Stubber(ram) ram_stubber.activate() yield ram_stubber ram_stubber.deactivate()
def test_anon_access(self): """ Test anonymous call w/ ALLOW_ANONYMOUS_ACCESS """ bucket = "bucket" key = ".quilt/packages/manifest_hash" params = dict( bucket=bucket, manifest=key, ) expected_args = { 'Bucket': bucket, 'Key': key, 'Expression': "SELECT SUBSTRING(s.logical_key, 1) AS logical_key FROM s3object s", 'ExpressionType': 'SQL', 'InputSerialization': { 'CompressionType': 'NONE', 'JSON': { 'Type': 'LINES' } }, 'OutputSerialization': { 'JSON': { 'RecordDelimiter': '\n' } }, } env_patcher = patch.dict( os.environ, { 'AWS_ACCESS_KEY_ID': 'test_key', 'AWS_SECRET_ACCESS_KEY': 'test_secret', 'ALLOW_ANONYMOUS_ACCESS': '1' }) env_patcher.start() mock_s3 = boto3.client('s3') response = { 'ETag': '12345', 'VersionId': '1.0', 'ContentLength': 123, } expected_params = { 'Bucket': bucket, 'Key': key, } s3_stubber = Stubber(mock_s3) s3_stubber.activate() s3_stubber.add_response('head_object', response, expected_params) with patch.object(mock_s3, 'select_object_content', side_effect=[ self.s3response, self.s3response_meta ]) as client_patch, patch('boto3.Session.client', return_value=mock_s3): response = lambda_handler(self._make_event(params), None) print(response) assert response['statusCode'] == 200 folder = json.loads(read_body(response))['contents'] print(folder) assert len(folder['prefixes']) == 1 assert len(folder['objects']) == 1 assert folder['objects'][0]['logical_key'] == 'foo.csv' assert folder['prefixes'][0]['logical_key'] == 'bar/' s3_stubber.deactivate() env_patcher.stop()
def test_event_good(mocker): #-------------------------- # Test data # event = utils.load_test_data(test_data + 'cis26.json', my_region) sns_message = { 'Note': '"Enable Access Logging on CloudTrail logs bucket" remediation was successfully invoked via AWS Systems Manager', 'State': 'RESOLVED', 'Account': '111111111111', 'Remediation': 'Enable Access Logging on CloudTrail logs bucket', 'AffectedObject': 'CloudTrail: cloudtrail-awslogs-111111111111-cv5ddz5l-isengard-do-not-delete', 'metrics_data': mocker.ANY } post_metrics_expected_parms = { 'Solution': 'SO0111', 'UUID': '12345678-1234-1234-1234-123412341234', 'TimeStamp': mocker.ANY, 'Data': { 'generator_id': 'arn:aws:securityhub:::ruleset/cis-aws-foundations-benchmark/v/1.2.0/rule/2.6', 'type': '2.6 Ensure S3 bucket access logging is enabled on the CloudTrail S3 bucket', 'productArn': mocker.ANY, 'finding_triggered_by': 'Security Hub Findings - Custom Action', 'region': mocker.ANY, 'status': 'RESOLVED' }, 'Version': 'v1.2.0TEST' } ssmc = boto3.client('ssm', region_name=my_region) ssmc_s = Stubber(ssmc) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_uuid) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_version) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_uuid) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_version) ssmc_s.activate() mocker.patch('lib.metrics.Metrics.connect_to_ssm', return_value=ssmc) post_metrics = mocker.patch('lib.metrics.Metrics.post_metrics_to_api', return_value=None) mocker.patch('lib.awsapi_helpers.BotoSession.__init__', return_value=None) mocker.patch('lib.awsapi_helpers.AWSClient.connect', return_value=None) awsc = [boto3.client('ssm'), boto3.client('s3')] def mock_select(thing1, thing2): if thing2 == 'ssm': return awsc[0] elif thing2 == 's3': return awsc[1] else: return # Mock the boto client and replace the BotoSession client with our stub awsc_s = Stubber(awsc[0]) awsc_s.add_response('start_automation_execution', {}) awsc_s.activate() awss3_s = Stubber(awsc[1]) awss3_s.add_response('create_bucket', {}) awss3_s.add_response('put_bucket_encryption', {}) awss3_s.activate() mocker.patch('lib.awsapi_helpers.BotoSession.client', new=mock_select) sns = mocker.patch('lib.awsapi_helpers.AWSClient.postit', return_value=None) # Mock Notifier init = mocker.patch('lib.sechub_findings.Finding.flag') resolve = mocker.patch('lib.sechub_findings.Finding.resolve') # Prevent flushing to logs mocker.patch('lib.applogger.LogHandler.flush', return_value=None) #-------------------------- # Run the lambda # cis26.lambda_handler(event, None) init.assert_called_once_with( 'INITIAL: "Enable Access Logging on CloudTrail logs bucket" remediation started' ) resolve.assert_called_once_with( 'RESOLVED: "Enable Access Logging on CloudTrail logs bucket" remediation was successfully invoked via AWS Systems Manager' ) sns.assert_called_with('SO0111-SHARR_Topic', sns_message, my_region) post_metrics.assert_called_with(post_metrics_expected_parms)
def qldb_stubber(): qldb_stubber = Stubber(qldb) qldb_stubber.activate() yield qldb_stubber qldb_stubber.deactivate()
def _get_test_plan(self, max_changes): provider = Route53Provider('test', 'abc', '123', max_changes) # Use the stubber stubber = Stubber(provider._conn) stubber.activate() got = Zone('unit.tests.', []) list_hosted_zones_resp = { 'HostedZones': [], 'Marker': 'm', 'IsTruncated': False, 'MaxItems': '100', } stubber.add_response('list_hosted_zones', list_hosted_zones_resp, {}) create_hosted_zone_resp = { 'HostedZone': { 'Name': 'unit.tests.', 'Id': 'z42', 'CallerReference': 'abc', }, 'ChangeInfo': { 'Id': 'a12', 'Status': 'PENDING', 'SubmittedAt': '2017-01-29T01:02:03Z', 'Comment': 'hrm', }, 'DelegationSet': { 'Id': 'b23', 'CallerReference': 'blip', 'NameServers': [ 'n12.unit.tests.', ], }, 'Location': 'us-east-1', } stubber.add_response('create_hosted_zone', create_hosted_zone_resp, { 'Name': got.name, 'CallerReference': ANY, }) stubber.add_response( 'list_health_checks', { 'HealthChecks': self.health_checks, 'IsTruncated': False, 'MaxItems': '100', 'Marker': '', }) stubber.add_response('delete_health_check', {}, { 'HealthCheckId': ANY, }) stubber.add_response('delete_health_check', {}, { 'HealthCheckId': ANY, }) stubber.add_response('delete_health_check', {}, { 'HealthCheckId': ANY, }) stubber.add_response( 'change_resource_record_sets', { 'ChangeInfo': { 'Id': 'id', 'Status': 'PENDING', 'SubmittedAt': '2017-01-29T01:02:03Z', } }, { 'HostedZoneId': 'z42', 'ChangeBatch': ANY }) plan = provider.plan(self.expected) return provider, plan
class TestADICommandS3(TransactionTestCase): fixtures = ('base/addon_3615', 'base/featured', 'addons/persona', 'base/appversion.json') date = '2014-07-10' stats_source = 's3' def add_response(self, stat): stat_path = os.path.join(hive_folder, 'src', '%s.hive' % stat) data = get_stats_data(stat_path) response = { 'Body': data, } expected_params = {'Bucket': 'test-bucket', 'Key': os.path.join('amo_stats', stat, self.date, '000000_0'), 'Range': ANY} self.stubber.add_response('get_object', response, expected_params) def setUp(self): self.client = boto3.client('s3') self.stubber = Stubber(self.client) self.stubber.activate() def tearDown(self): self.stubber.deactivate() @override_settings(AWS_STATS_S3_BUCKET='test-bucket') @mock.patch( 'olympia.stats.management.commands.update_counts_from_file.' 'save_stats_to_file') @mock.patch('olympia.stats.management.commands.boto3') def test_update_counts_from_s3(self, mock_boto3, mock_save_stats_to_file): stats = ['app', 'locale', 'os', 'status', 'version'] for x in range(2): for stat in stats: self.add_response('update_counts_by_%s' % stat) mock_boto3.client.return_value = self.client management.call_command('update_counts_from_file', date=self.date, stats_source=self.stats_source) assert UpdateCount.objects.all().count() == 1 update_count = UpdateCount.objects.last() # should be identical to `statuses.userEnabled` assert update_count.count == 4 assert update_count.date == date(2014, 7, 10) assert update_count.versions == {u'3.8': 2, u'3.7': 3} assert update_count.statuses == {u'userDisabled': 1, u'userEnabled': 4} application = u'{ec8030f7-c20a-464f-9b0e-13a3a9e97384}' assert update_count.applications[application] == {u'3.6': 18} assert update_count.oses == {u'WINNT': 5} assert update_count.locales == {u'en-us': 1, u'en-US': 4} # save_stats_to_file is called with a non-saved model. assert isinstance(mock_save_stats_to_file.call_args[0][0], UpdateCount) @override_settings(AWS_STATS_S3_BUCKET='test-bucket') @mock.patch( 'olympia.stats.management.commands.download_counts_from_file.' 'save_stats_to_file') @mock.patch('olympia.stats.management.commands.boto3') def test_download_counts_from_s3(self, mock_boto3, mock_save_stats_to_file): for x in range(2): self.add_response('download_counts') mock_boto3.client.return_value = self.client management.call_command('download_counts_from_file', date=self.date, stats_source=self.stats_source) assert DownloadCount.objects.all().count() == 2 download_count = DownloadCount.objects.get(addon_id=3615) assert download_count.count == 3 assert download_count.date == date(2014, 7, 10) assert download_count.sources == {u'search': 2, u'cb-dl-bob': 1} # save_stats_to_file is called with a non-saved model. assert isinstance( mock_save_stats_to_file.call_args[0][0], DownloadCount) @override_settings(AWS_STATS_S3_BUCKET='test-bucket') @mock.patch( 'olympia.stats.management.commands.theme_update_counts_from_file.' 'save_stats_to_file') @mock.patch('olympia.stats.management.commands.boto3') def test_theme_update_counts_from_s3(self, mock_boto3, mock_save_stats_to_file): for x in range(2): self.add_response('theme_update_counts') mock_boto3.client.return_value = self.client management.call_command('theme_update_counts_from_file', date=self.date, stats_source=self.stats_source) assert ThemeUpdateCount.objects.all().count() == 1 # Persona 813 has addon id 15663: we need the count to be the sum of # the "old" request on the persona_id 813 (only the one with the source # "gp") and the "new" request on the addon_id 15663. tuc2 = ThemeUpdateCount.objects.get(addon_id=15663) assert tuc2.count == 15 assert mock_save_stats_to_file.call_count == 1 # save_stats_to_file is called with a non-saved model. assert isinstance( mock_save_stats_to_file.call_args[0][0], ThemeUpdateCount)
class TestADICommandS3(TransactionTestCase): fixtures = ('base/addon_3615', 'base/featured', 'addons/persona', 'base/appversion.json') date = '2014-07-10' stats_source = 's3' def add_response(self, stat): stat_path = os.path.join(hive_folder, 'src', '%s.hive' % stat) data = get_stats_data(stat_path) response = { 'Body': data, } expected_params = {'Bucket': 'test-bucket', 'Key': os.path.join('amo_stats', stat, self.date, '000000_0'), 'Range': ANY} self.stubber.add_response('get_object', response, expected_params) def setUp(self): self.client = boto3.client('s3') self.stubber = Stubber(self.client) self.stubber.activate() def tearDown(self): self.stubber.deactivate() @override_settings(AWS_STATS_S3_BUCKET='test-bucket') @mock.patch('olympia.stats.management.commands.boto3') def test_update_counts_from_s3(self, mock_boto3): stats = ['app', 'locale', 'os', 'status', 'version'] for x in range(2): for stat in stats: self.add_response('update_counts_by_%s' % stat) mock_boto3.client.return_value = self.client management.call_command('update_counts_from_file', date=self.date, stats_source=self.stats_source) assert UpdateCount.objects.all().count() == 1 update_count = UpdateCount.objects.last() # should be identical to `statuses.userEnabled` assert update_count.count == 4 assert update_count.date == date(2014, 7, 10) assert update_count.versions == {u'3.8': 2, u'3.7': 3} assert update_count.statuses == {u'userDisabled': 1, u'userEnabled': 4} application = u'{ec8030f7-c20a-464f-9b0e-13a3a9e97384}' assert update_count.applications[application] == {u'3.6': 18} assert update_count.oses == {u'WINNT': 5} assert update_count.locales == {u'en-us': 1, u'en-US': 4} @override_settings(AWS_STATS_S3_BUCKET='test-bucket') @mock.patch('olympia.stats.management.commands.boto3') def test_download_counts_from_s3(self, mock_boto3): for x in range(2): self.add_response('download_counts') mock_boto3.client.return_value = self.client management.call_command('download_counts_from_file', date=self.date, stats_source=self.stats_source) assert DownloadCount.objects.all().count() == 2 download_count = DownloadCount.objects.get(addon_id=3615) assert download_count.count == 3 assert download_count.date == date(2014, 7, 10) assert download_count.sources == {u'search': 2, u'cb-dl-bob': 1} @override_settings(AWS_STATS_S3_BUCKET='test-bucket') @mock.patch('olympia.stats.management.commands.boto3') def test_theme_update_counts_from_s3(self, mock_boto3): for x in range(2): self.add_response('theme_update_counts') mock_boto3.client.return_value = self.client management.call_command('theme_update_counts_from_file', date=self.date, stats_source=self.stats_source) assert ThemeUpdateCount.objects.all().count() == 1 # Persona 813 has addon id 15663: we need the count to be the sum of # the "old" request on the persona_id 813 (only the one with the source # "gp") and the "new" request on the addon_id 15663. tuc2 = ThemeUpdateCount.objects.get(addon_id=15663) assert tuc2.count == 15 @override_settings(AWS_STATS_S3_BUCKET='test-bucket') @mock.patch('olympia.stats.management.commands.boto3') def test_lwt_stats_go_to_migrated_static_theme(self, mock_boto3): lwt = Addon.objects.get(id=15663) lwt.delete() static_theme = addon_factory(type=amo.ADDON_STATICTHEME) MigratedLWT.objects.create( lightweight_theme=lwt, static_theme=static_theme) for x in range(2): self.add_response('theme_update_counts') mock_boto3.client.return_value = self.client management.call_command('theme_update_counts_from_file', date=self.date, stats_source=self.stats_source) assert ThemeUpdateCount.objects.all().count() == 0 assert UpdateCount.objects.all().count() == 1 assert UpdateCount.objects.get(addon_id=static_theme.id).count == 15 @override_settings(AWS_STATS_S3_BUCKET='test-bucket') @mock.patch('olympia.stats.management.commands.boto3') def test_lwt_stats_go_to_migrated_with_stats_already(self, mock_boto3): lwt = Addon.objects.get(id=15663) lwt.delete() static_theme = addon_factory(type=amo.ADDON_STATICTHEME) MigratedLWT.objects.create( lightweight_theme=lwt, static_theme=static_theme) UpdateCount.objects.create( addon=static_theme, count=123, date=date(2014, 7, 10)) for x in range(2): self.add_response('theme_update_counts') mock_boto3.client.return_value = self.client management.call_command('theme_update_counts_from_file', date=self.date, stats_source=self.stats_source) assert ThemeUpdateCount.objects.all().count() == 0 assert UpdateCount.objects.all().count() == 1 assert UpdateCount.objects.get(addon_id=static_theme.id).count == 138
class TestStubber(unittest.TestCase): def setUp(self): session = botocore.session.get_session() config = botocore.config.Config(signature_version=botocore.UNSIGNED) self.client = session.create_client('s3', config=config) self.stubber = Stubber(self.client) def test_stubber_returns_response(self): service_response = {'ResponseMetadata': {'foo': 'bar'}} self.stubber.add_response('list_objects', service_response) self.stubber.activate() response = self.client.list_objects(Bucket='foo') self.assertEqual(response, service_response) def test_activated_stubber_errors_with_no_registered_stubs(self): self.stubber.activate() with self.assertRaises(StubResponseError): self.client.list_objects(Bucket='foo') def test_stubber_errors_when_stubs_are_used_up(self): self.stubber.add_response('list_objects', {}) self.stubber.activate() self.client.list_objects(Bucket='foo') with self.assertRaises(StubResponseError): self.client.list_objects(Bucket='foo') def test_client_error_response(self): error_code = "AccessDenied" error_message = "Access Denied" self.stubber.add_client_error( 'list_objects', error_code, error_message) self.stubber.activate() with self.assertRaises(ClientError): self.client.list_objects(Bucket='foo') def test_expected_params_success(self): service_response = {} expected_params = {'Bucket': 'foo'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # This should be called successfully with no errors being thrown # for mismatching expected params. response = self.client.list_objects(Bucket='foo') self.assertEqual(response, service_response) def test_expected_params_fail(self): service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # This should call should raise an for mismatching expected params. with self.assertRaises(StubResponseError): self.client.list_objects(Bucket='foo') def test_expected_params_mixed_with_errors_responses(self): # Add an error response error_code = "AccessDenied" error_message = "Access Denied" self.stubber.add_client_error( 'list_objects', error_code, error_message) # Add a response with incorrect expected params service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # The first call should throw and error as expected. with self.assertRaises(ClientError): self.client.list_objects(Bucket='foo') # The second call should throw an error for unexpected parameters with self.assertRaisesRegexp(StubResponseError, 'Expected parameters'): self.client.list_objects(Bucket='foo') def test_can_continue_to_call_after_expected_params_fail(self): service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # Throw an error for unexpected parameters with self.assertRaises(StubResponseError): self.client.list_objects(Bucket='foo') # The stubber should still have the responses queued up # even though the original parameters did not match the expected ones. self.client.list_objects(Bucket='bar') self.stubber.assert_no_pending_responses() def test_still_relies_on_param_validation_with_expected_params(self): service_response = {} expected_params = {'Buck': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # Throw an error for invalid parameters with self.assertRaises(ParamValidationError): self.client.list_objects(Buck='bar')
def test_event_good(mocker): # Read test data event = utils.load_test_data(test_data + 'cis_1-3-iamuser1.json', my_region) sns_message = { 'Note': 'Access key over 90 days old found: AKIAGHJGJFGHJFGETHFG', 'State': 'INFO', 'Account': '111111111111', 'Remediation': 'Deactivate unused keys over 90 days old', 'AffectedObject': 'Access Key: AKIAGHJGJFGHJFGETHFG', 'metrics_data': {'status': 'INFO'}, } iam_keys = { "AccessKeyMetadata": [ { "UserName": "******", "AccessKeyId": "AKIAADFHWEREFGFHSDDF", "Status": "Active", "CreateDate": "2015-05-22T14:43:16+00:00" }, { "UserName": "******", "AccessKeyId": "AKIAGHJGJFGHJFGETHFG", "Status": "Active", "CreateDate": "2020-05-15T15:20:04+00:00" } ] } # Mock the constructor. We don't need the session created mocker.patch('lib.awsapi_helpers.BotoSession.__init__', return_value=None) mocker.patch('lib.awsapi_helpers.AWSClient.connect', return_value=None) # sess = BotoSession() iamc = boto3.client('iam', region_name=my_region) iamr = boto3.resource('iam', region_name=my_region) iamc_s = Stubber(iamc) iamr_s = Stubber(iamr.meta.client) iamc_s.add_response( 'list_access_keys', iam_keys ) iamr_s.add_response( 'update_access_key', {} ) iam_keys['AccessKeyMetadata'][0]['Status'] = 'Inactive' iamc_s.add_response( 'list_access_keys', iam_keys ) iamr_s.add_response( 'update_access_key', {} ) iam_keys['AccessKeyMetadata'][0]['Status'] = 'Inactive' iamc_s.add_response( 'list_access_keys', iam_keys ) iamc_s.activate() iamr_s.activate() # Replace BotoSession client/resource with our stub mocker.patch('lib.awsapi_helpers.BotoSession.client', return_value=iamc) mocker.patch('lib.awsapi_helpers.BotoSession.resource', return_value=iamr) sns = mocker.patch('lib.awsapi_helpers.AWSClient.postit', return_value=None) # Mock Notifier resolve = mocker.patch('lib.sechub_findings.Finding.resolve') flag = mocker.patch('lib.sechub_findings.Finding.flag') mocker.patch('lib.applogger.LogHandler.flush', return_value=None) cis1314.lambda_handler(event, None) flag.assert_called_once_with( 'INITIAL: "Deactivate unused keys over 90 days old" remediation started' ) resolve.assert_called_once_with( 'RESOLVED: Remediation completed successfully, create new access keys using IAM console.' ) sns.assert_called_with('SO0111-SHARR_Topic', sns_message, my_region)
def test_s3_delete(self): client = get_s3_client() stubber = Stubber(client) stubber.activate() stubber.add_response('put_object', self._base_response, dict(Bucket=ANY, Key=self._path, Body='test')) list_objects_response = { 'IsTruncated': False, 'Name': 'test-bucket', 'MaxKeys': 1000, 'Prefix': '', 'Contents': [{ u'LastModified': datetime.datetime(2016, 9, 23, 11, 17, 14), u'ETag': '"20d2cb13afb394301bbea0bcff19e12b"', u'StorageClass': 'STANDARD', u'Key': self._path, u'Owner': { u'DisplayName': 'test', u'ID': '31d89f79718dbd4435290740e6fa5e41cffafa7d9a3c323c85b525342e6341ae' }, u'Size': 77824 }], 'EncodingType': 'url', 'ResponseMetadata': { 'RequestId': 'abc123', 'HTTPStatusCode': 200, 'HostId': 'abc123' }, 'Marker': '' } stubber.add_response('list_objects', list_objects_response, { 'Bucket': ANY, 'Prefix': self._uuid }) stubber.add_response('delete_object', self._base_response, { 'Bucket': ANY, 'Key': self._path }) with patch('eventkit_cloud.utils.s3.open', mock_open(read_data='test'), create=True) as mock_open_obj: upload_to_s3(self._uuid, self._filename, self._filename, client=client) delete_from_s3(self._uuid, client=client)
'DisplayName': 'string', 'ID': 'string' } }, ], 'Name': 'string', 'EncodingType': 'url', 'KeyCount': 123, 'ContinuationToken': 'string' } s3_stubber.add_response('list_objects_v2', list_objects_response) s3_stubber.activate() mock_sm_client = boto3.client('secretsmanager') sm_stubber = Stubber(mock_sm_client) mock_secret_value_response = { 'ARN': 'arn:aws:secretsmanager:eu-west-7:123456789012:secret:tutorials/MyFirstSecret-jiObOV', 'Name': 'string', 'VersionId': 'EXAMPLE1-90ab-cdef-fedc-ba987EXAMPLE', 'SecretBinary': b'{"azkaban_username": "******", "azkaban_password": "******"}', 'CreatedDate': datetime(2015, 1, 1) } sm_stubber.add_response('get_secret_value', mock_secret_value_response) sm_stubber.add_response('get_secret_value', mock_secret_value_response) sm_stubber.activate()
def sns_stubber(): sns_stubber = Stubber(sns) sns_stubber.activate() yield sns_stubber sns_stubber.deactivate()
class StaxAuthTests(unittest.TestCase): """ Inherited class to run all unit tests for this module """ def setUp(self): self.cognito_client = botocore.session.get_session().create_client( "cognito-identity", region_name="ap-southeast-2", config=BotoConfig(signature_version=UNSIGNED), ) self.cognito_stub = Stubber(self.cognito_client) self.aws_srp_client = botocore.session.get_session().create_client( "cognito-idp", region_name="ap-southeast-2", config=BotoConfig(signature_version=UNSIGNED), ) self.aws_srp_stubber = Stubber(self.aws_srp_client) def tearDown(self): self.cognito_stub.deactivate() self.aws_srp_stubber.deactivate() def testStaxAuthInit(self): """ Test to initialise StaxAuth """ sa = StaxAuth("ApiAuth") self.assertEqual(sa.aws_region, "ap-southeast-2") def testToken(self): """ Test valid JWT is returned """ sa = StaxAuth("ApiAuth") self.stub_aws_srp(sa, "valid_username") token = sa.id_token_from_cognito( username="******", password="******", srp_client=self.aws_srp_client, ) self.assertEqual(token, "valid_token") def testTokenClient(self): """ Test the AWSSRP client is invoked and throws an error """ sa = StaxAuth("ApiAuth") with self.assertRaises(InvalidCredentialsException): sa.id_token_from_cognito(username="******", password="******") def testCredentialErrors(self): """ Test that boto errors are caught and converted to InvalidCredentialExceptions """ sa = StaxAuth("ApiAuth") # Test with invalid username password self.stub_aws_srp(sa, "bad_password", "NotAuthorizedException") user_not_found_success = False try: sa.id_token_from_cognito( username="******", password="******", srp_client=self.aws_srp_client, ) except InvalidCredentialsException as e: self.assertIn("Please check your Secret Key is correct", str(e)) user_not_found_success = True self.assertTrue(user_not_found_success) # Test with no access self.stub_aws_srp(sa, "no_access", "UserNotFoundException") no_access_success = False try: sa.id_token_from_cognito(username="******", password="******", srp_client=self.aws_srp_client) except InvalidCredentialsException as e: self.assertIn( "Please check your Access Key, that you have created your Api Token and that you are using the right STAX REGION", str(e), ) no_access_success = True self.assertTrue(no_access_success) # Test Unknown Error self.stub_aws_srp(sa, "Unknown", "UnitTesting") with self.assertRaises(InvalidCredentialsException): sa.id_token_from_cognito(username="******", password="******", srp_client=self.aws_srp_client) def testCreds(self): """ Test valid credentials are returned """ sa = StaxAuth("ApiAuth") token = jwt.encode({"sub": "unittest"}, "secret", algorithm="HS256") jwt_token = jwt.decode(token, verify=False) self.stub_cognito_creds(sa, jwt_token.get("sub")) creds = sa.sts_from_cognito_identity_pool(jwt_token.get("sub"), self.cognito_client) self.assertIn("Credentials", creds) self.assertTrue(creds.get("IdentityId").startswith("ap-southeast-2")) def testCredsClient(self): """ Test the cognito client is invoked and throws an error """ sa = StaxAuth("ApiAuth") # Test Invalid Credentials token = jwt.encode({"sub": "unittest"}, "secret", algorithm="HS256") jwt_token = jwt.decode(token, verify=False) with self.assertRaises(InvalidCredentialsException): sa.sts_from_cognito_identity_pool(jwt_token.get("sub")) # Test "Couldn't verify signed token" retry expected_parameters = { "IdentityPoolId": sa.identity_pool, "Logins": { f"cognito-idp.{sa.aws_region}.amazonaws.com/{sa.user_pool}": "unittest" } } for i in range(sa.max_retries): self.cognito_stub.add_client_error( "get_id", service_error_code="NotAuthorizedException", service_message= "Invalid login token. Couldn't verify signed token.", expected_params=expected_parameters, ) self.cognito_stub.activate() with self.assertRaises(InvalidCredentialsException) as e: sa.sts_from_cognito_identity_pool( jwt_token.get("sub"), cognito_client=self.cognito_client) self.assertEqual( str(e.exception), "InvalidCredentialsException: Retries Exceeded: Unexpected Client Error" ) self.assertEqual(len(self.cognito_stub._queue), 0) def testAuthErrors(self): """ Test that errors are thrown when keys are invalid """ sa = StaxAuth("ApiAuth") # Test with no username with self.assertRaises(InvalidCredentialsException): sa.requests_auth(username=None, password="******") # Test with no username with self.assertRaises(InvalidCredentialsException): sa.requests_auth(username="******", password=None) def stub_aws_srp(self, sa, username, error_code=None): expected_parameters = { "AuthFlow": "USER_SRP_AUTH", "AuthParameters": { "SRP_A": ANY, "USERNAME": username }, "ClientId": sa.client_id, } if error_code: self.aws_srp_stubber.add_client_error( "initiate_auth", service_error_code=error_code, expected_params=expected_parameters, ) else: self.aws_srp_stubber.add_response( "initiate_auth", { "ChallengeParameters": { "USER_ID_FOR_SRP": "user", "SALT": "4", "SRP_B": "5", "SECRET_BLOCK": "secblock", }, "ChallengeName": "PASSWORD_VERIFIER", }, expected_parameters, ) self.aws_srp_stubber.add_response( "respond_to_auth_challenge", { "AuthenticationResult": { "IdToken": "valid_token" }, }, { "ClientId": sa.client_id, "ChallengeName": ANY, "ChallengeResponses": ANY, }, ) self.aws_srp_stubber.activate() def stub_cognito_creds(self, sa, token: str): id_response = {"IdentityId": "ap-southeast-2"} id_params = { "IdentityPoolId": sa.identity_pool, "Logins": { f"cognito-idp.{sa.aws_region}.amazonaws.com/{sa.user_pool}": token }, } self.cognito_stub.add_response("get_id", id_response, id_params) id_creds_response = { "IdentityId": id_response["IdentityId"], "Credentials": { "AccessKeyId": "ASIAX000000000000000", "SecretKey": "0000000000000000000000000000000000000000", "SessionToken": "a-totally-valid-JWT", "Expiration": datetime(2020, 1, 14, 11, 52, 26), }, } id_creds_params = { "IdentityId": id_response["IdentityId"], "Logins": { f"cognito-idp.{sa.aws_region}.amazonaws.com/{sa.user_pool}": token }, } self.cognito_stub.add_response("get_credentials_for_identity", id_creds_response, id_creds_params) self.cognito_stub.activate() @responses.activate def testSigV4Headers(self): """ Test sigv4 signed auth headers """ # Get signed auth headers sa = StaxAuth("ApiAuth") id_creds = { "Credentials": { "AccessKeyId": "ASIAX000000000000000", "SecretKey": "0000000000000000000000000000000000000000", "SessionToken": "a-totally-valid-JWT", "Expiration": datetime(2020, 1, 14, 11, 52, 26), } } auth = sa.sigv4_signed_auth_headers(id_creds) # Mock request response_dict = {"Status": "OK"} responses.add( responses.GET, f"{Config.api_base_url()}/auth", json=response_dict, status=200, ) response = requests.get(f"{Config.api_base_url()}/auth", auth=auth) self.assertEqual(response.json(), response_dict) self.assertIn("Authorization", response.request.headers) def testApiTokenAuthNotExpired(self): """ Test credentials have not expired """ StaxConfig = Config StaxConfig.expiration = datetime.now(timezone.utc) + timedelta(hours=8) self.assertIsNotNone(StaxConfig.expiration) ApiTokenAuth.requests_auth("username", "password") self.assertIsNotNone(StaxConfig.auth) def testApiTokenAuth(self): """ Test generating new credentials """ sa = StaxAuth("ApiAuth") StaxConfig = Config StaxConfig.expiration = None token = jwt.encode({"sub": "valid_token"}, "secret", algorithm="HS256") jwt_token = jwt.decode(token, verify=False) self.stub_cognito_creds(sa, jwt_token.get("sub")) self.stub_aws_srp(sa, "username") ApiTokenAuth.requests_auth( "username", "password", srp_client=self.aws_srp_client, cognito_client=self.cognito_client, ) self.assertIsNotNone(StaxConfig.auth) @patch("test_auth.StaxAuth.requests_auth") def testApiTokenAuthExpiring(self, requests_auth_mock): """ Test credentials close to expiration get refreshed """ sa = StaxAuth("ApiAuth") StaxConfig = Config ## expiration 20 minutes in the future, no need to refresh StaxConfig.expiration = datetime.now( timezone.utc) + timedelta(minutes=20) ApiTokenAuth.requests_auth( "username", "password", srp_client=self.aws_srp_client, cognito_client=self.cognito_client, ) requests_auth_mock.assert_not_called() requests_auth_mock.reset_mock() ## expiration in 5 seconds from now, refresh to avoid token becoming stale used StaxConfig.expiration = datetime.now( timezone.utc) + timedelta(seconds=5) ApiTokenAuth.requests_auth( "username", "password", srp_client=self.aws_srp_client, cognito_client=self.cognito_client, ) requests_auth_mock.assert_called_once() requests_auth_mock.reset_mock() ## expiration in 5 minutes from now, refresh to avoid token becoming stale used ## override default triggering library to not refresh environ["TOKEN_EXPIRY_THRESHOLD_IN_MINS"] = "10" StaxConfig.expiration = datetime.now( timezone.utc) + timedelta(minutes=2) ApiTokenAuth.requests_auth( "username", "password", srp_client=self.aws_srp_client, cognito_client=self.cognito_client, ) requests_auth_mock.assert_called_once() def testRootAuthNotExpired(self): """ Test credentials have not expired """ StaxConfig = Config StaxConfig.expiration = datetime.now(timezone.utc) + timedelta(hours=8) self.assertIsNotNone(StaxConfig.expiration) RootAuth.requests_auth("username", "password") self.assertIsNotNone(StaxConfig.auth) def testRootAuth(self): """ Test generating new credentials """ sa = StaxAuth("JumaAuth") StaxConfig = Config StaxConfig.expiration = None token = jwt.encode({"sub": "valid_token"}, "secret", algorithm="HS256") jwt_token = jwt.decode(token, verify=False) self.stub_cognito_creds(sa, jwt_token.get("sub")) self.stub_aws_srp(sa, "username") RootAuth.requests_auth( "username", "password", srp_client=self.aws_srp_client, cognito_client=self.cognito_client, ) self.assertIsNotNone(StaxConfig.auth) def testApiAuth(self): """ Test auth through the Api class """ sa = StaxAuth("ApiAuth") StaxConfig = Config StaxConfig.expiration = None StaxConfig.access_key = "username" StaxConfig.secret_key = "password" token = jwt.encode({"sub": "valid_token"}, "secret", algorithm="HS256") jwt_token = jwt.decode(token, verify=False) self.stub_cognito_creds(sa, jwt_token.get("sub")) self.stub_aws_srp(sa, "username") Api._requests_auth = None Api._auth( srp_client=self.aws_srp_client, cognito_client=self.cognito_client, ) self.assertIsNotNone(Api._requests_auth)
class TestConfigServiceMetricsCollector(unittest.TestCase): def setUp(self): self.gdClient = boto3.client("config") self.gdStubber = Stubber(self.gdClient) self.gdStubber.activate() self.botoSessionMock = MagicMock() self.botoSessionMock.client.return_value = self.gdClient def testCollectShouldReturnCurrentFindingsMetricFromSingleRegionWithSingleDetectorOnSuccess( self): # Mock Config Service self.gdStubber.add_response("get_discovered_resource_counts", { "totalDiscoveredResources": 0, "resourceCounts": [] }, {}) self.gdStubber.add_response( "get_compliance_summary_by_resource_type", { "ComplianceSummariesByResourceType": [{ "ComplianceSummary": { "CompliantResourceCount": { "CappedCount": 2, "CapExceeded": False }, "NonCompliantResourceCount": { "CappedCount": 11, "CapExceeded": False }, "ComplianceSummaryTimestamp": 1565880614.85 } }] }, {}) # Collect metrics with patch("boto3.session.Session", return_value=self.botoSessionMock): collector = ConfigServiceMetricsCollector(regions=["eu-west-1"]) metrics = collector.collect() currentResourcesMetric = metrics[0] self.assertEqual(currentResourcesMetric.name, "aws_config_current_resources") self.assertEqual(currentResourcesMetric.type, "gauge") self.assertEqual(len(currentResourcesMetric.samples), 3) self.assertEqual(currentResourcesMetric.samples[0].value, 1) self.assertEqual(currentResourcesMetric.samples[0].labels, { "region": "eu-west-1", "severity": "low" }) self.assertEqual(currentResourcesMetric.samples[1].value, 2) self.assertEqual(currentResourcesMetric.samples[1].labels, { "region": "eu-west-1", "severity": "medium" }) self.assertEqual(currentResourcesMetric.samples[2].value, 3) self.assertEqual(currentResourcesMetric.samples[2].labels, { "region": "eu-west-1", "severity": "high" }) scrapeErrorsMetric = metrics[1] self.assertEqual(scrapeErrorsMetric.name, "aws_configservice_scrape_errors") self.assertEqual(scrapeErrorsMetric.type, "counter") self.assertEqual(len(scrapeErrorsMetric.samples), 1) self.assertEqual(scrapeErrorsMetric.samples[0].value, 0) self.assertEqual(scrapeErrorsMetric.samples[0].labels, {"region": "eu-west-1"}) self.gdStubber.assert_no_pending_responses()
class TestStubber(unittest.TestCase): def setUp(self): session = botocore.session.get_session() config = botocore.config.Config(signature_version=botocore.UNSIGNED, s3={'addressing_style': 'path'}) self.client = session.create_client('s3', region_name='us-east-1', config=config) self.stubber = Stubber(self.client) def test_stubber_returns_response(self): service_response = {'ResponseMetadata': {'foo': 'bar'}} self.stubber.add_response('list_objects', service_response) self.stubber.activate() response = self.client.list_objects(Bucket='foo') self.assertEqual(response, service_response) def test_context_manager_returns_response(self): service_response = {'ResponseMetadata': {'foo': 'bar'}} self.stubber.add_response('list_objects', service_response) with self.stubber: response = self.client.list_objects(Bucket='foo') self.assertEqual(response, service_response) def test_activated_stubber_errors_with_no_registered_stubs(self): self.stubber.activate() # Params one per line for readability. with self.assertRaisesRegexp(UnStubbedResponseError, "Unexpected API Call"): self.client.list_objects(Bucket='asdfasdfasdfasdf', Delimiter='asdfasdfasdfasdf', Prefix='asdfasdfasdfasdf', EncodingType='url') def test_stubber_errors_when_stubs_are_used_up(self): self.stubber.add_response('list_objects', {}) self.stubber.activate() self.client.list_objects(Bucket='foo') with self.assertRaises(UnStubbedResponseError): self.client.list_objects(Bucket='foo') def test_client_error_response(self): error_code = "AccessDenied" error_message = "Access Denied" self.stubber.add_client_error('list_objects', error_code, error_message) self.stubber.activate() with self.assertRaises(ClientError): self.client.list_objects(Bucket='foo') def test_can_add_expected_params_to_client_error(self): self.stubber.add_client_error('list_objects', 'Error', 'error', expected_params={'Bucket': 'foo'}) self.stubber.activate() with self.assertRaises(ClientError): self.client.list_objects(Bucket='foo') def test_can_expected_param_fails_in_client_error(self): self.stubber.add_client_error('list_objects', 'Error', 'error', expected_params={'Bucket': 'foo'}) self.stubber.activate() # We expect an AssertionError instead of a ClientError # because we're calling the operation with the wrong # param value. with self.assertRaises(AssertionError): self.client.list_objects(Bucket='wrong-argument-value') def test_expected_params_success(self): service_response = {} expected_params = {'Bucket': 'foo'} self.stubber.add_response('list_objects', service_response, expected_params) self.stubber.activate() # This should be called successfully with no errors being thrown # for mismatching expected params. response = self.client.list_objects(Bucket='foo') self.assertEqual(response, service_response) def test_expected_params_fail(self): service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response('list_objects', service_response, expected_params) self.stubber.activate() # This should call should raise an for mismatching expected params. with self.assertRaisesRegexp(StubResponseError, "{'Bucket': 'bar'},\n"): self.client.list_objects(Bucket='foo') def test_expected_params_mixed_with_errors_responses(self): # Add an error response error_code = "AccessDenied" error_message = "Access Denied" self.stubber.add_client_error('list_objects', error_code, error_message) # Add a response with incorrect expected params service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response('list_objects', service_response, expected_params) self.stubber.activate() # The first call should throw and error as expected. with self.assertRaises(ClientError): self.client.list_objects(Bucket='foo') # The second call should throw an error for unexpected parameters with self.assertRaisesRegexp(StubResponseError, 'Expected parameters'): self.client.list_objects(Bucket='foo') def test_can_continue_to_call_after_expected_params_fail(self): service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response('list_objects', service_response, expected_params) self.stubber.activate() # Throw an error for unexpected parameters with self.assertRaises(StubResponseError): self.client.list_objects(Bucket='foo') # The stubber should still have the responses queued up # even though the original parameters did not match the expected ones. self.client.list_objects(Bucket='bar') self.stubber.assert_no_pending_responses() def test_still_relies_on_param_validation_with_expected_params(self): service_response = {} expected_params = {'Buck': 'bar'} self.stubber.add_response('list_objects', service_response, expected_params) self.stubber.activate() # Throw an error for invalid parameters with self.assertRaises(ParamValidationError): self.client.list_objects(Buck='bar') def test_any_ignores_param_for_validation(self): service_response = {} expected_params = {'Bucket': stub.ANY} self.stubber.add_response('list_objects', service_response, expected_params) self.stubber.add_response('list_objects', service_response, expected_params) try: with self.stubber: self.client.list_objects(Bucket='foo') self.client.list_objects(Bucket='bar') except StubAssertionError: self.fail("stub.ANY failed to ignore parameter for validation.") def test_mixed_any_and_concrete_params(self): service_response = {} expected_params = {'Bucket': stub.ANY, 'Key': 'foo.txt'} self.stubber.add_response('head_object', service_response, expected_params) self.stubber.add_response('head_object', service_response, expected_params) try: with self.stubber: self.client.head_object(Bucket='foo', Key='foo.txt') self.client.head_object(Bucket='bar', Key='foo.txt') except StubAssertionError: self.fail("stub.ANY failed to ignore parameter for validation.") def test_nested_any_param(self): service_response = {} expected_params = { 'Bucket': 'foo', 'Key': 'bar.txt', 'Metadata': { 'MyMeta': stub.ANY, } } self.stubber.add_response('put_object', service_response, expected_params) self.stubber.add_response('put_object', service_response, expected_params) try: with self.stubber: self.client.put_object(Bucket='foo', Key='bar.txt', Metadata={ 'MyMeta': 'Foo', }) self.client.put_object(Bucket='foo', Key='bar.txt', Metadata={ 'MyMeta': 'Bar', }) except StubAssertionError: self.fail( "stub.ANY failed to ignore nested parameter for validation.") def test_ANY_repr(self): self.assertEqual(repr(stub.ANY), '<ANY>') def test_none_param(self): service_response = {} expected_params = {'Buck': None} self.stubber.add_response('list_objects', service_response, expected_params) self.stubber.activate() # Throw an error for invalid parameters with self.assertRaises(StubAssertionError): self.client.list_objects(Buck='bar') def test_many_expected_params(self): service_response = {} expected_params = { 'Bucket': 'mybucket', 'Prefix': 'myprefix', 'Delimiter': '/', 'EncodingType': 'url' } self.stubber.add_response('list_objects', service_response, expected_params) try: with self.stubber: self.client.list_objects(**expected_params) except StubAssertionError: self.fail( "Stubber inappropriately raised error for same parameters.") def test_no_stub_for_presign_url(self): try: with self.stubber: url = self.client.generate_presigned_url( ClientMethod='get_object', Params={ 'Bucket': 'mybucket', 'Key': 'mykey' }) self.assertEqual(url, 'https://s3.amazonaws.com/mybucket/mykey') except StubResponseError: self.fail( 'Stubbed responses should not be required for generating ' 'presigned requests') def test_can_stub_with_presign_url_mixed_in(self): desired_response = {} expected_params = { 'Bucket': 'mybucket', 'Prefix': 'myprefix', } self.stubber.add_response('list_objects', desired_response, expected_params) with self.stubber: url = self.client.generate_presigned_url(ClientMethod='get_object', Params={ 'Bucket': 'myotherbucket', 'Key': 'myotherkey' }) self.assertEqual( url, 'https://s3.amazonaws.com/myotherbucket/myotherkey') actual_response = self.client.list_objects(**expected_params) self.assertEqual(desired_response, actual_response) self.stubber.assert_no_pending_responses() def test_parse_get_bucket_location(self): error_code = "NoSuchBucket" error_message = "The specified bucket does not exist" self.stubber.add_client_error('get_bucket_location', error_code, error_message) self.stubber.activate() with self.assertRaises(ClientError): self.client.get_bucket_location(Bucket='foo') def test_parse_get_bucket_location_returns_response(self): service_response = {"LocationConstraint": "us-west-2"} self.stubber.add_response('get_bucket_location', service_response) self.stubber.activate() response = self.client.get_bucket_location(Bucket='foo') self.assertEqual(response, service_response)
class TestStubber(unittest.TestCase): def setUp(self): session = botocore.session.get_session() config = botocore.config.Config(signature_version=botocore.UNSIGNED) self.client = session.create_client('s3', config=config) self.stubber = Stubber(self.client) def test_stubber_returns_response(self): service_response = {'ResponseMetadata': {'foo': 'bar'}} self.stubber.add_response('list_objects', service_response) self.stubber.activate() response = self.client.list_objects(Bucket='foo') self.assertEqual(response, service_response) def test_context_manager_returns_response(self): service_response = {'ResponseMetadata': {'foo': 'bar'}} self.stubber.add_response('list_objects', service_response) with self.stubber: response = self.client.list_objects(Bucket='foo') self.assertEqual(response, service_response) def test_activated_stubber_errors_with_no_registered_stubs(self): self.stubber.activate() # Params one per line for readability. with self.assertRaisesRegexp(StubResponseError, "'Bucket': 'asdfasdfasdfasdf',\n"): self.client.list_objects( Bucket='asdfasdfasdfasdf', Delimiter='asdfasdfasdfasdf', Prefix='asdfasdfasdfasdf', EncodingType='url') def test_stubber_errors_when_stubs_are_used_up(self): self.stubber.add_response('list_objects', {}) self.stubber.activate() self.client.list_objects(Bucket='foo') with self.assertRaises(StubResponseError): self.client.list_objects(Bucket='foo') def test_client_error_response(self): error_code = "AccessDenied" error_message = "Access Denied" self.stubber.add_client_error( 'list_objects', error_code, error_message) self.stubber.activate() with self.assertRaises(ClientError): self.client.list_objects(Bucket='foo') def test_expected_params_success(self): service_response = {} expected_params = {'Bucket': 'foo'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # This should be called successfully with no errors being thrown # for mismatching expected params. response = self.client.list_objects(Bucket='foo') self.assertEqual(response, service_response) def test_expected_params_fail(self): service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # This should call should raise an for mismatching expected params. with self.assertRaisesRegexp(StubResponseError, "{'Bucket': 'bar'},\n"): self.client.list_objects(Bucket='foo') def test_expected_params_mixed_with_errors_responses(self): # Add an error response error_code = "AccessDenied" error_message = "Access Denied" self.stubber.add_client_error( 'list_objects', error_code, error_message) # Add a response with incorrect expected params service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # The first call should throw and error as expected. with self.assertRaises(ClientError): self.client.list_objects(Bucket='foo') # The second call should throw an error for unexpected parameters with self.assertRaisesRegexp(StubResponseError, 'Expected parameters'): self.client.list_objects(Bucket='foo') def test_can_continue_to_call_after_expected_params_fail(self): service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # Throw an error for unexpected parameters with self.assertRaises(StubResponseError): self.client.list_objects(Bucket='foo') # The stubber should still have the responses queued up # even though the original parameters did not match the expected ones. self.client.list_objects(Bucket='bar') self.stubber.assert_no_pending_responses() def test_still_relies_on_param_validation_with_expected_params(self): service_response = {} expected_params = {'Buck': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # Throw an error for invalid parameters with self.assertRaises(ParamValidationError): self.client.list_objects(Buck='bar') def test_any_ignores_param_for_validation(self): service_response = {} expected_params = {'Bucket': stub.ANY} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.add_response( 'list_objects', service_response, expected_params) try: with self.stubber: self.client.list_objects(Bucket='foo') self.client.list_objects(Bucket='bar') except StubAssertionError: self.fail("stub.ANY failed to ignore parameter for validation.") def test_mixed_any_and_concrete_params(self): service_response = {} expected_params = {'Bucket': stub.ANY, 'Key': 'foo.txt'} self.stubber.add_response( 'head_object', service_response, expected_params) self.stubber.add_response( 'head_object', service_response, expected_params) try: with self.stubber: self.client.head_object(Bucket='foo', Key='foo.txt') self.client.head_object(Bucket='bar', Key='foo.txt') except StubAssertionError: self.fail("stub.ANY failed to ignore parameter for validation.") def test_none_param(self): service_response = {} expected_params = {'Buck': None} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # Throw an error for invalid parameters with self.assertRaises(StubAssertionError): self.client.list_objects(Buck='bar') def test_many_expected_params(self): service_response = {} expected_params = { 'Bucket': 'mybucket', 'Prefix': 'myprefix', 'Delimiter': '/', 'EncodingType': 'url' } self.stubber.add_response( 'list_objects', service_response, expected_params) try: with self.stubber: self.client.list_objects(**expected_params) except StubAssertionError: self.fail( "Stubber inappropriately raised error for same parameters.")
def test_event_good(mocker): #-------------------------- # Test data # event = utils.load_test_data(test_data + 'cis43.json', my_region) sns_message = { 'Note': '"Remove all rules from the default security group" remediation was successful', 'State': 'RESOLVED', 'Account': '111111111111', 'Remediation': 'Remove all rules from the default security group', 'AffectedObject': 'Security Group: sg-02cfbecbc814a3c24', 'metrics_data': { 'status': 'RESOLVED' } } desc_sg = { "SecurityGroups": [{ "Description": "Default SG", "GroupName": "SC-111111111111-pp-gz465ubujkfrs-SandboxSecurityGroup-175ZDF23V5MGX", "IpPermissions": [{ "FromPort": 80, "IpProtocol": "tcp", "IpRanges": [{ "CidrIp": "0.0.0.0/0" }], "Ipv6Ranges": [], "PrefixListIds": [], "ToPort": 80, "UserIdGroupPairs": [] }, { "FromPort": 9000, "IpProtocol": "tcp", "IpRanges": [{ "CidrIp": "72.21.198.65/32" }], "Ipv6Ranges": [], "PrefixListIds": [], "ToPort": 9000, "UserIdGroupPairs": [] }, { "FromPort": 22, "IpProtocol": "tcp", "IpRanges": [{ "CidrIp": "0.0.0.0/0" }], "Ipv6Ranges": [], "PrefixListIds": [], "ToPort": 22, "UserIdGroupPairs": [] }], "OwnerId": "123412341234", "GroupId": "sg-006bf520b9581b2d9", "IpPermissionsEgress": [{ "IpProtocol": "-1", "IpRanges": [{ "CidrIp": "0.0.0.0/0" }], "Ipv6Ranges": [], "PrefixListIds": [], "UserIdGroupPairs": [] }], "Tags": [], "VpcId": "vpc-11111113" }] } #-------------------------- # Mock/stub # mocker.patch('lib.awsapi_helpers.BotoSession.__init__', return_value=None) mocker.patch('lib.awsapi_helpers.AWSClient.connect', return_value=None) # Mock the boto client and replace the BotoSession client with our stub awsc = boto3.resource('ec2') awsc_s = Stubber(awsc.meta.client) awsc_s.add_response('describe_security_groups', desc_sg) awsc_s.add_response('revoke_security_group_ingress', {}) awsc_s.add_response('revoke_security_group_egress', {}) awsc_s.activate() mocker.patch('lib.awsapi_helpers.BotoSession.resource', return_value=awsc) sns = mocker.patch('lib.awsapi_helpers.AWSClient.postit', return_value=None) # Mock Notifier init = mocker.patch('lib.sechub_findings.Finding.flag') resolve = mocker.patch('lib.sechub_findings.Finding.resolve') # Prevent flushing to logs mocker.patch('lib.applogger.LogHandler.flush', return_value=None) #-------------------------- # Run the lambda # cis43.lambda_handler(event, None) init.assert_called_once_with( 'INITIAL: "Remove all rules from the default security group" remediation started' ) resolve.assert_called_once_with( 'RESOLVED: "Remove all rules from the default security group" remediation was successful' ) sns.assert_called_with('SO0111-SHARR_Topic', sns_message, my_region)
class TestStubber(unittest.TestCase): def setUp(self): session = botocore.session.get_session() config = botocore.config.Config( signature_version=botocore.UNSIGNED, s3={'addressing_style': 'path'} ) self.client = session.create_client( 's3', region_name='us-east-1', config=config) self.stubber = Stubber(self.client) def test_stubber_returns_response(self): service_response = {'ResponseMetadata': {'foo': 'bar'}} self.stubber.add_response('list_objects', service_response) self.stubber.activate() response = self.client.list_objects(Bucket='foo') self.assertEqual(response, service_response) def test_context_manager_returns_response(self): service_response = {'ResponseMetadata': {'foo': 'bar'}} self.stubber.add_response('list_objects', service_response) with self.stubber: response = self.client.list_objects(Bucket='foo') self.assertEqual(response, service_response) def test_activated_stubber_errors_with_no_registered_stubs(self): self.stubber.activate() # Params one per line for readability. with self.assertRaisesRegexp(UnStubbedResponseError, "Unexpected API Call"): self.client.list_objects( Bucket='asdfasdfasdfasdf', Delimiter='asdfasdfasdfasdf', Prefix='asdfasdfasdfasdf', EncodingType='url') def test_stubber_errors_when_stubs_are_used_up(self): self.stubber.add_response('list_objects', {}) self.stubber.activate() self.client.list_objects(Bucket='foo') with self.assertRaises(UnStubbedResponseError): self.client.list_objects(Bucket='foo') def test_client_error_response(self): error_code = "AccessDenied" error_message = "Access Denied" self.stubber.add_client_error( 'list_objects', error_code, error_message) self.stubber.activate() with self.assertRaises(ClientError): self.client.list_objects(Bucket='foo') def test_can_add_expected_params_to_client_error(self): self.stubber.add_client_error( 'list_objects', 'Error', 'error', expected_params={'Bucket': 'foo'} ) self.stubber.activate() with self.assertRaises(ClientError): self.client.list_objects(Bucket='foo') def test_can_expected_param_fails_in_client_error(self): self.stubber.add_client_error( 'list_objects', 'Error', 'error', expected_params={'Bucket': 'foo'} ) self.stubber.activate() # We expect an AssertionError instead of a ClientError # because we're calling the operation with the wrong # param value. with self.assertRaises(AssertionError): self.client.list_objects(Bucket='wrong-argument-value') def test_expected_params_success(self): service_response = {} expected_params = {'Bucket': 'foo'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # This should be called successfully with no errors being thrown # for mismatching expected params. response = self.client.list_objects(Bucket='foo') self.assertEqual(response, service_response) def test_expected_params_fail(self): service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # This should call should raise an for mismatching expected params. with self.assertRaisesRegexp(StubResponseError, "{'Bucket': 'bar'},\n"): self.client.list_objects(Bucket='foo') def test_expected_params_mixed_with_errors_responses(self): # Add an error response error_code = "AccessDenied" error_message = "Access Denied" self.stubber.add_client_error( 'list_objects', error_code, error_message) # Add a response with incorrect expected params service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # The first call should throw and error as expected. with self.assertRaises(ClientError): self.client.list_objects(Bucket='foo') # The second call should throw an error for unexpected parameters with self.assertRaisesRegexp(StubResponseError, 'Expected parameters'): self.client.list_objects(Bucket='foo') def test_can_continue_to_call_after_expected_params_fail(self): service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # Throw an error for unexpected parameters with self.assertRaises(StubResponseError): self.client.list_objects(Bucket='foo') # The stubber should still have the responses queued up # even though the original parameters did not match the expected ones. self.client.list_objects(Bucket='bar') self.stubber.assert_no_pending_responses() def test_still_relies_on_param_validation_with_expected_params(self): service_response = {} expected_params = {'Buck': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # Throw an error for invalid parameters with self.assertRaises(ParamValidationError): self.client.list_objects(Buck='bar') def test_any_ignores_param_for_validation(self): service_response = {} expected_params = {'Bucket': stub.ANY} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.add_response( 'list_objects', service_response, expected_params) try: with self.stubber: self.client.list_objects(Bucket='foo') self.client.list_objects(Bucket='bar') except StubAssertionError: self.fail("stub.ANY failed to ignore parameter for validation.") def test_mixed_any_and_concrete_params(self): service_response = {} expected_params = {'Bucket': stub.ANY, 'Key': 'foo.txt'} self.stubber.add_response( 'head_object', service_response, expected_params) self.stubber.add_response( 'head_object', service_response, expected_params) try: with self.stubber: self.client.head_object(Bucket='foo', Key='foo.txt') self.client.head_object(Bucket='bar', Key='foo.txt') except StubAssertionError: self.fail("stub.ANY failed to ignore parameter for validation.") def test_nested_any_param(self): service_response = {} expected_params = { 'Bucket': 'foo', 'Key': 'bar.txt', 'Metadata': { 'MyMeta': stub.ANY, } } self.stubber.add_response( 'put_object', service_response, expected_params) self.stubber.add_response( 'put_object', service_response, expected_params) try: with self.stubber: self.client.put_object( Bucket='foo', Key='bar.txt', Metadata={ 'MyMeta': 'Foo', } ) self.client.put_object( Bucket='foo', Key='bar.txt', Metadata={ 'MyMeta': 'Bar', } ) except StubAssertionError: self.fail( "stub.ANY failed to ignore nested parameter for validation.") def test_ANY_repr(self): self.assertEqual(repr(stub.ANY), '<ANY>') def test_none_param(self): service_response = {} expected_params = {'Buck': None} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # Throw an error for invalid parameters with self.assertRaises(StubAssertionError): self.client.list_objects(Buck='bar') def test_many_expected_params(self): service_response = {} expected_params = { 'Bucket': 'mybucket', 'Prefix': 'myprefix', 'Delimiter': '/', 'EncodingType': 'url' } self.stubber.add_response( 'list_objects', service_response, expected_params) try: with self.stubber: self.client.list_objects(**expected_params) except StubAssertionError: self.fail( "Stubber inappropriately raised error for same parameters.") def test_no_stub_for_presign_url(self): try: with self.stubber: url = self.client.generate_presigned_url( ClientMethod='get_object', Params={ 'Bucket': 'mybucket', 'Key': 'mykey' } ) self.assertEqual( url, 'https://s3.amazonaws.com/mybucket/mykey') except StubResponseError: self.fail( 'Stubbed responses should not be required for generating ' 'presigned requests' ) def test_can_stub_with_presign_url_mixed_in(self): desired_response = {} expected_params = { 'Bucket': 'mybucket', 'Prefix': 'myprefix', } self.stubber.add_response( 'list_objects', desired_response, expected_params) with self.stubber: url = self.client.generate_presigned_url( ClientMethod='get_object', Params={ 'Bucket': 'myotherbucket', 'Key': 'myotherkey' } ) self.assertEqual( url, 'https://s3.amazonaws.com/myotherbucket/myotherkey') actual_response = self.client.list_objects(**expected_params) self.assertEqual(desired_response, actual_response) self.stubber.assert_no_pending_responses()
def ddb_stubber(): ddb_stubber = Stubber(ItemsModel.get_table().meta.client) ddb_stubber.activate() yield ddb_stubber ddb_stubber.deactivate()
def test_event_good(mocker): #-------------------------- # Test data # event = utils.load_test_data(test_data + 'cis43.json', my_region) sns_message = { 'Note': '"Remove all rules from the default security group" remediation was successful', 'State': 'RESOLVED', 'Account': '111111111111', 'Remediation': 'Remove all rules from the default security group', 'AffectedObject': 'Security Group: sg-02cfbecbc814a3c24', 'metrics_data': mocker.ANY } desc_sg = { "SecurityGroups": [{ "Description": "Default SG", "GroupName": "SC-111111111111-pp-gz465ubujkfrs-SandboxSecurityGroup-175ZDF23V5MGX", "IpPermissions": [{ "FromPort": 80, "IpProtocol": "tcp", "IpRanges": [{ "CidrIp": "0.0.0.0/0" }], "Ipv6Ranges": [], "PrefixListIds": [], "ToPort": 80, "UserIdGroupPairs": [] }, { "FromPort": 9000, "IpProtocol": "tcp", "IpRanges": [{ "CidrIp": "72.21.198.65/32" }], "Ipv6Ranges": [], "PrefixListIds": [], "ToPort": 9000, "UserIdGroupPairs": [] }, { "FromPort": 22, "IpProtocol": "tcp", "IpRanges": [{ "CidrIp": "0.0.0.0/0" }], "Ipv6Ranges": [], "PrefixListIds": [], "ToPort": 22, "UserIdGroupPairs": [] }], "OwnerId": "123412341234", "GroupId": "sg-006bf520b9581b2d9", "IpPermissionsEgress": [{ "IpProtocol": "-1", "IpRanges": [{ "CidrIp": "0.0.0.0/0" }], "Ipv6Ranges": [], "PrefixListIds": [], "UserIdGroupPairs": [] }], "Tags": [], "VpcId": "vpc-11111113" }] } post_metrics_expected_parms = { 'Solution': 'SO0111', 'UUID': '12345678-1234-1234-1234-123412341234', 'TimeStamp': mocker.ANY, 'Data': { 'generator_id': 'arn:aws:securityhub:::ruleset/cis-aws-foundations-benchmark/v/1.2.0/rule/4.3', 'type': '4.3 Ensure the default security group of every VPC restricts all traffic', 'productArn': mocker.ANY, 'finding_triggered_by': 'Security Hub Findings - Custom Action', 'region': mocker.ANY, 'status': 'RESOLVED' }, 'Version': 'v1.2.0TEST' } ssmc = boto3.client('ssm', region_name=my_region) ssmc_s = Stubber(ssmc) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_uuid) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_version) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_uuid) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_version) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_uuid) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_version) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_uuid) ssmc_s.activate() mocker.patch('lib.metrics.Metrics.connect_to_ssm', return_value=ssmc) post_metrics = mocker.patch('lib.metrics.Metrics.post_metrics_to_api', return_value=None) mocker.patch('lib.awsapi_helpers.BotoSession.__init__', return_value=None) mocker.patch('lib.awsapi_helpers.AWSClient.connect', return_value=None) # Mock the boto client and replace the BotoSession client with our stub awsc = boto3.resource('ec2') awsc_s = Stubber(awsc.meta.client) awsc_s.add_response('describe_security_groups', desc_sg) awsc_s.add_response('revoke_security_group_ingress', {}) awsc_s.add_response('revoke_security_group_egress', {}) awsc_s.activate() mocker.patch('lib.awsapi_helpers.BotoSession.resource', return_value=awsc) sns = mocker.patch('lib.awsapi_helpers.AWSClient.postit', return_value=None) # Mock Notifier init = mocker.patch('lib.sechub_findings.Finding.flag') resolve = mocker.patch('lib.sechub_findings.Finding.resolve') # Prevent flushing to logs mocker.patch('lib.applogger.LogHandler.flush', return_value=None) #-------------------------- # Run the lambda # cis43.lambda_handler(event, None) init.assert_called_once_with( 'INITIAL: "Remove all rules from the default security group" remediation started' ) resolve.assert_called_once_with( 'RESOLVED: "Remove all rules from the default security group" remediation was successful' ) sns.assert_called_with('SO0111-SHARR_Topic', sns_message, my_region) post_metrics.assert_called_with(post_metrics_expected_parms)
class PackagePromoteTestBase(unittest.TestCase): credentials = { 'aws_access_key_id': 'test_aws_access_key_id', 'aws_secret_access_key': 'test_aws_secret_access_key', 'aws_session_token': 'test_aws_session_token', } handler = staticmethod(t4_lambda_pkgpush.promote_package) parent_bucket = 'parent-bucket' src_registry = f's3://{parent_bucket}' parent_pkg_name = 'parent/pkg-name' parent_commit_message = 'parent commit message' dst_bucket = 'dest-bucket' dst_registry = f's3://{dst_bucket}' dst_pkg_name = 'dest/pkg-name' dst_pkg_loc_params = { 'registry': dst_registry, 'name': dst_pkg_name, } mock_timestamp = 1600298935.9767091 mock_timestamp_pointer_name = '1600298935' file_size = 1 files_number = 2 @classmethod def get_file_data(cls, pk: PhysicalKey): return hashlib.sha256(str(pk).encode()).digest()[:1] * cls.file_size @classmethod def get_file_hash(cls, pk: PhysicalKey): return hashlib.sha256(cls.get_file_data(pk)).hexdigest() @classmethod def get_file_meta(cls, pk: PhysicalKey): return {f'meta-{pk}': f'value-{pk}'} @classmethod def get_pkg_entry(cls, path): pk = PhysicalKey.from_url( f's3://{cls.parent_bucket}/{path}?versionId=obj{path}Version') return PackageEntry( pk, cls.file_size, { 'type': 'SHA256', 'value': cls.get_file_hash(pk) }, cls.get_file_meta(pk), ) @classmethod def prepare_prefix_pkg_entries(cls, prefix, files_range, lk_prefix=''): return { lk_prefix + str(x): cls.get_pkg_entry(f'{prefix}{x}') for x in files_range } @classmethod def get_pkg_entries(cls): return cls.prepare_prefix_pkg_entries('test/pkg/', range(cls.files_number)) @classmethod def setUpClass(cls): super().setUpClass() pkg = Package() pkg._set_commit_message(cls.parent_commit_message) pkg._workflow = { 'config': f's3://{cls.parent_bucket}/.quilt/workflows/config.yml?versionId=configVersion', 'id': 'gamma', 'schemas': { 'top-secret': f's3://{cls.parent_bucket}/top-secret.schema.json?versionId=schemaVersion' }, } pkg.set_meta({'meta': 'old meta'}) cls.entries = cls.get_pkg_entries() for lk, entry in cls.entries.items(): pkg.set(lk, entry) manifest_buf = io.BytesIO() pkg._dump(manifest_buf) cls.parent_manifest = manifest_buf.getvalue() cls.parent_top_hash = pkg.top_hash cls.src_params = { 'parent': { 'registry': cls.src_registry, 'name': cls.parent_pkg_name, 'top_hash': cls.parent_top_hash, }, } @staticmethod def make_lambda_s3_stubber(): return Stubber(t4_lambda_pkgpush.s3) def setUp(self): super().setUp() self.headers = { 'content-type': 'application/json', } self.s3_stubber = Stubber(boto3.client('s3')) self.s3_stubber.activate() self.addCleanup(self.s3_stubber.deactivate) s3_client_patcher = mock.patch( 'quilt3.data_transfer.S3ClientProvider.find_correct_client', lambda *args, **kwargs: self.s3_stubber.client, ) s3_client_patcher.start() self.addCleanup(s3_client_patcher.stop) user_session_mock = mock.NonCallableMagicMock( spec_set=boto3.session.Session) user_session_mock.client.return_value = self.s3_stubber.client get_user_boto_session_patcher = mock.patch( 't4_lambda_pkgpush.get_user_boto_session', return_value=user_session_mock, ) self.get_user_boto_session_mock = get_user_boto_session_patcher.start() self.addCleanup(get_user_boto_session_patcher.stop) def calculate_pkg_hashes_side_effect(s3_client, pkg): for lk, entry in pkg.walk(): if entry.hash is None: entry.hash = { 'type': 'SHA256', 'value': self.get_file_hash(entry.physical_key), } calculate_pkg_hashes_patcher = mock.patch.object( t4_lambda_pkgpush, 'calculate_pkg_hashes', side_effect=calculate_pkg_hashes_side_effect, ) calculate_pkg_hashes_patcher.start() self.addCleanup(calculate_pkg_hashes_patcher.stop) @contextlib.contextmanager def mock_successors(self, successors): workflow_config_mock = mock.MagicMock() workflow_config_mock.config = { 'successors': successors, } src_registry = get_package_registry(self.src_registry) def side_effect(registry_url): if registry_url == self.src_registry: return src_registry return mock.DEFAULT with mock.patch.object(src_registry, 'get_workflow_config', return_value=workflow_config_mock), \ mock.patch('t4_lambda_pkgpush.get_package_registry', side_effect=side_effect, wraps=get_package_registry): yield def make_request_wrapper(self, params, *, credentials, **kwargs): return { "params": params, "credentials": credentials, } def make_request_base(self, data, *, credentials, **kwargs): result = self.handler( self.make_request_wrapper(data, credentials=credentials, **kwargs), None, ) # Check that result can be serialized to JSON. json.dumps(result) return result @mock.patch('time.time', mock.MagicMock(return_value=mock_timestamp)) def make_request(self, params, **kwargs): self.get_user_boto_session_mock.reset_mock() with mock.patch('quilt3.telemetry.reset_session_id') as reset_session_id_mock, \ calculate_sha256_patcher(return_value=[]) as calculate_sha256_mock: response = self.make_request_base(params, credentials=self.credentials, **kwargs) self.get_user_boto_session_mock.assert_called_once_with( **self.credentials, ) reset_session_id_mock.assert_called_once_with() if calculate_sha256_mock.called: calculate_sha256_mock.assert_called_once_with([], []) return response def setup_s3_load_pkg_source(self): self.s3_stubber.add_response( 'head_object', service_response={ 'VersionId': 'parentManifestVersion', 'ContentLength': len(self.parent_manifest), }, expected_params={ 'Bucket': self.parent_bucket, 'Key': f'.quilt/packages/{self.parent_top_hash}', }, ) self.s3_stubber.add_response( 'get_object', service_response={ 'VersionId': 'manifestVersion', 'ContentLength': len(self.parent_manifest), 'Body': io.BytesIO(self.parent_manifest), }, expected_params={ 'Bucket': self.parent_bucket, 'Key': f'.quilt/packages/{self.parent_top_hash}', 'VersionId': 'parentManifestVersion', }, ) def setup_s3(self, expected_pkg, *, copy_data): manifest = io.BytesIO() expected_pkg.dump(manifest) top_hash = expected_pkg.top_hash self.setup_s3_load_pkg_source() if copy_data: for src, (lk, dst) in zip(self.entries.values(), expected_pkg.walk()): self.s3_stubber.add_response( method='copy_object', service_response={ 'VersionId': 'dst_' + src.physical_key.version_id, }, expected_params={ 'CopySource': { 'Bucket': src.physical_key.bucket, 'Key': src.physical_key.path, 'VersionId': src.physical_key.version_id, }, 'Bucket': self.dst_bucket, 'Key': f'{self.dst_pkg_name}/{lk}', }) # Push new manifest. self.s3_stubber.add_response( 'put_object', service_response={}, expected_params={ 'Bucket': self.dst_bucket, 'Key': f'.quilt/packages/{top_hash}', 'Body': manifest.getvalue(), }, ) self.s3_stubber.add_response( 'put_object', service_response={}, expected_params={ 'Body': top_hash.encode(), 'Bucket': self.dst_bucket, 'Key': f'.quilt/named_packages/{self.dst_pkg_name}/{str(int(self.mock_timestamp))}' }) self.s3_stubber.add_response( 'put_object', service_response={ 'ResponseMetadata': { 'RequestId': 'foo' }, }, expected_params={ 'Body': top_hash.encode(), 'Bucket': self.dst_bucket, 'Key': f'.quilt/named_packages/{self.dst_pkg_name}/latest' }) def prepare_pkg(self, *, copy_data): expected_pkg = Package() pkg_entries = self.entries.items() if copy_data: pkg_entries = [( lk, e.with_physical_key( PhysicalKey(self.dst_bucket, f'{self.dst_pkg_name}/{lk}', 'dst_' + e.physical_key.version_id)), ) for lk, e in pkg_entries] for lk, entry in pkg_entries: expected_pkg.set(lk, entry) expected_pkg._set_commit_message(None) return expected_pkg
def test_not_remediated(mocker): #-------------------------- # Test data # event = utils.load_test_data(test_data + 'cis29.json', my_region) post_metrics_expected_parms = { 'Solution': 'SO0111', 'UUID': '12345678-1234-1234-1234-123412341234', 'TimeStamp': mocker.ANY, 'Data': { 'generator_id': 'arn:aws:securityhub:::ruleset/cis-aws-foundations-benchmark/v/1.2.0/rule/2.9', 'type': '2.9 Ensure VPC flow logging is enabled in all VPCs', 'productArn': mocker.ANY, 'finding_triggered_by': 'Security Hub Findings - Custom Action', 'region': mocker.ANY, 'status': 'FAILED' }, 'Version': 'v1.2.0TEST' } os.environ['AWS_SESSION_TOKEN'] = 'FAKETOKEN' os.environ['FLOW_LOG_ROLE_ARN'] = 'FAKELOGROLEARN' ssmc = boto3.client('ssm', region_name=my_region) ssmc_s = Stubber(ssmc) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_uuid) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_version) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_uuid) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_version) ssmc_s.activate() mocker.patch('lib.metrics.Metrics.connect_to_ssm', return_value=ssmc) post_metrics = mocker.patch('lib.metrics.Metrics.post_metrics_to_api', return_value=None) mocker.patch('lib.awsapi_helpers.BotoSession.__init__', return_value=None) mocker.patch('lib.awsapi_helpers.AWSClient.connect', return_value=None) awsc = [boto3.client('logs'), boto3.client('ec2')] def mock_select(thing1, thing2): if thing2 == 'logs': return awsc[0] else: return awsc[1] awsc_s = Stubber(awsc[0]) awsc_s.add_response('create_log_group', {}) awsc_s.activate() aws2c_s = Stubber(awsc[1]) aws2c_s.add_response('create_flow_logs', {}) aws2c_s.add_response('describe_flow_logs', {'FlowLogs': []}) aws2c_s.activate() # redirect to mock_select above to return the proper stub mocker.patch('lib.awsapi_helpers.BotoSession.client', new=mock_select) # Mock notifications init = mocker.patch('lib.sechub_findings.Finding.flag') resolve = mocker.patch('lib.sechub_findings.Finding.resolve') update = mocker.patch('lib.sechub_findings.Finding.update_text') mocker.patch('lib.applogger.LogHandler.flush', return_value=None) #-------------------------- # Run the lambda # cis29.lambda_handler(event, None) init.assert_called_once_with( 'INITIAL: "Enable VPC flow logging in all VPCs" remediation started') update.assert_called_once_with( 'FAILED: "Enable VPC flow logging in all VPCs" remediation failed. Please remediate manually', status='FAILED') resolve.assert_not_called() post_metrics.assert_called_with(post_metrics_expected_parms)
class TestBatchWatcherDaemon(UploadTestCaseUsingMockAWS): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def setUp(self): super().setUp() self.environment = { 'API_KEY': 'test' } self.environmentor = EnvironmentSetup(self.environment) self.environmentor.enter() self.batch_watcher = BatchWatcher() self.mock_batch_client = Stubber(self.batch_watcher.batch_client) self.mock_ec2_client = Stubber(self.batch_watcher.ec2_client) self.mock_lambda_client = Stubber(self.batch_watcher.lambda_client) def tearDown(self): self.environmentor.exit() super().tearDown() @patch('upload.lambdas.batch_watcher.batch_watcher.UploadDB.run_query') def test_find_incomplete_batch_jobs(self, mock_run_query): mock_run_query.return_value = QueryResult() csum_jobs, val_jobs = self.batch_watcher.find_incomplete_batch_jobs() self.assertEqual(csum_jobs, [{"id": "123", "job_id": "123", "file_id": "test/test"}]) self.assertEqual(val_jobs, [{"id": "123", "job_id": "123", "file_id": "test/test"}]) def test_find_and_kill_deployment_batch_instances(self): describe_params = { "Filters": [ {"Name": 'key-name', "Values": [f"hca-upload-{self.deployment_stage}"]}, {"Name": 'instance-state-name', "Values": ["running"]} ] } describe_output = { "Reservations": [{ "Instances": [ { "InstanceId": "instance_one" }, { "InstanceId": "instance_two" } ] }] } instance_ids = ["instance_one", "instance_two"] terminate_params = { "InstanceIds": instance_ids } self.mock_ec2_client.add_response("describe_instances", describe_output, describe_params) self.mock_ec2_client.add_response("terminate_instances", {}, terminate_params) self.mock_ec2_client.activate() killed_instance_ids = self.batch_watcher.find_and_kill_deployment_batch_instances() self.assertEqual(killed_instance_ids, instance_ids) @patch('upload.lambdas.batch_watcher.batch_watcher.UploadDB.run_query_with_params') @patch('upload.lambdas.batch_watcher.batch_watcher.BatchWatcher.schedule_validation_job') def test_schedule_job_with_validation(self, mock_schedule_validation_job, mock_run_query): row = { "id": "123", "file_id": "test_area/test_id", "job_id": "124", "docker_image": "test_docker_image", "original_validation_id": "567" } self.batch_watcher.schedule_job(row, "validation") mock_schedule_validation_job.assert_called_with("test_area", "test_id", "test_docker_image", "567") row = { "id": "123", "file_id": "test_area/test_id", "job_id": "124", "docker_image": "test_docker_image", "original_validation_id": None } self.batch_watcher.schedule_job(row, "validation") mock_schedule_validation_job.assert_called_with("test_area", "test_id", "test_docker_image", "123") @patch('upload.lambdas.batch_watcher.batch_watcher.UploadDB.run_query_with_params') @patch('upload.lambdas.batch_watcher.batch_watcher.BatchWatcher.invoke_checksum_lambda') def test_schedule_job_with_checksum(self, mock_invoke_csum_lambda, mock_run_query): row = { "id": "123", "file_id": "test_area/test_id", "job_id": "124" } self.batch_watcher.schedule_job(row, "checksum") mock_invoke_csum_lambda.assert_called_with("test_area/test_id") def test_invoke_checksum_lambda(self): payload = { 'Records': [{ 'eventName': 'ObjectCreated:Put', "s3": { "bucket": { "name": f"org-humancellatlas-upload-{self.deployment_stage}" }, "object": { "key": "test_area/test_file_id" } } }] } lambda_params = { "FunctionName": f"dcp-upload-csum-{self.deployment_stage}", "InvocationType": "Event", "Payload": json.dumps(payload).encode() } self.mock_lambda_client.add_response('invoke', {}, lambda_params) self.mock_lambda_client.activate() self.batch_watcher.invoke_checksum_lambda("test_area/test_file_id") self.mock_lambda_client.deactivate() def test_should_instances_be_killed_true(self): test_one_rows = [ { "id": "123", "file_id": "test/test", "job_id": "124" }, { "id": "234", "file_id": "test/test", "job_id": "235" } ] output_one = { "jobs": [{ "status": "FAILED", "jobName": "test", "jobId": "test", "jobQueue": "test", "startedAt": 1234, "jobDefinition": "test" }] } output_two = { "jobs": [{ "status": "SUCCEEDED", "jobName": "test", "jobId": "test", "jobQueue": "test", "startedAt": 1234, "jobDefinition": "test" }] } self.mock_batch_client.add_response('describe_jobs', output_one, {"jobs": ["124"]}) self.mock_batch_client.add_response('describe_jobs', output_two, {"jobs": ["235"]}) self.mock_batch_client.activate() kill_instances = self.batch_watcher.should_instances_be_killed(test_one_rows) self.assertEqual(kill_instances, True) self.mock_batch_client.deactivate() def test_should_instances_be_killed_false(self): test_rows = [ { "id": "345", "file_id": "test/test", "job_id": "346" }, { "id": "456", "file_id": "test/test", "job_id": "457" } ] output_two = { "jobs": [{ "status": "SUCCEEDED", "jobName": "test", "jobId": "test", "jobQueue": "test", "startedAt": 1234, "jobDefinition": "test" }] } self.mock_batch_client.add_response('describe_jobs', output_two, {"jobs": ["346"]}) self.mock_batch_client.add_response('describe_jobs', output_two, {"jobs": ["457"]}) self.mock_batch_client.activate() kill_instances = self.batch_watcher.should_instances_be_killed(test_rows) self.assertEqual(kill_instances, False) self.mock_batch_client.deactivate() def test_get_job_status(self): output_two = { "jobs": [{ "status": "SUCCEEDED", "jobName": "test", "jobId": "test", "jobQueue": "test", "startedAt": 1234, "jobDefinition": "test" }] } self.mock_batch_client.add_response('describe_jobs', output_two, {"jobs": ["346"]}) self.mock_batch_client.activate() status = self.batch_watcher._get_job_status("346") self.assertEqual(status, "SUCCEEDED") self.mock_batch_client.deactivate()
def test_event_good(mocker): #-------------------------- # Test data # event = utils.load_test_data(test_data + 'cis29.json', my_region) sns_message = { 'Note': '"Enable VPC flow logging in all VPCs" remediation was successful', 'State': 'RESOLVED', 'Account': '111111111111', 'Remediation': 'Enable VPC flow logging in all VPCs', 'AffectedObject': 'VPC Flow Logs for VPC: vpc-d1a07fba', 'metrics_data': mocker.ANY } os.environ['AWS_SESSION_TOKEN'] = 'FAKETOKEN' os.environ['FLOW_LOG_ROLE_ARN'] = 'FAKELOGROLEARN' post_metrics_expected_parms = { 'Solution': 'SO0111', 'UUID': '12345678-1234-1234-1234-123412341234', 'TimeStamp': mocker.ANY, 'Data': { 'generator_id': 'arn:aws:securityhub:::ruleset/cis-aws-foundations-benchmark/v/1.2.0/rule/2.9', 'type': '2.9 Ensure VPC flow logging is enabled in all VPCs', 'productArn': mocker.ANY, 'finding_triggered_by': 'Security Hub Findings - Custom Action', 'region': mocker.ANY, 'status': 'RESOLVED' }, 'Version': 'v1.2.0TEST' } ssmc = boto3.client('ssm', region_name=my_region) ssmc_s = Stubber(ssmc) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_uuid) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_version) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_uuid) ssmc_s.add_response('get_parameter', mock_ssm_get_parameter_version) ssmc_s.activate() mocker.patch('lib.metrics.Metrics.connect_to_ssm', return_value=ssmc) post_metrics = mocker.patch('lib.metrics.Metrics.post_metrics_to_api', return_value=None) # Mock the constructor. We don't need the session created mocker.patch('lib.awsapi_helpers.BotoSession.__init__', return_value=None) mocker.patch('lib.awsapi_helpers.AWSClient.connect', return_value=None) awsc = [boto3.client('logs'), boto3.client('ec2')] def mock_select(thing1, thing2): if thing2 == 'logs': return awsc[0] else: return awsc[1] # Mock the boto clients and replace the BotoSession client with our stub awsc_s = Stubber(awsc[0]) awsc_s.add_response('create_log_group', {}) awsc_s.activate() aws2c_s = Stubber(awsc[1]) aws2c_s.add_response('create_flow_logs', {}) aws2c_s.add_response('describe_flow_logs', {'FlowLogs': [{ 'FlowLogStatus': 'ACTIVE' }]}) aws2c_s.activate() sns = mocker.patch('lib.awsapi_helpers.AWSClient.postit', return_value=None) # redirect to mock_select above to return the proper stub mocker.patch('lib.awsapi_helpers.BotoSession.client', new=mock_select) # Mock notifications init = mocker.patch('lib.sechub_findings.Finding.flag') resolve = mocker.patch('lib.sechub_findings.Finding.resolve') mocker.patch('lib.applogger.LogHandler.flush', return_value=None) #-------------------------- # Run the lambda # cis29.lambda_handler(event, None) init.assert_called_once_with( 'INITIAL: "Enable VPC flow logging in all VPCs" remediation started') resolve.assert_called_once_with( 'RESOLVED: "Enable VPC flow logging in all VPCs" remediation was successful' ) sns.assert_called_with('SO0111-SHARR_Topic', sns_message, my_region) post_metrics.assert_called_with(post_metrics_expected_parms)
class TestStubber(unittest.TestCase): def setUp(self): session = botocore.session.get_session() config = botocore.client.Config(signature_version=botocore.UNSIGNED) self.client = session.create_client('s3', config=config) self.stubber = Stubber(self.client) def test_stubber_returns_response(self): service_response = {'ResponseMetadata': {'foo': 'bar'}} self.stubber.add_response('list_objects', service_response) self.stubber.activate() response = self.client.list_objects(Bucket='foo') self.assertEqual(response, service_response) def test_activated_stubber_errors_with_no_registered_stubs(self): self.stubber.activate() with self.assertRaises(StubResponseError): self.client.list_objects(Bucket='foo') def test_stubber_errors_when_stubs_are_used_up(self): self.stubber.add_response('list_objects', {}) self.stubber.activate() self.client.list_objects(Bucket='foo') with self.assertRaises(StubResponseError): self.client.list_objects(Bucket='foo') def test_client_error_response(self): error_code = "AccessDenied" error_message = "Access Denied" self.stubber.add_client_error( 'list_objects', error_code, error_message) self.stubber.activate() with self.assertRaises(ClientError): self.client.list_objects(Bucket='foo') def test_expected_params_success(self): service_response = {} expected_params = {'Bucket': 'foo'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # This should be called successfully with no errors being thrown # for mismatching expected params. response = self.client.list_objects(Bucket='foo') self.assertEqual(response, service_response) def test_expected_params_fail(self): service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # This should call should raise an for mismatching expected params. with self.assertRaises(StubResponseError): self.client.list_objects(Bucket='foo') def test_expected_params_mixed_with_errors_responses(self): # Add an error response error_code = "AccessDenied" error_message = "Access Denied" self.stubber.add_client_error( 'list_objects', error_code, error_message) # Add a response with incorrect expected params service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # The first call should throw and error as expected. with self.assertRaises(ClientError): self.client.list_objects(Bucket='foo') # The second call should throw an error for unexpected parameters with self.assertRaisesRegexp(StubResponseError, 'Expected parameters'): self.client.list_objects(Bucket='foo') def test_can_continue_to_call_after_expected_params_fail(self): service_response = {} expected_params = {'Bucket': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # Throw an error for unexpected parameters with self.assertRaises(StubResponseError): self.client.list_objects(Bucket='foo') # The stubber should still have the responses queued up # even though the original parameters did not match the expected ones. self.client.list_objects(Bucket='bar') self.stubber.assert_no_pending_responses() def test_still_relies_on_param_validation_with_expected_params(self): service_response = {} expected_params = {'Buck': 'bar'} self.stubber.add_response( 'list_objects', service_response, expected_params) self.stubber.activate() # Throw an error for invalid parameters with self.assertRaises(ParamValidationError): self.client.list_objects(Buck='bar')
@pytest.mark.parametrize("case", ALIAS_CASES) def test_can_use_original_name(case): session = botocore.session.get_session() _can_use_parameter_in_client_call(session, case, False) def _can_use_parameter_in_client_call(session, case, use_alias=True): client = session.create_client(case['service'], region_name='us-east-1', aws_access_key_id='foo', aws_secret_access_key='bar') stubber = Stubber(client) stubber.activate() operation = case['operation'] params = case.get('extra_args', {}) params = params.copy() param_name = case['original_name'] if use_alias: param_name = case['new_name'] params[param_name] = case['parameter_value'] stubbed_response = case.get('stubbed_response', {}) stubber.add_response(operation, stubbed_response) try: getattr(client, operation)(**params) except ParamValidationError as e: raise AssertionError( 'Expecting %s to be valid parameter for %s.%s but received ' '%s.' % (case['new_name'], case['service'], case['operation'], e))
class TestGuardDutyMetricsCollector(unittest.TestCase): def setUp(self): self.gdClient = boto3.client("guardduty") self.gdStubber = Stubber(self.gdClient) self.gdStubber.activate() self.botoSessionMock = MagicMock() self.botoSessionMock.client.return_value = self.gdClient def testCollectShouldReturnCurrentFindingsMetricFromSingleRegionWithSingleDetectorOnSuccess( self): # Mock GuardDuty self.gdStubber.add_response("list_detectors", {"DetectorIds": ["eu-detector-1"]}, {}) self.gdStubber.add_response( "get_findings_statistics", { "FindingStatistics": { "CountBySeverity": { "2.0": 1, "4.0": 2, "7.0": 3, } } }, { "DetectorId": "eu-detector-1", "FindingCriteria": { "Criterion": { "service.archived": { "Eq": ["false"] } } }, "FindingStatisticTypes": ["COUNT_BY_SEVERITY"] }) # Collect metrics with patch("boto3.session.Session", return_value=self.botoSessionMock): collector = GuardDutyMetricsCollector(regions=["eu-west-1"]) metrics = collector.collect() findingsMetric = metrics[0] self.assertEqual(findingsMetric.name, "aws_guardduty_current_findings") self.assertEqual(findingsMetric.type, "gauge") self.assertEqual(len(findingsMetric.samples), 3) self.assertEqual(findingsMetric.samples[0].value, 1) self.assertEqual(findingsMetric.samples[0].labels, { "region": "eu-west-1", "severity": "low" }) self.assertEqual(findingsMetric.samples[1].value, 2) self.assertEqual(findingsMetric.samples[1].labels, { "region": "eu-west-1", "severity": "medium" }) self.assertEqual(findingsMetric.samples[2].value, 3) self.assertEqual(findingsMetric.samples[2].labels, { "region": "eu-west-1", "severity": "high" }) scrapeErrorsMetric = metrics[1] self.assertEqual(scrapeErrorsMetric.name, "aws_guardduty_scrape_errors") self.assertEqual(scrapeErrorsMetric.type, "counter") self.assertEqual(len(scrapeErrorsMetric.samples), 1) self.assertEqual(scrapeErrorsMetric.samples[0].value, 0) self.assertEqual(scrapeErrorsMetric.samples[0].labels, {"region": "eu-west-1"}) self.gdStubber.assert_no_pending_responses() def testCollectShouldReturnCurrentFindingsMetricFromSingleRegionWithMultipleDetectorsOnSuccess( self): # Mock GuardDuty self.gdStubber.add_response( "list_detectors", {"DetectorIds": ["eu-detector-1", "eu-detector-2"]}, {}) self.gdStubber.add_response( "get_findings_statistics", { "FindingStatistics": { "CountBySeverity": { "2.0": 1, "4.0": 2, "7.0": 3, } } }, { "DetectorId": "eu-detector-1", "FindingCriteria": { "Criterion": { "service.archived": { "Eq": ["false"] } } }, "FindingStatisticTypes": ["COUNT_BY_SEVERITY"] }) self.gdStubber.add_response( "get_findings_statistics", { "FindingStatistics": { "CountBySeverity": { "2.5": 4, "5.2": 5, "8.1": 6, } } }, { "DetectorId": "eu-detector-2", "FindingCriteria": { "Criterion": { "service.archived": { "Eq": ["false"] } } }, "FindingStatisticTypes": ["COUNT_BY_SEVERITY"] }) # Collect metrics with patch("boto3.session.Session", return_value=self.botoSessionMock): collector = GuardDutyMetricsCollector(regions=["eu-west-1"]) metrics = collector.collect() findingsMetric = metrics[0] self.assertEqual(findingsMetric.name, "aws_guardduty_current_findings") self.assertEqual(findingsMetric.type, "gauge") self.assertEqual(len(findingsMetric.samples), 3) self.assertEqual(findingsMetric.samples[0].value, 5) self.assertEqual(findingsMetric.samples[0].labels, { "region": "eu-west-1", "severity": "low" }) self.assertEqual(findingsMetric.samples[1].value, 7) self.assertEqual(findingsMetric.samples[1].labels, { "region": "eu-west-1", "severity": "medium" }) self.assertEqual(findingsMetric.samples[2].value, 9) self.assertEqual(findingsMetric.samples[2].labels, { "region": "eu-west-1", "severity": "high" }) scrapeErrorsMetric = metrics[1] self.assertEqual(scrapeErrorsMetric.name, "aws_guardduty_scrape_errors") self.assertEqual(scrapeErrorsMetric.type, "counter") self.assertEqual(len(scrapeErrorsMetric.samples), 1) self.assertEqual(scrapeErrorsMetric.samples[0].value, 0) self.assertEqual(scrapeErrorsMetric.samples[0].labels, {"region": "eu-west-1"}) self.gdStubber.assert_no_pending_responses() def testCollectShouldReturnCurrentFindingsMetricFromMultipleRegionsOnSuccess( self): # Mock GuardDuty self.gdStubber.add_response("list_detectors", {"DetectorIds": ["eu-detector-1"]}, {}) self.gdStubber.add_response( "get_findings_statistics", { "FindingStatistics": { "CountBySeverity": { "2.0": 1, "4.0": 2, "7.0": 3, } } }, { "DetectorId": "eu-detector-1", "FindingCriteria": { "Criterion": { "service.archived": { "Eq": ["false"] } } }, "FindingStatisticTypes": ["COUNT_BY_SEVERITY"] }) self.gdStubber.add_response("list_detectors", {"DetectorIds": ["us-detector-1"]}, {}) self.gdStubber.add_response( "get_findings_statistics", { "FindingStatistics": { "CountBySeverity": { "2.5": 4, "5.2": 5, "8.1": 6, } } }, { "DetectorId": "us-detector-1", "FindingCriteria": { "Criterion": { "service.archived": { "Eq": ["false"] } } }, "FindingStatisticTypes": ["COUNT_BY_SEVERITY"] }) # Collect metrics with patch("boto3.session.Session", return_value=self.botoSessionMock): collector = GuardDutyMetricsCollector( regions=["eu-west-1", "us-east-1"]) metrics = collector.collect() findingsMetric = metrics[0] self.assertEqual(findingsMetric.name, "aws_guardduty_current_findings") self.assertEqual(findingsMetric.type, "gauge") self.assertEqual(len(findingsMetric.samples), 6) self.assertEqual(findingsMetric.samples[0].value, 1) self.assertEqual(findingsMetric.samples[0].labels, { "region": "eu-west-1", "severity": "low" }) self.assertEqual(findingsMetric.samples[1].value, 2) self.assertEqual(findingsMetric.samples[1].labels, { "region": "eu-west-1", "severity": "medium" }) self.assertEqual(findingsMetric.samples[2].value, 3) self.assertEqual(findingsMetric.samples[2].labels, { "region": "eu-west-1", "severity": "high" }) self.assertEqual(findingsMetric.samples[3].value, 4) self.assertEqual(findingsMetric.samples[3].labels, { "region": "us-east-1", "severity": "low" }) self.assertEqual(findingsMetric.samples[4].value, 5) self.assertEqual(findingsMetric.samples[4].labels, { "region": "us-east-1", "severity": "medium" }) self.assertEqual(findingsMetric.samples[5].value, 6) self.assertEqual(findingsMetric.samples[5].labels, { "region": "us-east-1", "severity": "high" }) scrapeErrorsMetric = metrics[1] self.assertEqual(scrapeErrorsMetric.name, "aws_guardduty_scrape_errors") self.assertEqual(scrapeErrorsMetric.type, "counter") self.assertEqual(len(scrapeErrorsMetric.samples), 2) self.assertEqual(scrapeErrorsMetric.samples[0].value, 0) self.assertEqual(scrapeErrorsMetric.samples[0].labels, {"region": "eu-west-1"}) self.assertEqual(scrapeErrorsMetric.samples[1].value, 0) self.assertEqual(scrapeErrorsMetric.samples[1].labels, {"region": "us-east-1"}) self.gdStubber.assert_no_pending_responses() def testCollectShouldSkipCurrentFindingsMetricOnFailingFetchingStatisticsFromARegionOutOfMultipleRegions( self): # Mock GuardDuty self.gdStubber.add_response("list_detectors", {"DetectorIds": ["eu-detector-1"]}, {}) self.gdStubber.add_client_error("get_findings_statistics") self.gdStubber.add_response("list_detectors", {"DetectorIds": ["us-detector-1"]}, {}) self.gdStubber.add_response( "get_findings_statistics", { "FindingStatistics": { "CountBySeverity": { "2.5": 4, "5.2": 5, "8.1": 6, } } }, { "DetectorId": "us-detector-1", "FindingCriteria": { "Criterion": { "service.archived": { "Eq": ["false"] } } }, "FindingStatisticTypes": ["COUNT_BY_SEVERITY"] }) # Collect metrics with patch("boto3.session.Session", return_value=self.botoSessionMock): collector = GuardDutyMetricsCollector( regions=["eu-west-1", "us-east-1"]) metrics = collector.collect() findingsMetric = metrics[0] self.assertEqual(findingsMetric.name, "aws_guardduty_current_findings") self.assertEqual(findingsMetric.type, "gauge") self.assertEqual(len(findingsMetric.samples), 3) self.assertEqual(findingsMetric.samples[0].value, 4) self.assertEqual(findingsMetric.samples[0].labels, { "region": "us-east-1", "severity": "low" }) self.assertEqual(findingsMetric.samples[1].value, 5) self.assertEqual(findingsMetric.samples[1].labels, { "region": "us-east-1", "severity": "medium" }) self.assertEqual(findingsMetric.samples[2].value, 6) self.assertEqual(findingsMetric.samples[2].labels, { "region": "us-east-1", "severity": "high" }) scrapeErrorsMetric = metrics[1] self.assertEqual(scrapeErrorsMetric.name, "aws_guardduty_scrape_errors") self.assertEqual(scrapeErrorsMetric.type, "counter") self.assertEqual(len(scrapeErrorsMetric.samples), 2) self.assertEqual(scrapeErrorsMetric.samples[0].value, 1) self.assertEqual(scrapeErrorsMetric.samples[0].labels, {"region": "eu-west-1"}) self.assertEqual(scrapeErrorsMetric.samples[1].value, 0) self.assertEqual(scrapeErrorsMetric.samples[1].labels, {"region": "us-east-1"}) self.gdStubber.assert_no_pending_responses() def testCollectShouldNeverDecreaseScrapeErrorsOnSubsequentCalls(self): # Mock GuardDuty self.gdStubber.add_response("list_detectors", {"DetectorIds": ["us-detector-1"]}, {}) self.gdStubber.add_client_error("get_findings_statistics") # Collect metrics with patch("boto3.session.Session", return_value=self.botoSessionMock): collector = GuardDutyMetricsCollector(regions=["us-east-1"]) metrics = collector.collect() findingsMetric = metrics[0] self.assertEqual(findingsMetric.name, "aws_guardduty_current_findings") self.assertEqual(findingsMetric.type, "gauge") self.assertEqual(len(findingsMetric.samples), 0) scrapeErrorsMetric = metrics[1] self.assertEqual(scrapeErrorsMetric.name, "aws_guardduty_scrape_errors") self.assertEqual(scrapeErrorsMetric.type, "counter") self.assertEqual(len(scrapeErrorsMetric.samples), 1) self.assertEqual(scrapeErrorsMetric.samples[0].value, 1) self.assertEqual(scrapeErrorsMetric.samples[0].labels, {"region": "us-east-1"}) with patch("boto3.session.Session", return_value=self.botoSessionMock): metrics = collector.collect() findingsMetric = metrics[0] self.assertEqual(findingsMetric.name, "aws_guardduty_current_findings") self.assertEqual(findingsMetric.type, "gauge") self.assertEqual(len(findingsMetric.samples), 0) scrapeErrorsMetric = metrics[1] self.assertEqual(scrapeErrorsMetric.name, "aws_guardduty_scrape_errors") self.assertEqual(scrapeErrorsMetric.type, "counter") self.assertEqual(len(scrapeErrorsMetric.samples), 1) self.assertEqual(scrapeErrorsMetric.samples[0].value, 2) self.assertEqual(scrapeErrorsMetric.samples[0].labels, {"region": "us-east-1"}) self.gdStubber.assert_no_pending_responses()
def test_event_bad(mocker): # Read test data event = utils.load_test_data(test_data + 'cis_1-3.json', my_region) iam_keys = { "AccessKeyMetadata": [ { "UserName": "******", "AccessKeyId": "AKIAADFHWEREFGFHSDDF", "Status": "Active", "CreateDate": "2015-05-22T14:43:16+00:00" }, { "UserName": "******", "AccessKeyId": "AKIAGHJGJFGHJFGETHFG", "Status": "Active", "CreateDate": "2020-05-15T15:20:04+00:00" } ] } # Mock the constructor. We don't need the session created mocker.patch('lib.awsapi_helpers.BotoSession.__init__', return_value=None) # create client and resource directly through boto3 iamc = boto3.client('iam', region_name=my_region) iamr = boto3.resource('iam', region_name=my_region) # stub the client iamc_s = Stubber(iamc) iamr_s = Stubber(iamr.meta.client) iamc_s.add_response( 'list_access_keys', iam_keys ) iam_keys['AccessKeyMetadata'][0]['Status'] = 'Inactive' iamc_s.add_response( 'list_access_keys', iam_keys ) iamc_s.activate() iamr_s.activate() # Mock the client and resource to return the clients we created mocker.patch('lib.awsapi_helpers.BotoSession.client', return_value=iamc) mocker.patch('lib.awsapi_helpers.BotoSession.resource', return_value=iamr) mocker.patch('lib.awsapi_helpers.AWSClient.postit', return_value=None) # Mock flush so we don't mocker.patch('lib.applogger.LogHandler.flush', return_value=None) # Mock Notifier init = mocker.patch('lib.sechub_findings.Finding.flag') update = mocker.patch('lib.sechub_findings.Finding.update_text') cis1314.lambda_handler(event, None) init.assert_called_with( 'INITIAL: "Deactivate unused keys over 90 days old" remediation started' ) update.assert_called_once_with( 'FAILED: "Deactivate unused keys over 90 days old" remediation failed. Please remediate manually' )
class TestStubber(unittest.TestCase): def setUp(self): self.event_emitter = hooks.HierarchicalEmitter() self.client = mock.Mock() self.client.meta.events = self.event_emitter self.client.meta.method_to_api_mapping.get.return_value = 'foo' self.stubber = Stubber(self.client) self.validate_parameters_mock = mock.Mock() self.validate_parameters_patch = mock.patch( 'botocore.stub.validate_parameters', self.validate_parameters_mock) self.validate_parameters_patch.start() def tearDown(self): self.validate_parameters_patch.stop() def emit_get_response_event(self, model=None, request_dict=None, signer=None, context=None): if model is None: model = mock.Mock() model.name = 'foo' handler, response = self.event_emitter.emit_until_response( event_name='before-call.myservice.foo', model=model, params=request_dict, request_signer=signer, context=context) return response def test_stubber_registers_events(self): self.event_emitter = mock.Mock() self.client.meta.events = self.event_emitter self.stubber.activate() # This just ensures that we register at the correct event # and nothing more self.event_emitter.register_first.assert_called_with( 'before-parameter-build.*.*', mock.ANY, unique_id=mock.ANY) self.event_emitter.register.assert_called_with( 'before-call.*.*', mock.ANY, unique_id=mock.ANY) def test_stubber_unregisters_events(self): self.event_emitter = mock.Mock() self.client.meta.events = self.event_emitter self.stubber.activate() self.stubber.deactivate() self.event_emitter.unregister.assert_any_call( 'before-parameter-build.*.*', mock.ANY, unique_id=mock.ANY) self.event_emitter.unregister.assert_any_call( 'before-call.*.*', mock.ANY, unique_id=mock.ANY) def test_add_response(self): response = {'foo': 'bar'} self.stubber.add_response('foo', response) with self.assertRaises(AssertionError): self.stubber.assert_no_pending_responses() def test_add_response_fails_when_missing_client_method(self): del self.client.foo with self.assertRaises(ValueError): self.stubber.add_response('foo', {}) def test_validates_service_response(self): self.stubber.add_response('foo', {}) self.assertTrue(self.validate_parameters_mock.called) def test_validate_ignores_response_metadata(self): service_response = {'ResponseMetadata': {'foo': 'bar'}} service_model = ServiceModel({ 'documentation': '', 'operations': { 'foo': { 'name': 'foo', 'input': {'shape': 'StringShape'}, 'output': {'shape': 'StringShape'} } }, 'shapes': { 'StringShape': {'type': 'string'} } }) op_name = service_model.operation_names[0] output_shape = service_model.operation_model(op_name).output_shape self.client.meta.service_model = service_model self.stubber.add_response('TestOperation', service_response) self.validate_parameters_mock.assert_called_with( {}, output_shape) # Make sure service response hasn't been mutated self.assertEqual( service_response, {'ResponseMetadata': {'foo': 'bar'}}) def test_validates_on_empty_output_shape(self): service_model = ServiceModel({ 'documentation': '', 'operations': { 'foo': { 'name': 'foo' } } }) self.client.meta.service_model = service_model with self.assertRaises(ParamValidationError): self.stubber.add_response('TestOperation', {'foo': 'bar'}) def test_get_response(self): service_response = {'bar': 'baz'} self.stubber.add_response('foo', service_response) self.stubber.activate() response = self.emit_get_response_event() self.assertEqual(response[1], service_response) self.assertEqual(response[0].status_code, 200) def test_get_client_error_response(self): error_code = "foo" service_message = "bar" self.stubber.add_client_error('foo', error_code, service_message) self.stubber.activate() response = self.emit_get_response_event() self.assertEqual(response[1]['Error']['Message'], service_message) self.assertEqual(response[1]['Error']['Code'], error_code) def test_get_response_errors_with_no_stubs(self): self.stubber.activate() with self.assertRaises(StubResponseError): self.emit_get_response_event() def test_assert_no_responses_remaining(self): self.stubber.add_response('foo', {}) with self.assertRaises(AssertionError): self.stubber.assert_no_pending_responses()
class TestIndex(TestCase): def setUp(self): self.requests_mock = responses.RequestsMock( assert_all_requests_are_fired=False) self.requests_mock.start() # Create a dummy S3 client that (hopefully) can't do anything. self.s3_client = boto3.client( 's3', config=Config(signature_version=UNSIGNED)) self.s3_client_patcher = patch(__name__ + '.index.make_s3_client', return_value=self.s3_client) self.s3_client_patcher.start() self.s3_stubber = Stubber(self.s3_client) self.s3_stubber.activate() self.env_patcher = patch.dict( os.environ, { 'ES_HOST': 'example.com', 'AWS_ACCESS_KEY_ID': 'test_key', 'AWS_SECRET_ACCESS_KEY': 'test_secret', 'AWS_DEFAULT_REGION': 'ng-north-1', }) self.env_patcher.start() def tearDown(self): self.env_patcher.stop() self.s3_stubber.assert_no_pending_responses() self.s3_stubber.deactivate() self.s3_client_patcher.stop() self.requests_mock.stop() def _get_contents(self, name, ext): return index.get_contents( 'test-bucket', name, ext, etag='etag', version_id=None, s3_client=self.s3_client, size=123, ) def test_infer_extensions(self): """ensure we are guessing file types well""" # parquet assert index.infer_extensions("s3/some/file.c000", ".c000") == ".parquet", \ "Expected .c0000 to infer as .parquet" # parquet, nonzero part number assert index.infer_extensions("s3/some/file.c001", ".c001") == ".parquet", \ "Expected .c0001 to infer as .parquet" # -c0001 file assert index.infer_extensions("s3/some/file-c0001", "") == ".parquet", \ "Expected -c0001 to infer as .parquet" # -c00111 file (should never happen) assert index.infer_extensions("s3/some/file-c000121", "") == "", \ "Expected -c000121 not to infer as .parquet" # .txt file, should be unchanged assert index.infer_extensions("s3/some/file-c0000.txt", ".txt") == ".txt", \ "Expected .txt to infer as .txt" def test_delete_event(self): """ Check that the indexer doesn't blow up on delete events. """ # don't mock head or get; they should never be called for deleted objects self._test_index_event("ObjectRemoved:Delete", mock_head=False, mock_object=False) def test_delete_marker_event(self): """ Common event in versioned; buckets, should no-op """ # don't mock head or get; this event should never call them self._test_index_event( "ObjectRemoved:DeleteMarkerCreated", # we should never call Elastic in this case mock_elastic=False, mock_head=False, mock_object=False) def test_test_event(self): """ Check that the indexer does not barf when it gets an S3 test notification. """ event = { "Records": [{ "body": json.dumps({ "Message": json.dumps({ "Service": "Amazon S3", "Event": "s3:TestEvent", "Time": "2014-10-13T15:57:02.089Z", "Bucket": "test-bucket", "RequestId": "5582815E1AEA5ADF", "HostId": "8cLeGAmw098X5cv4Zkwcmo8vvZa3eH3eKxsPzbB9wrR+YstdA6Knx4Ip8EXAMPLE" }) }) }] } index.handler(event, None) def test_index_file(self): """test indexing a single file""" self._test_index_event("ObjectCreated:Put") @patch(__name__ + '.index.get_contents') def test_index_exception(self, get_mock): """test indexing a single file that throws an exception""" class ContentException(Exception): pass get_mock.side_effect = ContentException("Unable to get contents") with pytest.raises(ContentException): # get_mock already mocks get_object, so don't mock it in _test_index_event self._test_index_event("ObjectCreated:Put", mock_object=False) def _test_index_event(self, event_name, mock_elastic=True, mock_head=True, mock_object=True): """ Reusable helper function to test indexing a single text file. """ assert event_name in RECORDS, f"unexpected event: {event_name}" records = { "Records": [{ "body": json.dumps({ "Message": json.dumps({"Records": [RECORDS[event_name]]}) }) }] } now = index.now_like_boto3() metadata = { 'helium': json.dumps({ 'comment': 'blah', 'user_meta': { 'foo': 'bar' }, 'x': 'y' }) } if mock_head: self.s3_stubber.add_response(method='head_object', service_response={ 'Metadata': metadata, 'ContentLength': 100, 'LastModified': now, }, expected_params={ 'Bucket': 'test-bucket', 'Key': 'hello world.txt', 'IfMatch': '123456', }) if mock_object: self.s3_stubber.add_response( method='get_object', service_response={ 'Metadata': metadata, 'ContentLength': 100, 'LastModified': now, 'Body': BytesIO(b'Hello World!'), }, expected_params={ 'Bucket': 'test-bucket', 'Key': 'hello world.txt', 'IfMatch': '123456', 'Range': f'bytes=0-{index.ELASTIC_LIMIT_BYTES}', }) def es_callback(request): response_key = 'delete' if event_name == index.OBJECT_DELETE else 'index' actions = [json.loads(line) for line in request.body.splitlines()] expected = [{ response_key: { '_index': 'test-bucket', '_type': '_doc', '_id': 'hello world.txt:None' } }, { 'comment': 'blah', 'content': '' if not mock_object else 'Hello World!', 'etag': '123456', 'event': event_name, 'ext': '.txt', 'key': 'hello world.txt', 'last_modified': now.isoformat(), 'meta_text': 'blah {"x": "y"} {"foo": "bar"}', 'size': 100, 'target': '', 'updated': ANY, 'version_id': None }] if response_key == 'delete': # delete events do not include request body expected.pop() assert actions == expected, "Unexpected request to ElasticSearch" response = {'items': [{response_key: {'status': 200}}]} return (200, {}, json.dumps(response)) if mock_elastic: self.requests_mock.add_callback(responses.POST, 'https://example.com:443/_bulk', callback=es_callback, content_type='application/json') index.handler(records, MockContext()) def test_unsupported_contents(self): assert self._get_contents('foo.exe', '.exe') == "" assert self._get_contents('foo.exe.gz', '.exe.gz') == "" def test_get_plain_text(self): self.s3_stubber.add_response( method='get_object', service_response={ 'Metadata': {}, 'ContentLength': 123, 'Body': BytesIO(b'Hello World!\nThere is more to know.'), }, expected_params={ 'Bucket': 'test-bucket', 'Key': 'foo.txt', 'IfMatch': 'etag', 'Range': f'bytes=0-{index.ELASTIC_LIMIT_BYTES}', }) contents = index.get_plain_text('test-bucket', 'foo.txt', compression=None, etag='etag', version_id=None, s3_client=self.s3_client, size=123) assert contents == "Hello World!\nThere is more to know." def test_text_contents(self): self.s3_stubber.add_response( method='get_object', service_response={ 'Metadata': {}, 'ContentLength': 123, 'Body': BytesIO(b'Hello World!'), }, expected_params={ 'Bucket': 'test-bucket', 'Key': 'foo.txt', 'IfMatch': 'etag', 'Range': f'bytes=0-{index.ELASTIC_LIMIT_BYTES}', }) assert self._get_contents('foo.txt', '.txt') == "Hello World!" def test_gzipped_text_contents(self): self.s3_stubber.add_response( method='get_object', service_response={ 'Metadata': {}, 'ContentLength': 123, 'Body': BytesIO(compress(b'Hello World!')), }, expected_params={ 'Bucket': 'test-bucket', 'Key': 'foo.txt.gz', 'IfMatch': 'etag', 'Range': f'bytes=0-{index.ELASTIC_LIMIT_BYTES}', }) assert self._get_contents('foo.txt.gz', '.txt.gz') == "Hello World!" def test_notebook_contents(self): notebook = (BASE_DIR / 'normal.ipynb').read_bytes() self.s3_stubber.add_response(method='get_object', service_response={ 'Metadata': {}, 'ContentLength': 123, 'Body': BytesIO(notebook), }, expected_params={ 'Bucket': 'test-bucket', 'Key': 'foo.ipynb', 'IfMatch': 'etag', }) assert "model.fit" in self._get_contents('foo.ipynb', '.ipynb') def test_gzipped_notebook_contents(self): notebook = compress((BASE_DIR / 'normal.ipynb').read_bytes()) self.s3_stubber.add_response(method='get_object', service_response={ 'Metadata': {}, 'ContentLength': 123, 'Body': BytesIO(notebook), }, expected_params={ 'Bucket': 'test-bucket', 'Key': 'foo.ipynb.gz', 'IfMatch': 'etag', }) assert "Model results visualization" in self._get_contents( 'foo.ipynb.gz', '.ipynb.gz') def test_parquet_contents(self): parquet = (BASE_DIR / 'amazon-reviews-1000.snappy.parquet').read_bytes() self.s3_stubber.add_response(method='get_object', service_response={ 'Metadata': {}, 'ContentLength': 123, 'Body': BytesIO(parquet), }, expected_params={ 'Bucket': 'test-bucket', 'Key': 'foo.parquet', 'IfMatch': 'etag', }) contents = self._get_contents('foo.parquet', '.parquet') size = len(contents.encode('utf-8', 'ignore')) assert size <= index.ELASTIC_LIMIT_BYTES # spot check for contents assert "This is not even worth the money." in contents assert "As for results; I felt relief almost immediately." in contents assert "R2LO11IPLTDQDX" in contents # see PRE conditions in conftest.py @pytest.mark.extended def test_parquet_extended(self): directory = (BASE_DIR / 'amazon-reviews-pds') files = directory.glob('**/*.parquet') for f in files: print(f"Testing {f}") parquet = f.read_bytes() self.s3_stubber.add_response(method='get_object', service_response={ 'Metadata': {}, 'ContentLength': 123, 'Body': BytesIO(parquet), }, expected_params={ 'Bucket': 'test-bucket', 'Key': 'foo.parquet', 'IfMatch': 'etag', })