def authenticate(self): # This value controls how long we sleep # between reauthenticating and getting a new set of creds # eventually this gets set by aws response self.flush_wait_time = 1800 if options.cloudtrail_arn not in [ '<cloudtrail_arn>', 'cloudtrail_arn' ]: client = boto3.client('sts', aws_access_key_id=options.accesskey, aws_secret_access_key=options.secretkey) response = client.assume_role( RoleArn=options.cloudtrail_arn, RoleSessionName='MozDef-CloudTrail-Reader', ) role_creds = { 'aws_access_key_id': response['Credentials']['AccessKeyId'], 'aws_secret_access_key': response['Credentials']['SecretAccessKey'], 'aws_session_token': response['Credentials']['SessionToken'] } current_time = toUTC(datetime.now()) # Let's remove 3 seconds from the flush wait time just in case self.flush_wait_time = (response['Credentials']['Expiration'] - current_time).seconds - 3 else: role_creds = {} role_creds['region_name'] = options.region self.s3_client = boto3.client('s3', **get_aws_credentials(**role_creds))
def run(self): while True: try: records = self.sqs_queue.receive_messages( MaxNumberOfMessages=options.prefetch) for msg in records: body_message = msg.body event = json.loads(body_message) if not event['Message']: logger.error( 'Invalid message format for cloudtrail SQS messages' ) logger.error('Malformed Message: %r' % body_message) continue if event['Message'] == 'CloudTrail validation message.': # We don't care about these messages continue message_json = json.loads(event['Message']) if 's3ObjectKey' not in message_json: logger.error( 'Invalid message format, expecting an s3ObjectKey in Message' ) logger.error('Malformed Message: %r' % body_message) continue s3_log_files = message_json['s3ObjectKey'] for log_file in s3_log_files: logger.debug('Downloading and parsing ' + log_file) s3_obj = self.s3_client.get_object( Bucket=message_json['s3Bucket'], Key=log_file) events = self.parse_s3_file(s3_obj) for event in events: self.on_message(event) msg.delete() except (SSLEOFError, SSLError, socket.error): logger.info('Received network related error...reconnecting') time.sleep(5) self.sqs_queue = connect_sqs( task_exchange=options.taskexchange, **get_aws_credentials(options.region, options.accesskey, options.secretkey)) time.sleep(options.sleep_time)
def main(): if hasUWSGI: logger.info("started as uwsgi mule {0}".format(uwsgi.mule_id())) else: logger.info('started without uwsgi') if options.mqprotocol not in ('sqs'): logger.error('Can only process SQS queues, terminating') sys.exit(1) sqs_queue = connect_sqs( task_exchange=options.taskexchange, **get_aws_credentials( options.region, options.accesskey, options.secretkey)) # consume our queue taskConsumer(sqs_queue, es, options).run()