def test_submission_namespace(datastore, sio): submission_queue = CommsQueue('submissions', private=True) monitoring = get_random_id() ingested = random_model_obj(SubmissionMessage).as_primitives() ingested['msg_type'] = "SubmissionIngested" received = random_model_obj(SubmissionMessage).as_primitives() received['msg_type'] = "SubmissionReceived" queued = random_model_obj(SubmissionMessage).as_primitives() queued['msg_type'] = "SubmissionQueued" started = random_model_obj(SubmissionMessage).as_primitives() started['msg_type'] = "SubmissionStarted" test_res_array = [] @sio.on('monitoring', namespace='/submissions') def on_monitoring(data): # Confirmation that we are waiting for status messages test_res_array.append(('on_monitoring', data == monitoring)) @sio.on('SubmissionIngested', namespace='/submissions') def on_submission_ingested(data): test_res_array.append( ('on_submission_ingested', data == ingested['msg'])) @sio.on('SubmissionReceived', namespace='/submissions') def on_submission_received(data): test_res_array.append( ('on_submission_received', data == received['msg'])) @sio.on('SubmissionQueued', namespace='/submissions') def on_submission_queued(data): test_res_array.append(('on_submission_queued', data == queued['msg'])) @sio.on('SubmissionStarted', namespace='/submissions') def on_submission_started(data): test_res_array.append( ('on_submission_started', data == started['msg'])) try: sio.emit('monitor', monitoring, namespace='/submissions') sio.sleep(1) submission_queue.publish(ingested) submission_queue.publish(received) submission_queue.publish(queued) submission_queue.publish(started) start_time = time.time() while len(test_res_array) < 5 and time.time() - start_time < 5: sio.sleep(0.1) assert len(test_res_array) == 5 for test, result in test_res_array: if not result: pytest.fail(f"{test} failed.") finally: sio.disconnect()
def save_alert(datastore, counter, logger, alert, psid): def create_alert(): msg_type = "AlertCreated" datastore.alert.save(alert['alert_id'], alert) logger.info(f"Alert {alert['alert_id']} has been created.") counter.increment('created') ret_val = 'create' return msg_type, ret_val if psid: try: msg_type = "AlertUpdated" perform_alert_update(datastore, logger, alert) counter.increment('updated') ret_val = 'update' except AlertMissingError as e: logger.info( f"{str(e)}. Creating a new alert [{alert['alert_id']}]...") msg_type, ret_val = create_alert() else: msg_type, ret_val = create_alert() msg = AlertMessage({ "msg": alert, "msg_type": msg_type, "sender": "alerter" }) CommsQueue('alerts').publish(msg.as_primitives()) return ret_val
def __init__(self, config=None): super().__init__('assemblyline.heartbeat_manager') self.config = config or forge.get_config() self.datastore = forge.get_datastore() self.metrics_queue = CommsQueue(METRICS_QUEUE) self.scheduler = BackgroundScheduler(daemon=True) self.hm = HeartbeatFormatter("heartbeat_manager", self.log, config=self.config) self.counters_lock = Lock() self.counters = {} self.rolling_window = {} self.window_ttl = {} self.ttl = self.config.core.metrics.export_interval * 2 self.window_size = int(60 / self.config.core.metrics.export_interval) if self.window_size != 60 / self.config.core.metrics.export_interval: self.log.warning( "Cannot calculate a proper window size for reporting heartbeats. " "Metrics reported during hearbeat will be wrong.") if self.config.core.metrics.apm_server.server_url is not None: self.log.info( f"Exporting application metrics to: {self.config.core.metrics.apm_server.server_url}" ) elasticapm.instrument() self.apm_client = elasticapm.Client( server_url=self.config.core.metrics.apm_server.server_url, service_name="heartbeat_manager") else: self.apm_client = None
def test_comms_queue(redis_connection): if redis_connection: from assemblyline.remote.datatypes.queues.comms import CommsQueue def publish_messages(message_list): time.sleep(0.1) with CommsQueue('test-comms-queue') as cq_p: for message in message_list: cq_p.publish(message) msg_list = ["bob", 1, {"bob": 1}, [1, 2, 3], None, "Nice!", "stop"] t = Thread(target=publish_messages, args=(msg_list, )) t.start() with CommsQueue('test-comms-queue') as cq: x = 0 for msg in cq.listen(): if msg == "stop": break assert msg == msg_list[x] x += 1 t.join() assert not t.is_alive()
def test_alert_created(datastore, client): alert_queue = CommsQueue('alerts', private=True) created = random_model_obj(AlertMessage) created.msg_type = "AlertCreated" updated = random_model_obj(AlertMessage) updated.msg_type = "AlertUpdated" test_res_array = [] def alerter_created_callback(data): test_res_array.append(('created', created['msg'] == data)) def alerter_updated_callback(data): test_res_array.append(('updated', updated['msg'] == data)) def publish_thread(): time.sleep(1) alert_queue.publish(created.as_primitives()) alert_queue.publish(updated.as_primitives()) threading.Thread(target=publish_thread).start() client.socketio.listen_on_alerts_messages( alert_created_callback=alerter_created_callback, alert_updated_callback=alerter_updated_callback, timeout=2) assert len(test_res_array) == 2 for test, result in test_res_array: if not result: pytest.fail("{} failed.".format(test))
def __init__(self, sender, log, config=None, redis=None): self.sender = sender self.log = log self.config = config or forge.get_config() self.datastore = forge.get_datastore(self.config) self.redis = redis or get_client( host=self.config.core.redis.nonpersistent.host, port=self.config.core.redis.nonpersistent.port, private=False, ) self.redis_persist = get_client( host=self.config.core.redis.persistent.host, port=self.config.core.redis.persistent.port, private=False, ) self.status_queue = CommsQueue(STATUS_QUEUE, self.redis) self.dispatch_active_hash = Hash(DISPATCH_TASK_HASH, self.redis_persist) self.dispatcher_submission_queue = NamedQueue(SUBMISSION_QUEUE, self.redis) self.ingest_scanning = Hash('m-scanning-table', self.redis_persist) self.ingest_unique_queue = PriorityQueue('m-unique', self.redis_persist) self.ingest_queue = NamedQueue(INGEST_QUEUE_NAME, self.redis_persist) self.ingest_complete_queue = NamedQueue(COMPLETE_QUEUE_NAME, self.redis) self.alert_queue = NamedQueue(ALERT_QUEUE_NAME, self.redis_persist) constants = forge.get_constants(self.config) self.c_rng = constants.PRIORITY_RANGES['critical'] self.h_rng = constants.PRIORITY_RANGES['high'] self.m_rng = constants.PRIORITY_RANGES['medium'] self.l_rng = constants.PRIORITY_RANGES['low'] self.c_s_at = self.config.core.ingester.sampling_at['critical'] self.h_s_at = self.config.core.ingester.sampling_at['high'] self.m_s_at = self.config.core.ingester.sampling_at['medium'] self.l_s_at = self.config.core.ingester.sampling_at['low'] self.to_expire = {k: 0 for k in metrics.EXPIRY_METRICS} if self.config.core.expiry.batch_delete: self.delete_query = f"expiry_ts:[* TO {self.datastore.ds.now}-{self.config.core.expiry.delay}" \ f"{self.datastore.ds.hour}/DAY]" else: self.delete_query = f"expiry_ts:[* TO {self.datastore.ds.now}-{self.config.core.expiry.delay}" \ f"{self.datastore.ds.hour}]" self.scheduler = BackgroundScheduler(daemon=True) self.scheduler.add_job( self._reload_expiry_queues, 'interval', seconds=self.config.core.metrics.export_interval * 4) self.scheduler.start()
def test_status_messages(datastore, client): status_queue = CommsQueue('status', private=True) test_res_array = [] alerter_hb_msg = random_model_obj(AlerterMessage).as_primitives() dispatcher_hb_msg = random_model_obj(DispatcherMessage).as_primitives() expiry_hb_msg = random_model_obj(ExpiryMessage).as_primitives() ingest_hb_msg = random_model_obj(IngestMessage).as_primitives() service_hb_msg = random_model_obj(ServiceMessage).as_primitives() service_timing_msg = random_model_obj(ServiceTimingMessage).as_primitives() def alerter_callback(data): test_res_array.append(('alerter', alerter_hb_msg['msg'] == data)) def dispatcher_callback(data): test_res_array.append(('dispatcher', dispatcher_hb_msg['msg'] == data)) def expiry_callback(data): test_res_array.append(('expiry', expiry_hb_msg['msg'] == data)) def ingest_callback(data): test_res_array.append(('ingest', ingest_hb_msg['msg'] == data)) def service_callback(data): test_res_array.append(('service', service_hb_msg['msg'] == data)) def service_timing_callback(data): test_res_array.append( ('service_timing', service_timing_msg['msg'] == data)) def publish_thread(): time.sleep(1) status_queue.publish(alerter_hb_msg) status_queue.publish(dispatcher_hb_msg) status_queue.publish(expiry_hb_msg) status_queue.publish(ingest_hb_msg) status_queue.publish(service_hb_msg) status_queue.publish(service_timing_msg) threading.Thread(target=publish_thread).start() client.socketio.listen_on_status_messages( alerter_msg_callback=alerter_callback, dispatcher_msg_callback=dispatcher_callback, expiry_msg_callback=expiry_callback, ingest_msg_callback=ingest_callback, service_msg_callback=service_callback, service_timing_msg_callback=service_timing_callback, timeout=2) assert len(test_res_array) == 6 for test, result in test_res_array: if not result: pytest.fail("{} failed.".format(test))
def test_alert_namespace(datastore, sio): alert_queue = CommsQueue('alerts', private=True) test_id = get_random_id() created = random_model_obj(AlertMessage) created.msg_type = "AlertCreated" updated = random_model_obj(AlertMessage) updated.msg_type = "AlertUpdated" test_res_array = [] @sio.on('monitoring', namespace='/alerts') def on_monitoring(data): # Confirmation that we are waiting for alerts test_res_array.append(('on_monitoring', data == test_id)) @sio.on('AlertCreated', namespace='/alerts') def on_alert_created(data): test_res_array.append( ('on_alert_created', data == created.as_primitives()['msg'])) @sio.on('AlertUpdated', namespace='/alerts') def on_alert_updated(data): test_res_array.append( ('on_alert_updated', data == updated.as_primitives()['msg'])) try: sio.emit('alert', test_id, namespace='/alerts') sio.sleep(1) alert_queue.publish(created.as_primitives()) alert_queue.publish(updated.as_primitives()) start_time = time.time() while len(test_res_array) < 3 or time.time() - start_time < 5: sio.sleep(0.1) assert len(test_res_array) == 3 for test, result in test_res_array: if not result: pytest.fail(f"{test} failed.") finally: sio.disconnect()
def test_submission_ingested(datastore, client): submission_queue = CommsQueue('submissions', private=True) test_res_array = [] completed = random_model_obj(SubmissionMessage).as_primitives() completed['msg_type'] = "SubmissionCompleted" ingested = random_model_obj(SubmissionMessage).as_primitives() ingested['msg_type'] = "SubmissionIngested" received = random_model_obj(SubmissionMessage).as_primitives() received['msg_type'] = "SubmissionReceived" started = random_model_obj(SubmissionMessage).as_primitives() started['msg_type'] = "SubmissionStarted" def completed_callback(data): test_res_array.append(('completed', completed['msg'] == data)) def ingested_callback(data): test_res_array.append(('ingested', ingested['msg'] == data)) def received_callback(data): test_res_array.append(('received', received['msg'] == data)) def started_callback(data): test_res_array.append(('started', started['msg'] == data)) def publish_thread(): time.sleep(1) submission_queue.publish(completed) submission_queue.publish(ingested) submission_queue.publish(received) submission_queue.publish(started) threading.Thread(target=publish_thread).start() client.socketio.listen_on_submissions( completed_callback=completed_callback, ingested_callback=ingested_callback, received_callback=received_callback, started_callback=started_callback, timeout=2) assert len(test_res_array) == 4 for test, result in test_res_array: if not result: pytest.fail("{} failed.".format(test))
def try_run(self): # If our connection to the metrics database requires a custom ca cert, prepare it ca_certs = None if self.config.core.metrics.elasticsearch.host_certificates: with tempfile.NamedTemporaryFile(delete=False) as ca_certs_file: ca_certs = ca_certs_file.name ca_certs_file.write(self.config.core.metrics.elasticsearch.host_certificates.encode()) self.metrics_queue = CommsQueue(METRICS_QUEUE) self.es = elasticsearch.Elasticsearch(hosts=self.elastic_hosts, connection_class=elasticsearch.RequestsHttpConnection, ca_certs=ca_certs) self.scheduler.add_job(self._create_aggregated_metrics, 'interval', seconds=60) self.scheduler.start() while self.running: for msg in self.metrics_queue.listen(): # APM Transaction start if self.apm_client: self.apm_client.begin_transaction('metrics') m_name = msg.pop('name', None) m_type = msg.pop('type', None) msg.pop('host', None) msg.pop('instance', None) self.log.debug(f"Received {m_type.upper()} metrics message") if not m_name or not m_type: # APM Transaction end if self.apm_client: self.apm_client.end_transaction('process_message', 'invalid_message') continue with self.counters_lock: c_key = (m_name, m_type) if c_key not in self.counters or m_type in NON_AGGREGATED: self.counters[c_key] = Counter(msg) else: self.counters[c_key].update(Counter(msg)) # APM Transaction end if self.apm_client: self.apm_client.end_transaction('process_message', 'success')
def save_alert(datastore, counter, logger, alert, psid): if psid: msg_type = "AlertUpdated" perform_alert_update(datastore, logger, alert) counter.increment('updated') ret_val = 'update' else: msg_type = "AlertCreated" datastore.alert.save(alert['alert_id'], alert) logger.info(f"Alert {alert['alert_id']} has been created.") counter.increment('created') ret_val = 'create' msg = AlertMessage({ "msg": alert, "msg_type": msg_type, "sender": "alerter" }) CommsQueue('alerts').publish(msg.as_primitives()) return ret_val
def monitor_system_status(self): q = CommsQueue('status', private=True) try: for msg in q.listen(): if self.stop: break message = msg['msg'] msg_type = msg['msg_type'] self.socketio.emit(msg_type, message, namespace=self.namespace) LOGGER.info( f"SocketIO:{self.namespace} - Sending {msg_type} event to all connected users." ) except Exception: LOGGER.exception(f"SocketIO:{self.namespace}") finally: LOGGER.info( f"SocketIO:{self.namespace} - No more users connected to status monitoring, exiting thread..." ) with self.connections_lock: self.background_task = None
def _test_message_through_queue(queue_name, test_message, redis): t = Thread(target=publish_message, args=(queue_name, test_message, redis)) try: t.start() with CommsQueue(queue_name) as cq: for msg in cq.listen(): loader_path = msg.get('msg_loader', None) if loader_path is None: raise ValueError( "Message does not have a message loader class path.") msg_obj = load_module_by_path(loader_path) obj = msg_obj(msg) assert obj == test_message break finally: t.join() assert not t.is_alive()
def monitor_alerts(self, user_info): sid = user_info['sid'] q = CommsQueue('alerts', private=True) try: for msg in q.listen(): if sid not in self.connections: break alert = msg['msg'] msg_type = msg['msg_type'] if classification.is_accessible( user_info['classification'], alert.get('classification', classification.UNRESTRICTED)): self.socketio.emit(msg_type, alert, room=sid, namespace=self.namespace) LOGGER.info( f"SocketIO:{self.namespace} - {user_info['display']} - " f"Sending {msg_type} event for alert matching ID: {alert['alert_id']}" ) if AUDIT: AUDIT_LOG.info( f"{user_info['uname']} [{user_info['classification']}]" f" :: AlertMonitoringNamespace.get_alert(alert_id={alert['alert_id']})" ) except Exception: LOGGER.exception( f"SocketIO:{self.namespace} - {user_info['display']}") finally: LOGGER.info( f"SocketIO:{self.namespace} - {user_info['display']} - Connection to client was terminated" )
def publish_messages(message_list): time.sleep(0.1) with CommsQueue('test-comms-queue') as cq_p: for message in message_list: cq_p.publish(message)
def publish_message(queue_name, test_message, redis): time.sleep(0.1) with CommsQueue(queue_name, redis) as cq: cq.publish(test_message.as_primitives())
def __init__(self, datastore, logger, classification=None, redis=None, persistent_redis=None, metrics_name='ingester'): self.datastore = datastore self.log = logger # Cache the user groups self.cache_lock = threading.RLock( ) # TODO are middle man instances single threaded now? self._user_groups = {} self._user_groups_reset = time.time() // HOUR_IN_SECONDS self.cache = {} self.notification_queues = {} self.whitelisted = {} self.whitelisted_lock = threading.RLock() # Create a config cache that will refresh config values periodically self.config = forge.CachedObject(forge.get_config) # Module path parameters are fixed at start time. Changing these involves a restart self.is_low_priority = load_module_by_path( self.config.core.ingester.is_low_priority) self.get_whitelist_verdict = load_module_by_path( self.config.core.ingester.get_whitelist_verdict) self.whitelist = load_module_by_path( self.config.core.ingester.whitelist) # Constants are loaded based on a non-constant path, so has to be done at init rather than load constants = forge.get_constants(self.config) self.priority_value = constants.PRIORITIES self.priority_range = constants.PRIORITY_RANGES self.threshold_value = constants.PRIORITY_THRESHOLDS # Connect to the redis servers self.redis = redis or get_client( host=self.config.core.redis.nonpersistent.host, port=self.config.core.redis.nonpersistent.port, private=False, ) self.persistent_redis = persistent_redis or get_client( host=self.config.core.redis.persistent.host, port=self.config.core.redis.persistent.port, private=False, ) # Classification engine self.ce = classification or forge.get_classification() # Metrics gathering factory self.counter = MetricsFactory(metrics_type='ingester', schema=Metrics, redis=self.redis, config=self.config, name=metrics_name) # State. The submissions in progress are stored in Redis in order to # persist this state and recover in case we crash. self.scanning = Hash('m-scanning-table', self.persistent_redis) # Input. The dispatcher creates a record when any submission completes. self.complete_queue = NamedQueue(_completeq_name, self.redis) # Internal. Dropped entries are placed on this queue. # self.drop_queue = NamedQueue('m-drop', self.persistent_redis) # Input. An external process places submission requests on this queue. self.ingest_queue = NamedQueue(INGEST_QUEUE_NAME, self.persistent_redis) # Output. Duplicate our input traffic into this queue so it may be cloned by other systems self.traffic_queue = CommsQueue('submissions', self.redis) # Internal. Unique requests are placed in and processed from this queue. self.unique_queue = PriorityQueue('m-unique', self.persistent_redis) # Internal, delay queue for retrying self.retry_queue = PriorityQueue('m-retry', self.persistent_redis) # Internal, timeout watch queue self.timeout_queue = PriorityQueue('m-timeout', self.redis) # Internal, queue for processing duplicates # When a duplicate file is detected (same cache key => same file, and same # submission parameters) the file won't be ingested normally, but instead a reference # will be written to a duplicate queue. Whenever a file is finished, in the complete # method, not only is the original ingestion finalized, but all entries in the duplicate queue # are finalized as well. This has the effect that all concurrent ingestion of the same file # are 'merged' into a single submission to the system. self.duplicate_queue = MultiQueue(self.persistent_redis) # Output. submissions that should have alerts generated self.alert_queue = NamedQueue(ALERT_QUEUE_NAME, self.persistent_redis) # Utility object to help submit tasks to dispatching self.submit_client = SubmissionClient(datastore=self.datastore, redis=self.redis)
def get_metrics_sink(redis=None): from assemblyline.remote.datatypes.queues.comms import CommsQueue return CommsQueue('assemblyline_metrics', host=redis)
#!/usr/bin/env python import sys from assemblyline.remote.datatypes.queues.comms import CommsQueue from pprint import pprint if __name__ == "__main__": queue_name = None if len(sys.argv) > 1: queue_name = sys.argv[1] if queue_name is None: print( "\nERROR: You must specify a queue name.\n\npubsub_reader.py [queue_name]" ) exit(1) print(f"Listening for messages on '{queue_name}' queue.") q = CommsQueue(queue_name) try: while True: for msg in q.listen(): pprint(msg) except KeyboardInterrupt: print('Exiting') finally: q.close()
def test_status_namspace(datastore, sio): status_queue = CommsQueue('status', private=True) monitoring = get_random_id() alerter_hb_msg = random_model_obj(AlerterMessage).as_primitives() dispatcher_hb_msg = random_model_obj(DispatcherMessage).as_primitives() expiry_hb_msg = random_model_obj(ExpiryMessage).as_primitives() ingest_hb_msg = random_model_obj(IngestMessage).as_primitives() service_hb_msg = random_model_obj(ServiceMessage).as_primitives() test_res_array = [] @sio.on('monitoring', namespace='/status') def on_monitoring(data): # Confirmation that we are waiting for status messages test_res_array.append(('on_monitoring', data == monitoring)) @sio.on('AlerterHeartbeat', namespace='/status') def on_alerter_heartbeat(data): test_res_array.append( ('on_alerter_heartbeat', data == alerter_hb_msg['msg'])) @sio.on('DispatcherHeartbeat', namespace='/status') def on_dispatcher_heartbeat(data): test_res_array.append( ('on_dispatcher_heartbeat', data == dispatcher_hb_msg['msg'])) @sio.on('ExpiryHeartbeat', namespace='/status') def on_expiry_heartbeat(data): test_res_array.append( ('on_expiry_heartbeat', data == expiry_hb_msg['msg'])) @sio.on('IngestHeartbeat', namespace='/status') def on_ingest_heartbeat(data): test_res_array.append( ('on_ingest_heartbeat', data == ingest_hb_msg['msg'])) @sio.on('ServiceHeartbeat', namespace='/status') def on_service_heartbeat(data): test_res_array.append( ('on_service_heartbeat', data == service_hb_msg['msg'])) try: sio.emit('monitor', monitoring, namespace='/status') sio.sleep(1) status_queue.publish(alerter_hb_msg) status_queue.publish(dispatcher_hb_msg) status_queue.publish(expiry_hb_msg) status_queue.publish(ingest_hb_msg) status_queue.publish(service_hb_msg) start_time = time.time() while len(test_res_array) < 6 and time.time() - start_time < 5: sio.sleep(0.1) assert len(test_res_array) == 6 for test, result in test_res_array: if not result: pytest.fail(f"{test} failed.") finally: sio.disconnect()
def get_submission_traffic_channel(): return CommsQueue('submissions', host=config.core.redis.nonpersistent.host, port=config.core.redis.nonpersistent.port)