class ExpiryManager(ServerBase):
def __init__(self, force_ilm=False):
self.config = forge.get_config()
if force_ilm:
self.config.datastore.ilm.enabled = True
super().__init__('assemblyline.expiry',
shutdown_timeout=self.config.core.expiry.sleep_time +
5)
self.datastore = forge.get_datastore(config=self.config,
archive_access=True)
self.hot_datastore = forge.get_datastore(config=self.config,
archive_access=False)
self.filestore = forge.get_filestore(config=self.config)
self.cachestore = FileStore(*self.config.filestore.cache)
self.expirable_collections = []
self.archiveable_collections = []
self.counter = MetricsFactory('expiry', Metrics)
self.counter_archive = MetricsFactory('archive', Metrics)
if self.config.datastore.ilm.enabled:
self.fs_hashmap = {
'file': self.archive_filestore_delete,
'cached_file': self.archive_cachestore_delete
}
else:
self.fs_hashmap = {
'file': self.filestore_delete,
'cached_file': self.cachestore_delete
}
for name, definition in self.datastore.ds.get_models().items():
if hasattr(definition, 'archive_ts'):
self.archiveable_collections.append(
getattr(self.datastore, name))
if hasattr(definition, 'expiry_ts'):
self.expirable_collections.append(getattr(
self.datastore, name))
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="expiry")
else:
self.apm_client = None
def stop(self):
if self.counter:
self.counter.stop()
if self.apm_client:
elasticapm.uninstrument()
super().stop()
def filestore_delete(self, sha256, _):
self.filestore.delete(sha256)
def archive_filestore_delete(self, sha256, expiry_time):
# If we are working with an archive, their may be a hot entry that expires later.
doc = self.hot_datastore.file.get_if_exists(sha256, as_obj=False)
if doc and doc['expiry_ts'] > expiry_time:
return
self.filestore.delete(sha256)
def cachestore_delete(self, sha256, _):
self.filestore.delete(sha256)
def archive_cachestore_delete(self, sha256, expiry_time):
doc = self.hot_datastore.cached_file.get_if_exists(sha256,
as_obj=False)
if doc and doc['expiry_ts'] > expiry_time:
return
self.cachestore.delete(sha256)
def run_expiry_once(self):
now = now_as_iso()
reached_max = False
# Expire data
for collection in self.expirable_collections:
# Call heartbeat pre-dated by 5 minutes. If a collection takes more than
# 5 minutes to expire, this container could be seen as unhealthy. The down
# side is if it is stuck on something it will be more than 5 minutes before
# the container is restarted.
self.heartbeat(int(time.time() + 5 * 60))
# Start of expiry transaction
if self.apm_client:
self.apm_client.begin_transaction("Delete expired documents")
if self.config.core.expiry.batch_delete:
computed_date = epoch_to_iso(
dm(f"{now}||-{self.config.core.expiry.delay}h/d").
float_timestamp)
else:
computed_date = epoch_to_iso(
dm(f"{now}||-{self.config.core.expiry.delay}h").
float_timestamp)
delete_query = f"expiry_ts:[* TO {computed_date}]"
if self.config.core.expiry.delete_storage and collection.name in self.fs_hashmap:
file_delete = True
sort = ["expiry_ts asc", "id asc"]
else:
file_delete = False
sort = None
number_to_delete = collection.search(
delete_query,
rows=0,
as_obj=False,
use_archive=True,
sort=sort,
track_total_hits=EXPIRY_SIZE)['total']
if self.apm_client:
elasticapm.label(query=delete_query)
elasticapm.label(number_to_delete=number_to_delete)
self.log.info(f"Processing collection: {collection.name}")
if number_to_delete != 0:
if file_delete:
with elasticapm.capture_span(
name='FILESTORE [ThreadPoolExecutor] :: delete()',
labels={
"num_files": number_to_delete,
"query": delete_query
}):
# Delete associated files
with concurrent.futures.ThreadPoolExecutor(
self.config.core.expiry.workers,
thread_name_prefix="file_delete") as executor:
for item in collection.search(
delete_query,
fl='id',
rows=number_to_delete,
sort=sort,
use_archive=True,
as_obj=False)['items']:
executor.submit(
self.fs_hashmap[collection.name],
item['id'], computed_date)
self.log.info(
f' Deleted associated files from the '
f'{"cachestore" if "cache" in collection.name else "filestore"}...'
)
# Proceed with deletion
collection.delete_by_query(
delete_query,
workers=self.config.core.expiry.workers,
sort=sort,
max_docs=number_to_delete)
else:
# Proceed with deletion
collection.delete_by_query(
delete_query, workers=self.config.core.expiry.workers)
if number_to_delete == EXPIRY_SIZE:
reached_max = True
self.counter.increment(f'{collection.name}',
increment_by=number_to_delete)
self.log.info(
f" Deleted {number_to_delete} items from the datastore..."
)
else:
self.log.debug(" Nothing to delete in this collection.")
# End of expiry transaction
if self.apm_client:
self.apm_client.end_transaction(collection.name, 'deleted')
return reached_max
def run_archive_once(self):
reached_max = False
if not self.config.datastore.ilm.enabled:
return reached_max
now = now_as_iso()
# Archive data
for collection in self.archiveable_collections:
# Call heartbeat pre-dated by 5 minutes. If a collection takes more than
# 5 minutes to expire, this container could be seen as unhealthy. The down
# side is if it is stuck on something it will be more than 5 minutes before
# the container is restarted.
self.heartbeat(int(time.time() + 5 * 60))
# Start of expiry transaction
if self.apm_client:
self.apm_client.begin_transaction("Archive older documents")
archive_query = f"archive_ts:[* TO {now}]"
sort = ["archive_ts asc", "id asc"]
number_to_archive = collection.search(
archive_query,
rows=0,
as_obj=False,
use_archive=False,
sort=sort,
track_total_hits=ARCHIVE_SIZE)['total']
if number_to_archive == ARCHIVE_SIZE:
reached_max = True
if self.apm_client:
elasticapm.label(query=archive_query)
elasticapm.label(number_to_archive=number_to_archive)
self.log.info(f"Processing collection: {collection.name}")
if number_to_archive != 0:
# Proceed with archiving
if collection.archive(archive_query,
max_docs=number_to_archive,
sort=sort):
self.counter_archive.increment(
f'{collection.name}', increment_by=number_to_archive)
self.log.info(
f" Archived {number_to_archive} documents...")
else:
self.log.warning(
f" Failed to properly archive {number_to_archive} documents..."
)
else:
self.log.debug(" Nothing to archive in this collection.")
# End of expiry transaction
if self.apm_client:
self.apm_client.end_transaction(collection.name, 'archived')
return reached_max
def try_run(self):
while self.running:
expiry_maxed_out = False
archive_maxed_out = False
try:
expiry_maxed_out = self.run_expiry_once()
except Exception as e:
self.log.exception(str(e))
try:
archive_maxed_out = self.run_archive_once()
except Exception as e:
self.log.exception(str(e))
if not expiry_maxed_out and not archive_maxed_out:
self.sleep_with_heartbeat(self.config.core.expiry.sleep_time)
class ExpiryManager(ServerBase):
def __init__(self):
self.config = forge.get_config()
super().__init__('assemblyline.expiry',
shutdown_timeout=self.config.core.expiry.sleep_time +
5)
self.datastore = forge.get_datastore(config=self.config,
archive_access=True)
self.filestore = forge.get_filestore(config=self.config)
self.cachestore = FileStore(*self.config.filestore.cache)
self.expirable_collections = []
self.archiveable_collections = []
self.counter = MetricsFactory('expiry', Metrics)
self.counter_archive = MetricsFactory('archive', Metrics)
self.fs_hashmap = {
'file': self.filestore.delete,
'cached_file': self.cachestore.delete
}
for name, definition in self.datastore.ds.get_models().items():
if hasattr(definition, 'archive_ts'):
self.archiveable_collections.append(
getattr(self.datastore, name))
if hasattr(definition, 'expiry_ts'):
self.expirable_collections.append(getattr(
self.datastore, name))
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="expiry")
else:
self.apm_client = None
def close(self):
if self.counter:
self.counter.stop()
if self.apm_client:
elasticapm.uninstrument()
def run_expiry_once(self):
now = now_as_iso()
delay = self.config.core.expiry.delay
hour = self.datastore.ds.hour
day = self.datastore.ds.day
# Expire data
for collection in self.expirable_collections:
# Call heartbeat pre-dated by 5 minutes. If a collection takes more than
# 5 minutes to expire, this container could be seen as unhealthy. The down
# side is if it is stuck on something it will be more than 5 minutes before
# the container is restarted.
self.heartbeat(int(time.time() + 5 * 60))
# Start of expiry transaction
if self.apm_client:
self.apm_client.begin_transaction("Delete expired documents")
if self.config.core.expiry.batch_delete:
delete_query = f"expiry_ts:[* TO {now}||-{delay}{hour}/{day}]"
else:
delete_query = f"expiry_ts:[* TO {now}||-{delay}{hour}]"
number_to_delete = collection.search(delete_query,
rows=0,
as_obj=False)['total']
if self.apm_client:
elasticapm.tag(query=delete_query)
elasticapm.tag(number_to_delete=number_to_delete)
self.log.info(f"Processing collection: {collection.name}")
if number_to_delete != 0:
if self.config.core.expiry.delete_storage and collection.name in self.fs_hashmap:
with elasticapm.capture_span(
name=f'FILESTORE [ThreadPoolExecutor] :: delete()',
labels={
"num_files": number_to_delete,
"query": delete_query
}):
# Delete associated files
with concurrent.futures.ThreadPoolExecutor(
self.config.core.expiry.workers) as executor:
res = {
item['id']: executor.submit(
self.fs_hashmap[collection.name],
item['id'])
for item in collection.stream_search(
delete_query, fl='id', as_obj=False)
}
for v in res.values():
v.result()
self.log.info(
f' Deleted associated files from the '
f'{"cachestore" if "cache" in collection.name else "filestore"}...'
)
# Proceed with deletion
collection.delete_matching(
delete_query, workers=self.config.core.expiry.workers)
self.counter.increment(f'{collection.name}',
increment_by=number_to_delete)
self.log.info(
f" Deleted {number_to_delete} items from the datastore..."
)
else:
self.log.debug(" Nothing to delete in this collection.")
# End of expiry transaction
if self.apm_client:
self.apm_client.end_transaction(collection.name, 'deleted')
def run_archive_once(self):
if not self.config.datastore.ilm.enabled:
return
now = now_as_iso()
# Archive data
for collection in self.archiveable_collections:
# Start of expiry transaction
if self.apm_client:
self.apm_client.begin_transaction("Archive older documents")
archive_query = f"archive_ts:[* TO {now}]"
number_to_archive = collection.search(archive_query,
rows=0,
as_obj=False,
use_archive=False)['total']
if self.apm_client:
elasticapm.tag(query=archive_query)
elasticapm.tag(number_to_archive=number_to_archive)
self.log.info(f"Processing collection: {collection.name}")
if number_to_archive != 0:
# Proceed with archiving
collection.archive(archive_query)
self.counter_archive.increment(f'{collection.name}',
increment_by=number_to_archive)
self.log.info(
f" Archived {number_to_archive} items to the time sliced storage..."
)
else:
self.log.debug(" Nothing to archive in this collection.")
# End of expiry transaction
if self.apm_client:
self.apm_client.end_transaction(collection.name, 'archived')
def try_run(self):
while self.running:
try:
self.run_expiry_once()
except Exception as e:
self.log.exception(str(e))
try:
self.run_archive_once()
except Exception as e:
self.log.exception(str(e))
self.sleep_with_heartbeat(self.config.core.expiry.sleep_time)
class WatcherServer(CoreBase):
def __init__(self, redis=None, redis_persist=None):
super().__init__('assemblyline.watcher',
redis=redis,
redis_persist=redis_persist)
# Watcher structures
self.hash = ExpiringHash(name=WATCHER_HASH,
ttl=MAX_TIMEOUT,
host=self.redis_persist)
self.queue = UniquePriorityQueue(WATCHER_QUEUE, self.redis_persist)
# Task management structures
self.running_tasks = ExpiringHash(
DISPATCH_RUNNING_TASK_HASH,
host=self.redis) # TODO, move to persistant?
self.scaler_timeout_queue = NamedQueue(SCALER_TIMEOUT_QUEUE,
host=self.redis_persist)
# Metrics tracking
self.counter = MetricsFactory(metrics_type='watcher',
schema=Metrics,
name='watcher',
redis=self.redis,
config=self.config)
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="watcher")
else:
self.apm_client = None
def try_run(self):
counter = self.counter
apm_client = self.apm_client
while self.running:
self.heartbeat()
# Download all messages from the queue that have expired
seconds, _ = retry_call(self.redis.time)
messages = self.queue.dequeue_range(0, seconds)
cpu_mark = time.process_time()
time_mark = time.time()
# Try to pass on all the messages to their intended recipient, try not to let
# the failure of one message from preventing the others from going through
for key in messages:
# Start of transaction
if apm_client:
apm_client.begin_transaction('process_messages')
message = self.hash.pop(key)
if message:
try:
if message['action'] == WatcherAction.TimeoutTask:
self.cancel_service_task(message['task_key'],
message['worker'])
else:
queue = NamedQueue(message['queue'], self.redis)
queue.push(message['message'])
self.counter.increment('expired')
# End of transaction (success)
if apm_client:
apm_client.end_transaction('watch_message',
'success')
except Exception as error:
# End of transaction (exception)
if apm_client:
apm_client.end_transaction('watch_message',
'error')
self.log.exception(error)
else:
# End of transaction (duplicate)
if apm_client:
apm_client.end_transaction('watch_message',
'duplicate')
self.log.warning(
f'Handled watch twice: {key} {len(key)} {type(key)}')
counter.increment_execution_time('cpu_seconds',
time.process_time() - cpu_mark)
counter.increment_execution_time('busy_seconds',
time.time() - time_mark)
if not messages:
time.sleep(0.1)
def cancel_service_task(self, task_key, worker):
# We believe a service task has timed out, try and read it from running tasks
# If we can't find the task in running tasks, it finished JUST before timing out, let it go
task = self.running_tasks.pop(task_key)
if not task:
return
# We can confirm that the task is ours now, even if the worker finished, the result will be ignored
task = Task(task)
self.log.info(
f"[{task.sid}] Service {task.service_name} timed out on {task.fileinfo.sha256}."
)
# Mark the previous attempt as invalid and redispatch it
dispatch_table = DispatchHash(task.sid, self.redis)
dispatch_table.fail_recoverable(task.fileinfo.sha256,
task.service_name)
dispatch_table.dispatch(task.fileinfo.sha256, task.service_name)
get_service_queue(task.service_name,
self.redis).push(task.priority, task.as_primitives())
# We push the task of killing the container off on the scaler, which already has root access
# the scaler can also double check that the service name and container id match, to be sure
# we aren't accidentally killing the wrong container
self.scaler_timeout_queue.push({
'service': task.service_name,
'container': worker
})
# Report to the metrics system that a recoverable error has occurred for that service
export_metrics_once(task.service_name,
ServiceMetrics,
dict(fail_recoverable=1),
host=worker,
counter_type='service')
def __init__(self,
datastore=None,
logger=None,
classification=None,
redis=None,
persistent_redis=None,
metrics_name='ingester',
config=None):
super().__init__('assemblyline.ingester',
logger,
redis=redis,
redis_persist=persistent_redis,
datastore=datastore,
config=config)
# Cache the user groups
self.cache_lock = threading.RLock()
self._user_groups = {}
self._user_groups_reset = time.time() // HOUR_IN_SECONDS
self.cache = {}
self.notification_queues = {}
self.whitelisted = {}
self.whitelisted_lock = threading.RLock()
# 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: dict[str, int] = constants.PRIORITIES
self.priority_range: dict[str, Tuple[int,
int]] = constants.PRIORITY_RANGES
self.threshold_value: dict[str, int] = constants.PRIORITY_THRESHOLDS
# 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.redis_persist)
# Input. The dispatcher creates a record when any submission completes.
self.complete_queue = NamedQueue(COMPLETE_QUEUE_NAME, self.redis)
# Input. An external process places submission requests on this queue.
self.ingest_queue = NamedQueue(INGEST_QUEUE_NAME, self.redis_persist)
# 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.redis_persist)
# Internal, delay queue for retrying
self.retry_queue = PriorityQueue('m-retry', self.redis_persist)
# Internal, timeout watch queue
self.timeout_queue: PriorityQueue[str] = 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.redis_persist)
# Output. submissions that should have alerts generated
self.alert_queue = NamedQueue(ALERT_QUEUE_NAME, self.redis_persist)
# Utility object to help submit tasks to dispatching
self.submit_client = SubmissionClient(datastore=self.datastore,
redis=self.redis)
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="ingester")
else:
self.apm_client = None
class Ingester(ThreadedCoreBase):
def __init__(self,
datastore=None,
logger=None,
classification=None,
redis=None,
persistent_redis=None,
metrics_name='ingester',
config=None):
super().__init__('assemblyline.ingester',
logger,
redis=redis,
redis_persist=persistent_redis,
datastore=datastore,
config=config)
# Cache the user groups
self.cache_lock = threading.RLock()
self._user_groups = {}
self._user_groups_reset = time.time() // HOUR_IN_SECONDS
self.cache = {}
self.notification_queues = {}
self.whitelisted = {}
self.whitelisted_lock = threading.RLock()
# 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: dict[str, int] = constants.PRIORITIES
self.priority_range: dict[str, Tuple[int,
int]] = constants.PRIORITY_RANGES
self.threshold_value: dict[str, int] = constants.PRIORITY_THRESHOLDS
# 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.redis_persist)
# Input. The dispatcher creates a record when any submission completes.
self.complete_queue = NamedQueue(COMPLETE_QUEUE_NAME, self.redis)
# Input. An external process places submission requests on this queue.
self.ingest_queue = NamedQueue(INGEST_QUEUE_NAME, self.redis_persist)
# 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.redis_persist)
# Internal, delay queue for retrying
self.retry_queue = PriorityQueue('m-retry', self.redis_persist)
# Internal, timeout watch queue
self.timeout_queue: PriorityQueue[str] = 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.redis_persist)
# Output. submissions that should have alerts generated
self.alert_queue = NamedQueue(ALERT_QUEUE_NAME, self.redis_persist)
# Utility object to help submit tasks to dispatching
self.submit_client = SubmissionClient(datastore=self.datastore,
redis=self.redis)
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="ingester")
else:
self.apm_client = None
def try_run(self):
threads_to_maintain = {
'Retries': self.handle_retries,
'Timeouts': self.handle_timeouts
}
threads_to_maintain.update({
f'Complete_{n}': self.handle_complete
for n in range(COMPLETE_THREADS)
})
threads_to_maintain.update(
{f'Ingest_{n}': self.handle_ingest
for n in range(INGEST_THREADS)})
threads_to_maintain.update(
{f'Submit_{n}': self.handle_submit
for n in range(SUBMIT_THREADS)})
self.maintain_threads(threads_to_maintain)
def handle_ingest(self):
cpu_mark = time.process_time()
time_mark = time.time()
# Move from ingest to unique and waiting queues.
# While there are entries in the ingest queue we consume chunk_size
# entries at a time and move unique entries to uniqueq / queued and
# duplicates to their own queues / waiting.
while self.running:
self.counter.increment_execution_time(
'cpu_seconds',
time.process_time() - cpu_mark)
self.counter.increment_execution_time('busy_seconds',
time.time() - time_mark)
message = self.ingest_queue.pop(timeout=1)
cpu_mark = time.process_time()
time_mark = time.time()
if not message:
continue
# Start of ingest message
if self.apm_client:
self.apm_client.begin_transaction('ingest_msg')
try:
if 'submission' in message:
# A retried task
task = IngestTask(message)
else:
# A new submission
sub = MessageSubmission(message)
task = IngestTask(dict(
submission=sub,
ingest_id=sub.sid,
))
task.submission.sid = None # Reset to new random uuid
# Write all input to the traffic queue
self.traffic_queue.publish(
SubmissionMessage({
'msg': sub,
'msg_type': 'SubmissionIngested',
'sender': 'ingester',
}).as_primitives())
except (ValueError, TypeError) as error:
self.counter.increment('error')
self.log.exception(
f"Dropped ingest submission {message} because {str(error)}"
)
# End of ingest message (value_error)
if self.apm_client:
self.apm_client.end_transaction('ingest_input',
'value_error')
continue
self.ingest(task)
# End of ingest message (success)
if self.apm_client:
self.apm_client.end_transaction('ingest_input', 'success')
def handle_submit(self):
time_mark, cpu_mark = time.time(), time.process_time()
while self.running:
# noinspection PyBroadException
try:
self.counter.increment_execution_time(
'cpu_seconds',
time.process_time() - cpu_mark)
self.counter.increment_execution_time('busy_seconds',
time.time() - time_mark)
# Check if there is room for more submissions
length = self.scanning.length()
if length >= self.config.core.ingester.max_inflight:
self.sleep(0.1)
time_mark, cpu_mark = time.time(), time.process_time()
continue
raw = self.unique_queue.blocking_pop(timeout=3)
time_mark, cpu_mark = time.time(), time.process_time()
if not raw:
continue
# Start of ingest message
if self.apm_client:
self.apm_client.begin_transaction('ingest_msg')
task = IngestTask(raw)
# Check if we need to drop a file for capacity reasons, but only if the
# number of files in flight is alreay over 80%
if length >= self.config.core.ingester.max_inflight * 0.8 and self.drop(
task):
# End of ingest message (dropped)
if self.apm_client:
self.apm_client.end_transaction(
'ingest_submit', 'dropped')
continue
if self.is_whitelisted(task):
# End of ingest message (whitelisted)
if self.apm_client:
self.apm_client.end_transaction(
'ingest_submit', 'whitelisted')
continue
# Check if this file has been previously processed.
pprevious, previous, score, scan_key = None, None, None, None
if not task.submission.params.ignore_cache:
pprevious, previous, score, scan_key = self.check(task)
else:
scan_key = self.stamp_filescore_key(task)
# If it HAS been previously processed, we are dealing with a resubmission
# finalize will decide what to do, and put the task back in the queue
# rewritten properly if we are going to run it again
if previous:
if not task.submission.params.services.resubmit and not pprevious:
self.log.warning(
f"No psid for what looks like a resubmission of "
f"{task.submission.files[0].sha256}: {scan_key}")
self.finalize(pprevious, previous, score, task)
# End of ingest message (finalized)
if self.apm_client:
self.apm_client.end_transaction(
'ingest_submit', 'finalized')
continue
# We have decided this file is worth processing
# Add the task to the scanning table, this is atomic across all submit
# workers, so if it fails, someone beat us to the punch, record the file
# as a duplicate then.
if not self.scanning.add(scan_key, task.as_primitives()):
self.log.debug('Duplicate %s',
task.submission.files[0].sha256)
self.counter.increment('duplicates')
self.duplicate_queue.push(_dup_prefix + scan_key,
task.as_primitives())
# End of ingest message (duplicate)
if self.apm_client:
self.apm_client.end_transaction(
'ingest_submit', 'duplicate')
continue
# We have managed to add the task to the scan table, so now we go
# ahead with the submission process
try:
self.submit(task)
# End of ingest message (submitted)
if self.apm_client:
self.apm_client.end_transaction(
'ingest_submit', 'submitted')
continue
except Exception as _ex:
# For some reason (contained in `ex`) we have failed the submission
# The rest of this function is error handling/recovery
ex = _ex
# traceback = _ex.__traceback__
self.counter.increment('error')
should_retry = True
if isinstance(ex, CorruptedFileStoreException):
self.log.error(
"Submission for file '%s' failed due to corrupted "
"filestore: %s" % (task.sha256, str(ex)))
should_retry = False
elif isinstance(ex, DataStoreException):
trace = exceptions.get_stacktrace_info(ex)
self.log.error("Submission for file '%s' failed due to "
"data store error:\n%s" %
(task.sha256, trace))
elif not isinstance(ex, FileStoreException):
trace = exceptions.get_stacktrace_info(ex)
self.log.error("Submission for file '%s' failed: %s" %
(task.sha256, trace))
task = IngestTask(self.scanning.pop(scan_key))
if not task:
self.log.error('No scanning entry for for %s', task.sha256)
# End of ingest message (no_scan_entry)
if self.apm_client:
self.apm_client.end_transaction(
'ingest_submit', 'no_scan_entry')
continue
if not should_retry:
# End of ingest message (cannot_retry)
if self.apm_client:
self.apm_client.end_transaction(
'ingest_submit', 'cannot_retry')
continue
self.retry(task, scan_key, ex)
# End of ingest message (retry)
if self.apm_client:
self.apm_client.end_transaction('ingest_submit', 'retried')
except Exception:
self.log.exception("Unexpected error")
# End of ingest message (exception)
if self.apm_client:
self.apm_client.end_transaction('ingest_submit',
'exception')
def handle_complete(self):
while self.running:
result = self.complete_queue.pop(timeout=3)
if not result:
continue
cpu_mark = time.process_time()
time_mark = time.time()
# Start of ingest message
if self.apm_client:
self.apm_client.begin_transaction('ingest_msg')
sub = DatabaseSubmission(result)
self.completed(sub)
# End of ingest message (success)
if self.apm_client:
elasticapm.label(sid=sub.sid)
self.apm_client.end_transaction('ingest_complete', 'success')
self.counter.increment_execution_time(
'cpu_seconds',
time.process_time() - cpu_mark)
self.counter.increment_execution_time('busy_seconds',
time.time() - time_mark)
def handle_retries(self):
tasks = []
while self.sleep(0 if tasks else 3):
cpu_mark = time.process_time()
time_mark = time.time()
# Start of ingest message
if self.apm_client:
self.apm_client.begin_transaction('ingest_retries')
tasks = self.retry_queue.dequeue_range(upper_limit=isotime.now(),
num=100)
for task in tasks:
self.ingest_queue.push(task)
# End of ingest message (success)
if self.apm_client:
elasticapm.label(retries=len(tasks))
self.apm_client.end_transaction('ingest_retries', 'success')
self.counter.increment_execution_time(
'cpu_seconds',
time.process_time() - cpu_mark)
self.counter.increment_execution_time('busy_seconds',
time.time() - time_mark)
def handle_timeouts(self):
timeouts = []
while self.sleep(0 if timeouts else 3):
cpu_mark = time.process_time()
time_mark = time.time()
# Start of ingest message
if self.apm_client:
self.apm_client.begin_transaction('ingest_timeouts')
timeouts = self.timeout_queue.dequeue_range(
upper_limit=isotime.now(), num=100)
for scan_key in timeouts:
# noinspection PyBroadException
try:
actual_timeout = False
# Remove the entry from the hash of submissions in progress.
entry = self.scanning.pop(scan_key)
if entry:
actual_timeout = True
self.log.error("Submission timed out for %s: %s",
scan_key, str(entry))
dup = self.duplicate_queue.pop(_dup_prefix + scan_key,
blocking=False)
if dup:
actual_timeout = True
while dup:
self.log.error("Submission timed out for %s: %s",
scan_key, str(dup))
dup = self.duplicate_queue.pop(_dup_prefix + scan_key,
blocking=False)
if actual_timeout:
self.counter.increment('timed_out')
except Exception:
self.log.exception("Problem timing out %s:", scan_key)
# End of ingest message (success)
if self.apm_client:
elasticapm.label(timeouts=len(timeouts))
self.apm_client.end_transaction('ingest_timeouts', 'success')
self.counter.increment_execution_time(
'cpu_seconds',
time.process_time() - cpu_mark)
self.counter.increment_execution_time('busy_seconds',
time.time() - time_mark)
def get_groups_from_user(self, username: str) -> List[str]:
# Reset the group cache at the top of each hour
if time.time() // HOUR_IN_SECONDS > self._user_groups_reset:
self._user_groups = {}
self._user_groups_reset = time.time() // HOUR_IN_SECONDS
# Get the groups for this user if not known
if username not in self._user_groups:
user_data = self.datastore.user.get(username)
if user_data:
self._user_groups[username] = user_data.groups
else:
self._user_groups[username] = []
return self._user_groups[username]
def ingest(self, task: IngestTask):
self.log.info(
f"[{task.ingest_id} :: {task.sha256}] Task received for processing"
)
# Load a snapshot of ingest parameters as of right now.
max_file_size = self.config.submission.max_file_size
param = task.params
self.counter.increment('bytes_ingested', increment_by=task.file_size)
self.counter.increment('submissions_ingested')
if any(len(file.sha256) != 64 for file in task.submission.files):
self.log.error(
f"[{task.ingest_id} :: {task.sha256}] Invalid sha256, skipped")
self.send_notification(task,
failure="Invalid sha256",
logfunc=self.log.warning)
return
# Clean up metadata strings, since we may delete some, iterate on a copy of the keys
for key in list(task.submission.metadata.keys()):
value = task.submission.metadata[key]
meta_size = len(value)
if meta_size > self.config.submission.max_metadata_length:
self.log.info(
f'[{task.ingest_id} :: {task.sha256}] '
f'Removing {key} from metadata because value is too big')
task.submission.metadata.pop(key)
if task.file_size > max_file_size and not task.params.ignore_size and not task.params.never_drop:
task.failure = f"File too large ({task.file_size} > {max_file_size})"
self._notify_drop(task)
self.counter.increment('skipped')
self.log.error(
f"[{task.ingest_id} :: {task.sha256}] {task.failure}")
return
# Set the groups from the user, if they aren't already set
if not task.params.groups:
task.params.groups = self.get_groups_from_user(
task.params.submitter)
# Check if this file is already being processed
self.stamp_filescore_key(task)
pprevious, previous, score = None, None, None
if not param.ignore_cache:
pprevious, previous, score, _ = self.check(task, count_miss=False)
# Assign priority.
low_priority = self.is_low_priority(task)
priority = param.priority
if priority < 0:
priority = self.priority_value['medium']
if score is not None:
priority = self.priority_value['low']
for level, threshold in self.threshold_value.items():
if score >= threshold:
priority = self.priority_value[level]
break
elif low_priority:
priority = self.priority_value['low']
# Reduce the priority by an order of magnitude for very old files.
current_time = now()
if priority and self.expired(
current_time - task.submission.time.timestamp(), 0):
priority = (priority / 10) or 1
param.priority = priority
# Do this after priority has been assigned.
# (So we don't end up dropping the resubmission).
if previous:
self.counter.increment('duplicates')
self.finalize(pprevious, previous, score, task)
# On cache hits of any kind we want to send out a completed message
self.traffic_queue.publish(
SubmissionMessage({
'msg': task.submission,
'msg_type': 'SubmissionCompleted',
'sender': 'ingester',
}).as_primitives())
return
if self.drop(task):
self.log.info(f"[{task.ingest_id} :: {task.sha256}] Dropped")
return
if self.is_whitelisted(task):
self.log.info(f"[{task.ingest_id} :: {task.sha256}] Whitelisted")
return
self.unique_queue.push(priority, task.as_primitives())
def check(
self,
task: IngestTask,
count_miss=True
) -> Tuple[Optional[str], Optional[str], Optional[float], str]:
key = self.stamp_filescore_key(task)
with self.cache_lock:
result = self.cache.get(key, None)
if result:
self.counter.increment('cache_hit_local')
self.log.info(
f'[{task.ingest_id} :: {task.sha256}] Local cache hit')
else:
result = self.datastore.filescore.get_if_exists(key)
if result:
self.counter.increment('cache_hit')
self.log.info(
f'[{task.ingest_id} :: {task.sha256}] Remote cache hit')
else:
if count_miss:
self.counter.increment('cache_miss')
return None, None, None, key
with self.cache_lock:
self.cache[key] = result
current_time = now()
age = current_time - result.time
errors = result.errors
if self.expired(age, errors):
self.log.info(
f"[{task.ingest_id} :: {task.sha256}] Cache hit dropped, cache has expired"
)
self.counter.increment('cache_expired')
self.cache.pop(key, None)
self.datastore.filescore.delete(key)
return None, None, None, key
elif self.stale(age, errors):
self.log.info(
f"[{task.ingest_id} :: {task.sha256}] Cache hit dropped, cache is stale"
)
self.counter.increment('cache_stale')
return None, None, result.score, key
return result.psid, result.sid, result.score, key
def stop(self):
super().stop()
if self.apm_client:
elasticapm.uninstrument()
self.submit_client.stop()
def stale(self, delta: float, errors: int):
if errors:
return delta >= self.config.core.ingester.incomplete_stale_after_seconds
else:
return delta >= self.config.core.ingester.stale_after_seconds
@staticmethod
def stamp_filescore_key(task: IngestTask, sha256: str = None) -> str:
if not sha256:
sha256 = task.submission.files[0].sha256
key = task.submission.scan_key
if not key:
key = task.params.create_filescore_key(sha256)
task.submission.scan_key = key
return key
def completed(self, sub: DatabaseSubmission):
"""Invoked when notified that a submission has completed."""
# There is only one file in the submissions we have made
sha256 = sub.files[0].sha256
scan_key = sub.scan_key
if not scan_key:
self.log.warning(
f"[{sub.metadata.get('ingest_id', 'unknown')} :: {sha256}] "
f"Submission missing scan key")
scan_key = sub.params.create_filescore_key(sha256)
raw = self.scanning.pop(scan_key)
psid = sub.params.psid
score = sub.max_score
sid = sub.sid
if not raw:
# Some other worker has already popped the scanning queue?
self.log.warning(
f"[{sub.metadata.get('ingest_id', 'unknown')} :: {sha256}] "
f"Submission completed twice")
return scan_key
task = IngestTask(raw)
task.submission.sid = sid
errors = sub.error_count
file_count = sub.file_count
self.counter.increment('submissions_completed')
self.counter.increment('files_completed', increment_by=file_count)
self.counter.increment('bytes_completed', increment_by=task.file_size)
with self.cache_lock:
fs = self.cache[scan_key] = FileScore({
'expiry_ts':
now(self.config.core.ingester.cache_dtl * 24 * 60 * 60),
'errors':
errors,
'psid':
psid,
'score':
score,
'sid':
sid,
'time':
now(),
})
self.datastore.filescore.save(scan_key, fs)
self.finalize(psid, sid, score, task)
def exhaust() -> Iterable[IngestTask]:
while True:
res = self.duplicate_queue.pop(_dup_prefix + scan_key,
blocking=False)
if res is None:
break
res = IngestTask(res)
res.submission.sid = sid
yield res
# You may be tempted to remove the assignment to dups and use the
# value directly in the for loop below. That would be a mistake.
# The function finalize may push on the duplicate queue which we
# are pulling off and so condensing those two lines creates a
# potential infinite loop.
dups = [dup for dup in exhaust()]
for dup in dups:
self.finalize(psid, sid, score, dup)
return scan_key
def send_notification(self, task: IngestTask, failure=None, logfunc=None):
if logfunc is None:
logfunc = self.log.info
if failure:
task.failure = failure
failure = task.failure
if failure:
logfunc("%s: %s", failure, str(task.json()))
if not task.submission.notification.queue:
return
note_queue = _notification_queue_prefix + task.submission.notification.queue
threshold = task.submission.notification.threshold
if threshold is not None and task.score is not None and task.score < threshold:
return
q = self.notification_queues.get(note_queue, None)
if not q:
self.notification_queues[note_queue] = q = NamedQueue(
note_queue, self.redis_persist)
q.push(task.as_primitives())
def expired(self, delta: float, errors) -> bool:
if errors:
return delta >= self.config.core.ingester.incomplete_expire_after_seconds
else:
return delta >= self.config.core.ingester.expire_after
def drop(self, task: IngestTask) -> bool:
priority = task.params.priority
sample_threshold = self.config.core.ingester.sampling_at
dropped = False
if priority <= _min_priority:
dropped = True
else:
for level, rng in self.priority_range.items():
if rng[0] <= priority <= rng[1] and level in sample_threshold:
dropped = must_drop(self.unique_queue.count(*rng),
sample_threshold[level])
break
if not dropped:
if task.file_size > self.config.submission.max_file_size or task.file_size == 0:
dropped = True
if task.params.never_drop or not dropped:
return False
task.failure = 'Skipped'
self._notify_drop(task)
self.counter.increment('skipped')
return True
def _notify_drop(self, task: IngestTask):
self.send_notification(task)
c12n = task.params.classification
expiry = now_as_iso(86400)
sha256 = task.submission.files[0].sha256
self.datastore.save_or_freshen_file(sha256, {'sha256': sha256},
expiry,
c12n,
redis=self.redis)
def is_whitelisted(self, task: IngestTask):
reason, hit = self.get_whitelist_verdict(self.whitelist, task)
hit = {x: dotdump(safe_str(y)) for x, y in hit.items()}
sha256 = task.submission.files[0].sha256
if not reason:
with self.whitelisted_lock:
reason = self.whitelisted.get(sha256, None)
if reason:
hit = 'cached'
if reason:
if hit != 'cached':
with self.whitelisted_lock:
self.whitelisted[sha256] = reason
task.failure = "Whitelisting due to reason %s (%s)" % (dotdump(
safe_str(reason)), hit)
self._notify_drop(task)
self.counter.increment('whitelisted')
return reason
def submit(self, task: IngestTask):
self.submit_client.submit(
submission_obj=task.submission,
completed_queue=COMPLETE_QUEUE_NAME,
)
self.timeout_queue.push(int(now(_max_time)), task.submission.scan_key)
self.log.info(
f"[{task.ingest_id} :: {task.sha256}] Submitted to dispatcher for analysis"
)
def retry(self, task: IngestTask, scan_key: str, ex):
current_time = now()
retries = task.retries + 1
if retries > _max_retries:
trace = ''
if ex:
trace = ': ' + get_stacktrace_info(ex)
self.log.error(
f'[{task.ingest_id} :: {task.sha256}] Max retries exceeded {trace}'
)
self.duplicate_queue.delete(_dup_prefix + scan_key)
elif self.expired(current_time - task.ingest_time.timestamp(), 0):
self.log.info(
f'[{task.ingest_id} :: {task.sha256}] No point retrying expired submission'
)
self.duplicate_queue.delete(_dup_prefix + scan_key)
else:
self.log.info(
f'[{task.ingest_id} :: {task.sha256}] Requeuing ({ex or "unknown"})'
)
task.retries = retries
self.retry_queue.push(int(now(_retry_delay)), task.as_primitives())
def finalize(self, psid: str, sid: str, score: float, task: IngestTask):
self.log.info(f"[{task.ingest_id} :: {task.sha256}] Completed")
if psid:
task.params.psid = psid
task.score = score
task.submission.sid = sid
selected = task.params.services.selected
resubmit_to = task.params.services.resubmit
resubmit_selected = determine_resubmit_selected(selected, resubmit_to)
will_resubmit = resubmit_selected and should_resubmit(score)
if will_resubmit:
task.extended_scan = 'submitted'
task.params.psid = None
if self.is_alert(task, score):
self.log.info(
f"[{task.ingest_id} :: {task.sha256}] Notifying alerter "
f"to {'update' if task.params.psid else 'create'} an alert")
self.alert_queue.push(task.as_primitives())
self.send_notification(task)
if will_resubmit:
self.log.info(
f"[{task.ingest_id} :: {task.sha256}] Resubmitted for extended analysis"
)
task.params.psid = sid
task.submission.sid = None
task.submission.scan_key = None
task.params.services.resubmit = []
task.params.services.selected = resubmit_selected
self.unique_queue.push(task.params.priority, task.as_primitives())
def is_alert(self, task: IngestTask, score: float) -> bool:
if not task.params.generate_alert:
return False
if score < self.config.core.alerter.threshold:
return False
return True
def try_run(self):
self.status = STATUSES.INITIALIZING
# Try to load service class
try:
service_class = load_module_by_path(SERVICE_PATH)
except Exception:
self.log.error("Could not find service in path.")
raise
# Load on-disk manifest for bootstrap/registration
service_manifest = self._load_manifest()
# Register the service
registration = self.tasking_client.register_service(service_manifest)
# Are we just registering?
if not registration['keep_alive'] or REGISTER_ONLY:
self.status = STATUSES.STOPPING
self.stop()
return
# Instantiate the service based of the registration results
self.service_config = registration.get('service_config', {})
self.service = service_class(config=self.service_config.get('config'))
self.service_name = self.service_config['name']
self.service_tool_version = self.service.get_tool_version()
self.metric_factory = MetricsFactory('service',
Metrics,
name=self.service_name,
export_zero=False,
redis=self.redis)
file_required = self.service_config.get('file_required', True)
# Start the service
self.service.start_service()
while self.running:
# Cleanup the working directory
self._cleanup_working_directory()
# Get a task
self.status = STATUSES.WAITING_FOR_TASK
task, _ = self.tasking_client.get_task(
self.client_id, self.service_name,
self.service_config['version'], self.service_tool_version,
self.metric_factory)
if not task:
continue
# Load Task
try:
# Inspect task to ensure submission parameters are given, add defaults where necessary
params = {
x['name']:
task['service_config'].get(x['name'], x['default'])
for x in service_manifest.get('submission_params', [])
}
task['service_config'] = params
service_task = ServiceTask(task)
self.log.info(f"[{service_task.sid}] New task received")
except ValueError as e:
self.log.error(f"Invalid task received: {str(e)}")
continue
# Download file if needed
if file_required:
self.status = STATUSES.DOWNLOADING_FILE
file_path = os.path.join(self.tasking_dir,
service_task.fileinfo.sha256)
received_file_sha256 = None
self.log.info(
f"[{service_task.sid}] Downloading file: {service_task.fileinfo.sha256}"
)
try:
self.filestore.download(service_task.fileinfo.sha256,
file_path)
received_file_sha256 = get_sha256_for_file(file_path)
except FileStoreException:
self.status = STATUSES.FILE_NOT_FOUND
self.log.error(
f"[{service_task.sid}] Requested file not found in the system: {service_task.fileinfo.sha256}"
)
# If the file retrieved is different from what we requested, report the error
if received_file_sha256 and received_file_sha256 != service_task.fileinfo.sha256:
self.status = STATUSES.ERROR_FOUND
self.log.error(
f"[{service_task.sid}] Downloaded ({received_file_sha256}) doesn't match "
f"requested ({service_task.fileinfo.sha256})")
# Process if we're not already in error
if self.status not in [
STATUSES.ERROR_FOUND, STATUSES.FILE_NOT_FOUND
]:
self.status = STATUSES.PROCESSING
self.service.handle_task(service_task)
# Check for the response from the service
result_json = os.path.join(
self.tasking_dir,
f"{service_task.sid}_{service_task.fileinfo.sha256}_result.json"
)
error_json = os.path.join(
self.tasking_dir,
f"{service_task.sid}_{service_task.fileinfo.sha256}_error.json"
)
if os.path.exists(result_json):
self.status = STATUSES.RESULT_FOUND
elif os.path.exists(error_json):
self.status = STATUSES.ERROR_FOUND
else:
self.status = STATUSES.ERROR_FOUND
error_json = None
# Handle the service response
if self.status == STATUSES.RESULT_FOUND:
self.log.info(
f"[{service_task.sid}] Task successfully completed")
try:
self._handle_task_result(result_json, service_task)
except RuntimeError as re:
if is_recoverable_runtime_error(re):
self.log.info(
f"[{service_task.sid}] Service trying to use a threadpool during shutdown, "
"sending recoverable error.")
self._handle_task_error(service_task)
else:
raise
elif self.status == STATUSES.ERROR_FOUND:
self.log.info(
f"[{service_task.sid}] Task completed with errors")
self._handle_task_error(service_task,
error_json_path=error_json)
elif self.status == STATUSES.FILE_NOT_FOUND:
self.log.info(
f"[{service_task.sid}] Task completed with errors due to missing file from filestore"
)
self._handle_task_error(service_task,
status="FAIL_NONRECOVERABLE",
error_type="EXCEPTION")