def test_poller_processor_handler(historical_sqs, historical_role, mock_lambda_environment, security_groups, swag_accounts): """Test the Poller's processing component that tasks the collector.""" # Mock this so it returns a `NextToken`: def mock_describe_security_groups(**kwargs): from cloudaux.aws.ec2 import describe_security_groups # Did we receive a NextToken? (this will happen on the second run through to verify that # this logic is being reached: if kwargs.get('NextToken'): assert kwargs['NextToken'] == 'MOARRESULTS' result = describe_security_groups(**kwargs) result['NextToken'] = 'MOARRESULTS' return result patch_sgs = patch('historical.security_group.poller.describe_security_groups', mock_describe_security_groups) patch_sgs.start() from historical.security_group.poller import poller_processor_handler as handler from historical.common import cloudwatch # Create the events and SQS records: messages = make_poller_events() event = json.loads(json.dumps(RecordsFactory(records=messages), default=serialize)) # Run the poller handler: handler(event, mock_lambda_environment) # Need to ensure that 3 total SGs were added into SQS: sqs = boto3.client("sqs", region_name="us-east-1") queue_url = get_queue_url(os.environ['POLLER_QUEUE_NAME']) messages = sqs.receive_message(QueueUrl=queue_url, MaxNumberOfMessages=10)['Messages'] assert len(messages) == 3 # Verify that the region is properly propagated through, and that we got the collected data: for msg in messages: body = json.loads(msg['Body']) assert cloudwatch.get_region(body) == 'us-east-1' assert body['detail']['collected']['OwnerId'] == '123456789012' assert not body['detail']['collected'].get('ResponseMetadata') # Now, verify that the pagination was sent in properly to SQS tasker queue: queue_url = get_queue_url(os.environ['POLLER_TASKER_QUEUE_NAME']) messages = sqs.receive_message(QueueUrl=queue_url, MaxNumberOfMessages=10)['Messages'] assert len(messages) == 1 assert json.loads(messages[0]['Body'])['NextToken'] == 'MOARRESULTS' # Re-run the poller: messages[0]['body'] = messages[0]['Body'] # Need to change the casing handler({'Records': messages}, mock_lambda_environment) patch_sgs.stop()
def poller_processor_handler(event, context): """ Historical S3 Poller Processor. This will receive events from the Poller Tasker, and will list all objects of a given technology for an account/region pair. This will generate `polling events` which simulate changes. These polling events contain configuration data such as the account/region defining where the collector should attempt to gather data from. """ log.debug('[@] Running Poller...') queue_url = get_queue_url(os.environ.get('POLLER_QUEUE_NAME', 'HistoricalS3Poller')) records = deserialize_records(event['Records']) for record in records: # Skip accounts that have role assumption errors: try: # List all buckets in the account: all_buckets = list_buckets(account_number=record['account_id'], assume_role=HISTORICAL_ROLE, session_name="historical-cloudwatch-s3list", region=record['region'])["Buckets"] events = [s3_polling_schema.serialize_me(record['account_id'], bucket) for bucket in all_buckets] produce_events(events, queue_url) except ClientError as e: log.error('[X] Unable to generate events for account. Account Id: {account_id} Reason: {reason}'.format( account_id=record['account_id'], reason=e)) log.debug('[@] Finished generating polling events. Events Created: {}'.format(len(record['account_id'])))
def poller_tasker_handler(event, context): """ Historical S3 Poller Tasker. The Poller is run at a set interval in order to ensure that changes do not go undetected by Historical. Historical pollers generate `polling events` which simulate changes. These polling events contain configuration data such as the account/region defining where the collector should attempt to gather data from. This is the entry point. This will task subsequent Poller lambdas to list all of a given resource in a select few AWS accounts. """ log.debug('[@] Running Poller Tasker...') queue_url = get_queue_url(os.environ.get('POLLER_TASKER_QUEUE_NAME', 'HistoricalS3PollerTasker')) poller_task_schema = HistoricalPollerTaskEventModel() events = [poller_task_schema.serialize_me(account['id'], CURRENT_REGION) for account in get_historical_accounts()] try: produce_events(events, queue_url) except ClientError as e: log.error('[X] Unable to generate poller tasker events! Reason: {reason}'.format(reason=e)) log.debug('[@] Finished tasking the pollers.')
def poller_tasker_handler(event, context): # pylint: disable=W0613 """ Historical Security Group Poller Tasker. The Poller is run at a set interval in order to ensure that changes do not go undetected by Historical. Historical pollers generate `polling events` which simulate changes. These polling events contain configuration data such as the account/region defining where the collector should attempt to gather data from. This is the entry point. This will task subsequent Poller lambdas to list all of a given resource in a select few AWS accounts. """ LOG.debug('[@] Running Poller Tasker...') queue_url = get_queue_url( os.environ.get('POLLER_TASKER_QUEUE_NAME', 'HistoricalSecurityGroupPollerTasker')) poller_task_schema = HistoricalPollerTaskEventModel() events = [] for account in get_historical_accounts(): for region in POLL_REGIONS: events.append( poller_task_schema.serialize_me(account['id'], region)) try: produce_events(events, queue_url, randomize_delay=RANDOMIZE_POLLER) except ClientError as exc: LOG.error( f'[X] Unable to generate poller tasker events! Reason: {exc}') LOG.debug('[@] Finished tasking the pollers.')
def handler(event, context): """ Historical Security Group Poller. This Poller is run at a set interval in order to ensure that changes do not go undetected by historical. Historical pollers generate `polling events` which simulate changes. These polling events contain configuration data such as the account/region defining where the collector should attempt to gather data from. """ log.debug('Running poller. Configuration: {}'.format(event)) queue_url = get_queue_url( os.environ.get('POLLER_QUEUE_NAME', 'HistoricalSecurityGroupPoller')) for account in get_historical_accounts(): for region in POLL_REGIONS: try: groups = describe_security_groups(account_number=account['id'], assume_role=HISTORICAL_ROLE, region=region) events = [ security_group_polling_schema.serialize(account['id'], g) for g in groups['SecurityGroups'] ] produce_events(events, queue_url) log.debug('Finished generating polling events. Account: {}/{} ' 'Events Created: {}'.format(account['id'], region, len(events))) except ClientError as e: log.warning( 'Unable to generate events for account/region. Account Id/Region: {account_id}/{region}' ' Reason: {reason}'.format(account_id=account['id'], region=region, reason=e))
def poller_processor_handler(event, context): # pylint: disable=W0613 """ Historical Security Group Poller Processor. This will receive events from the Poller Tasker, and will list all objects of a given technology for an account/region pair. This will generate `polling events` which simulate changes. These polling events contain configuration data such as the account/region defining where the collector should attempt to gather data from. """ LOG.debug('[@] Running Poller...') queue_url = get_queue_url(os.environ.get('POLLER_QUEUE_NAME', 'HistoricalVPCPoller')) records = deserialize_records(event['Records']) for record in records: # Skip accounts that have role assumption errors: try: vpcs = describe_vpcs( account_number=record['account_id'], assume_role=HISTORICAL_ROLE, region=record['region'] ) events = [VPC_POLLING_SCHEMA.serialize(record['account_id'], v) for v in vpcs] produce_events(events, queue_url, randomize_delay=RANDOMIZE_POLLER) LOG.debug(f"[@] Finished generating polling events. Account: {record['account_id']}/{record['region']} " f"Events Created: {len(events)}") except ClientError as exc: LOG.error(f"[X] Unable to generate events for account/region. Account Id/Region: {record['account_id']}" f"/{record['region']} Reason: {exc}")
def test_poller(historical_role, buckets, mock_lambda_environment, historical_sqs, swag_accounts): from historical.s3.poller import handler handler({}, None) # Need to ensure that 51 total buckets were added into SQS: sqs = boto3.client("sqs", region_name="us-east-1") queue_url = get_queue_url(os.environ['POLLER_QUEUE_NAME']) all_buckets = {"SWAG": True} for i in range(0, 50): all_buckets["testbucket{}".format(i)] = True # Loop through the queue and make sure all buckets are accounted for: for i in range(0, 6): messages = sqs.receive_message(QueueUrl=queue_url, MaxNumberOfMessages=10)['Messages'] message_ids = [] for m in messages: message_ids.append({ "Id": m['MessageId'], "ReceiptHandle": m['ReceiptHandle'] }) data = s3_polling_schema.loads(m['Body']).data assert all_buckets[data["detail"]["request_parameters"] ["bucket_name"]] assert datetime.strptime( data["detail"]["request_parameters"]["creation_date"], '%Y-%m-%dT%H:%M:%SZ') assert data["detail"]["event_source"] == "historical.s3.poller" # Remove from the dict (at the end, there should be 0 items left) del all_buckets[data["detail"]["request_parameters"] ["bucket_name"]] sqs.delete_message_batch(QueueUrl=queue_url, Entries=message_ids) assert len(all_buckets) == 0 # Check that an exception raised doesn't break things: import historical.s3.poller def mocked_poller(account, stream): raise ClientError({"Error": { "Message": "", "Code": "AccessDenied" }}, "sts:AssumeRole") old_method = historical.s3.poller.produce_events # For pytest inter-test issues... historical.s3.poller.produce_events = mocked_poller handler({}, None) historical.s3.poller.produce_events = old_method
def test_poller(historical_sqs, historical_role, mock_lambda_environment, vpcs, swag_accounts): from historical.vpc.poller import handler handler({}, None) # Need to ensure that 2 total VPCs were added into SQS: sqs = boto3.client("sqs", region_name="us-east-1") queue_url = get_queue_url(os.environ['POLLER_QUEUE_NAME']) messages = sqs.receive_message(QueueUrl=queue_url, MaxNumberOfMessages=10)['Messages'] assert len(messages) == 2
def make_poller_events(): """A sort-of fixture to make polling events for tests.""" from historical.security_group.poller import poller_tasker_handler as handler handler({}, None) # Need to ensure that all of the accounts and regions were properly tasked (only 1 region for S3): sqs = boto3.client("sqs", region_name="us-east-1") queue_url = get_queue_url(os.environ['POLLER_TASKER_QUEUE_NAME']) messages = sqs.receive_message(QueueUrl=queue_url, MaxNumberOfMessages=10)['Messages'] # 'Body' needs to be made into 'body' for proper parsing later: for m in messages: m['body'] = m.pop('Body') return messages
def test_poller_processor_handler(historical_sqs, historical_role, mock_lambda_environment, security_groups, swag_accounts): from historical.security_group.poller import poller_processor_handler as handler # Create the events and SQS records: messages = make_poller_events() event = json.loads(json.dumps(RecordsFactory(records=messages), default=serialize)) # Run the collector: handler(event, mock_lambda_environment) # Need to ensure that 3 total SGs were added into SQS: sqs = boto3.client("sqs", region_name="us-east-1") queue_url = get_queue_url(os.environ['POLLER_QUEUE_NAME']) messages = sqs.receive_message(QueueUrl=queue_url, MaxNumberOfMessages=10)['Messages'] assert len(messages) == 3
def poller_processor_handler(event, context): """ Historical Security Group Poller Processor. This will receive events from the Poller Tasker, and will list all objects of a given technology for an account/region pair. This will generate `polling events` which simulate changes. These polling events contain configuration data such as the account/region defining where the collector should attempt to gather data from. """ log.debug('[@] Running Poller...') queue_url = get_queue_url( os.environ.get('POLLER_QUEUE_NAME', 'HistoricalSecurityGroupPoller')) records = deserialize_records(event['Records']) for record in records: # Skip accounts that have role assumption errors: try: groups = describe_security_groups( account_number=record['account_id'], assume_role=HISTORICAL_ROLE, region=record['region']) events = [ security_group_polling_schema.serialize( record['account_id'], g) for g in groups['SecurityGroups'] ] produce_events(events, queue_url) log.debug('[@] Finished generating polling events. Account: {}/{} ' 'Events Created: {}'.format(record['account_id'], record['region'], len(events))) except ClientError as e: log.error( '[X] Unable to generate events for account/region. Account Id/Region: {account_id}/{region}' ' Reason: {reason}'.format(account_id=record['account_id'], region=record['region'], reason=e)) log.debug('[@] Finished generating polling events. Events Created: {}'. format(len(record['account_id'])))
def handler(event, context): """ Historical S3 Poller. This poller is run at a set interval in order to ensure that changes do not go undetected by historical. Historical pollers generate `polling events` which simulate changes. These polling events contain configuration data such as the account/region defining where the collector should attempt to gather data from. """ log.debug('Running poller. Configuration: {}'.format(event)) queue_url = get_queue_url( os.environ.get('POLLER_QUEUE_NAME', 'HistoricalS3Poller')) for account in get_historical_accounts(): # Skip accounts that have role assumption errors: try: # List all buckets in the account: all_buckets = list_buckets( account_number=account['id'], assume_role=HISTORICAL_ROLE, session_name="historical-cloudwatch-s3list", region=CURRENT_REGION)["Buckets"] events = [ s3_polling_schema.serialize_me(account['id'], bucket) for bucket in all_buckets ] produce_events(events, queue_url) except ClientError as e: log.warning( 'Unable to generate events for account. AccountId: {account_id} Reason: {reason}' .format(account_id=account['id'], reason=e)) log.debug( 'Finished generating polling events. Events Created: {}'.format( len(account['id'])))
def poller_processor_handler(event, context): # pylint: disable=W0613 """ Historical Security Group Poller Processor. This will receive events from the Poller Tasker, and will list all objects of a given technology for an account/region pair. This will generate `polling events` which simulate changes. These polling events contain configuration data such as the account/region defining where the collector should attempt to gather data from. """ LOG.debug('[@] Running Poller...') collector_poller_queue_url = get_queue_url( os.environ.get('POLLER_QUEUE_NAME', 'HistoricalSecurityGroupPoller')) takser_queue_url = get_queue_url( os.environ.get('POLLER_TASKER_QUEUE_NAME', 'HistoricalSecurityGroupPollerTasker')) poller_task_schema = HistoricalPollerTaskEventModel() records = deserialize_records(event['Records']) for record in records: # Skip accounts that have role assumption errors: try: # Did we get a NextToken? if record.get('NextToken'): LOG.debug( f"[@] Received pagination token: {record['NextToken']}") groups = describe_security_groups( account_number=record['account_id'], assume_role=HISTORICAL_ROLE, region=record['region'], MaxResults=200, NextToken=record['NextToken']) else: groups = describe_security_groups( account_number=record['account_id'], assume_role=HISTORICAL_ROLE, region=record['region'], MaxResults=200) # FIRST THINGS FIRST: Did we get a `NextToken`? If so, we need to enqueue that ASAP because # 'NextToken`s expire in 60 seconds! if groups.get('NextToken'): logging.debug( f"[-->] Pagination required {groups['NextToken']}. Tasking continuation." ) produce_events([ poller_task_schema.serialize_me( record['account_id'], record['region'], next_token=groups['NextToken']) ], takser_queue_url) # Task the collector to perform all the DDB logic -- this will pass in the collected data to the # collector in very small batches. events = [ SECURITY_GROUP_POLLING_SCHEMA.serialize( record['account_id'], g, record['region']) for g in groups['SecurityGroups'] ] produce_events(events, collector_poller_queue_url, batch_size=3) LOG.debug( f"[@] Finished generating polling events. Account: {record['account_id']}/{record['region']} " f"Events Created: {len(events)}") except ClientError as exc: LOG.error( f"[X] Unable to generate events for account/region. Account Id/Region: {record['account_id']}" f"/{record['region']} Reason: {exc}")
def test_poller_processor_handler(historical_role, buckets, mock_lambda_environment, historical_sqs, swag_accounts): """Test the Poller's processing component that tasks the collector.""" from historical.s3.poller import poller_processor_handler as handler # Create the events and SQS records: messages = make_poller_events() event = json.loads( json.dumps(RecordsFactory(records=messages), default=serialize)) # Run the collector: handler(event, None) # Need to ensure that 51 total buckets were added into SQS: sqs = boto3.client("sqs", region_name="us-east-1") queue_url = get_queue_url(os.environ['POLLER_QUEUE_NAME']) all_buckets = {"SWAG": True} for i in range(0, 50): all_buckets[f"testbucket{i}"] = True # Loop through the queue and make sure all buckets are accounted for: for i in range(0, 6): messages = sqs.receive_message(QueueUrl=queue_url, MaxNumberOfMessages=10)['Messages'] message_ids = [] for msg in messages: message_ids.append({ "Id": msg['MessageId'], "ReceiptHandle": msg['ReceiptHandle'] }) data = S3_POLLING_SCHEMA.loads(msg['Body']).data assert all_buckets[data["detail"]["request_parameters"] ["bucket_name"]] assert datetime.strptime( data["detail"]["request_parameters"]["creation_date"], '%Y-%m-%dT%H:%M:%SZ') assert data["detail"]["event_source"] == "historical.s3.poller" # Remove from the dict (at the end, there should be 0 items left) del all_buckets[data["detail"]["request_parameters"] ["bucket_name"]] sqs.delete_message_batch(QueueUrl=queue_url, Entries=message_ids) assert not all_buckets # Check that an exception raised doesn't break things: import historical.s3.poller def mocked_poller(account, stream, randomize_delay=0): # pylint: disable=W0613 raise ClientError({"Error": { "Message": "", "Code": "AccessDenied" }}, "sts:AssumeRole") old_method = historical.s3.poller.produce_events # For pytest inter-test issues... historical.s3.poller.produce_events = mocked_poller handler(event, None) historical.s3.poller.produce_events = old_method