def __init__(self, logger=None, datastore=None, redis=None, persistent_redis=None):
super().__init__('assemblyline.ingester.input', logger=logger)
config = forge.get_config()
# Connect to all sorts of things
datastore = datastore or forge.get_datastore(config)
classification_engine = forge.get_classification()
# Initialize the ingester specific resources
self.ingester = Ingester(datastore=datastore, classification=classification_engine, logger=self.log,
redis=redis, persistent_redis=persistent_redis)
if config.core.metrics.apm_server.server_url is not None:
self.log.info(f"Exporting application metrics to: {config.core.metrics.apm_server.server_url}")
elasticapm.instrument()
self.apm_client = elasticapm.Client(server_url=config.core.metrics.apm_server.server_url,
service_name="ingester")
else:
self.apm_client = None
def __init__(self):
super().__init__('assemblyline.workflow')
self.config = forge.get_config()
self.datastore = forge.get_datastore(self.config)
self.start_ts = f"{self.datastore.ds.now}/{self.datastore.ds.day}-1{self.datastore.ds.day}"
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="workflow")
else:
self.apm_client = None
def __init__(self, datastore=None, filestore=None):
super().__init__('assemblyline.randomservice')
self.config = forge.get_config()
self.datastore = datastore or forge.get_datastore()
self.filestore = filestore or forge.get_filestore()
self.client_id = get_random_id()
self.service_state_hash = ExpiringHash(SERVICE_STATE_HASH, ttl=30 * 60)
self.counters = {
n: MetricsFactory('service', Metrics, name=n, config=self.config)
for n in self.datastore.service_delta.keys()
}
self.queues = [
forge.get_service_queue(name)
for name in self.datastore.service_delta.keys()
]
self.dispatch_client = DispatchClient(self.datastore)
self.service_info = CachedObject(self.datastore.list_all_services,
kwargs={'as_obj': False})
def __init__(self,
datastore: AssemblylineDatastore = None,
filestore: FileStore = None,
config=None,
redis=None,
identify=None):
self.log = logging.getLogger('assemblyline.submission_client')
self.config = config or forge.CachedObject(forge.get_config)
self.datastore = datastore or forge.get_datastore(self.config)
self.filestore = filestore or forge.get_filestore(self.config)
self.redis = redis
if identify:
self.cleanup = False
else:
self.cleanup = True
self.identify = identify or forge.get_identify(
config=self.config, datastore=self.datastore, use_cache=True)
# A client for interacting with the dispatcher
self.dispatcher = DispatchClient(datastore, redis)
def __init__(self,
use_cache=True,
config=None,
datastore=None,
log=None) -> None:
self.log = log or logging.getLogger('assemblyline.identify')
self.config = None
self.datastore = None
self.use_cache = use_cache
self.custom = re.compile(r"^custom: ", re.IGNORECASE)
self.lock = threading.Lock()
self.yara_default_externals = {'mime': '', 'magic': '', 'type': ''}
# If cache is use, load the config and datastore objects to load potential items from cache
if self.use_cache:
self.log.info("Using cache with identify")
self.config = config or get_config()
self.datastore = datastore or get_datastore(config)
# Load all data for the first time
self._load_magic_file()
self._load_yara_file()
self._load_magic_patterns()
self._load_trusted_mimes()
# Register hot reloader
if self.use_cache:
self.reload_map = {
'magic': self._load_magic_file,
'mimes': self._load_trusted_mimes,
'patterns': self._load_magic_patterns,
'yara': self._load_yara_file
}
self.reload_watcher = EventWatcher()
self.reload_watcher.register('system.identify',
self._handle_reload_event)
self.reload_watcher.start()
else:
self.reload_watcher = None
self.reload_map = {}
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 __init__(self,
datastore=None,
redis=None,
redis_persist=None,
logger=None):
super().__init__('assemblyline.dispatcher.submissions', logger)
config = forge.get_config()
datastore = datastore or forge.get_datastore(config)
self.dispatcher = Dispatcher(logger=self.log,
redis=redis,
redis_persist=redis_persist,
datastore=datastore)
if config.core.metrics.apm_server.server_url is not None:
self.log.info(
f"Exporting application metrics to: {config.core.metrics.apm_server.server_url}"
)
elasticapm.instrument()
self.apm_client = elasticapm.Client(
server_url=config.core.metrics.apm_server.server_url,
service_name="dispatcher")
else:
self.apm_client = None
def __init__(self):
super().__init__('assemblyline.alerter')
# Publish counters to the metrics sink.
self.counter = MetricsFactory('alerter', Metrics)
self.datastore = forge.get_datastore(self.config)
self.persistent_redis = get_client(
host=self.config.core.redis.persistent.host,
port=self.config.core.redis.persistent.port,
private=False,
)
self.process_alert_message = forge.get_process_alert_message()
self.running = False
self.alert_queue = NamedQueue(ALERT_QUEUE_NAME, self.persistent_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="alerter")
else:
self.apm_client = None
def restore_worker(worker_id, instance_id, working_dir):
datastore = forge.get_datastore(archive_access=True)
done_queue = NamedQueue(f"r-done-{instance_id}", ttl=1800)
with open(os.path.join(working_dir, "backup.part%s" % worker_id),
"rb") as input_file:
for line in input_file:
bucket_name, key, data = json.loads(line)
success = True
try:
collection = datastore.get_collection(bucket_name)
collection.save(key, data)
except Exception:
success = False
done_queue.push({
"success": success,
"missing": False,
"bucket_name": bucket_name,
"key": key
})
done_queue.push({"stopped": True})
RATE_LIMITER = Counters(prefix="quota", host=redis, track_counters=True)
# End of Configuration
#################################################################
#################################################################
# Prepare loggers
config.logging.log_to_console = config.logging.log_to_console or DEBUG
al_log.init_logging('svc', config=config)
LOGGER = logging.getLogger('assemblyline.svc')
LOGGER.debug('Logger ready!')
# End of prepare logger
#################################################################
#################################################################
# Global instances
STORAGE = forge.get_datastore(config=config)
FILESTORE = forge.get_filestore(config=config)
LOCK = threading.Lock()
TASKING_CLIENT = TaskingClient(datastore=STORAGE,
filestore=FILESTORE,
redis=redis,
redis_persist=redis_persist)
SAFELIST_CLIENT = SafelistClient(datastore=STORAGE)
# End global
#################################################################
def __init__(self, log, **kw):
super(AL, self).__init__(log, **kw)
self.datastore = forge.get_datastore()
def create_extra_data(log=None, ds=None, fs=None):
ds = ds or forge.get_datastore()
fs = fs or forge.get_filestore()
log.info("\nCreating 10 Submissions...")
submissions = []
for _ in range(10):
s = create_submission(ds, fs, log=log)
submissions.append(s)
log.info("\nCreating 50 Alerts...")
create_alerts(ds, submission_list=submissions, log=log)
if __name__ == "__main__":
datastore = forge.get_datastore()
logger = PrintLogger()
create_basic_data(log=logger,
ds=datastore,
svc="nosvc" not in sys.argv,
sigs="nosigs" not in sys.argv,
reset="reset" in sys.argv)
if "full" in sys.argv:
create_extra_data(log=logger, ds=datastore)
if "alerts" in sys.argv:
create_alerts(datastore, alert_count=1000, log=logger)
logger.info("\nDone.")
def import_bundle(path,
working_dir=WORK_DIR,
min_classification=Classification.UNRESTRICTED,
allow_incomplete=False):
with forge.get_datastore() as datastore:
current_working_dir = os.path.join(working_dir, get_random_id())
res_file = os.path.join(current_working_dir, "results.json")
try:
os.makedirs(current_working_dir)
except Exception:
pass
# Extract the bundle
try:
subprocess.check_call(
["tar", "-zxf", path, "-C", current_working_dir])
except subprocess.CalledProcessError:
raise BundlingException(
"Bundle decompression failed. Not a valid bundle...")
with open(res_file, 'rb') as fh:
data = json.load(fh)
submission = data['submission']
results = data['results']
files = data['files']
errors = data['errors']
try:
sid = submission['sid']
# Check if we have all the service results
for res_key in submission['results']:
if res_key not in results['results'].keys(
) and not allow_incomplete:
raise IncompleteBundle(
"Incomplete results in bundle. Skipping %s..." % sid)
# Check if we have all files
for sha256 in list(set([x[:64] for x in submission['results']])):
if sha256 not in files['infos'].keys(
) and not allow_incomplete:
raise IncompleteBundle(
"Incomplete files in bundle. Skipping %s..." % sid)
# Check if we all errors
for err_key in submission['errors']:
if err_key not in errors['errors'].keys(
) and not allow_incomplete:
raise IncompleteBundle(
"Incomplete errors in bundle. Skipping %s..." % sid)
if datastore.submission.get(sid, as_obj=False):
raise SubmissionAlreadyExist("Submission %s already exists." %
sid)
# Make sure bundle's submission meets minimum classification and save the submission
submission['classification'] = Classification.max_classification(
submission['classification'], min_classification)
submission.update(
Classification.get_access_control_parts(
submission['classification']))
datastore.submission.save(sid, submission)
# Make sure files meet minimum classification and save the files
with forge.get_filestore() as filestore:
for f, f_data in files['infos'].items():
f_classification = Classification.max_classification(
f_data['classification'], min_classification)
datastore.save_or_freshen_file(f,
f_data,
f_data['expiry_ts'],
f_classification,
cl_engine=Classification)
try:
filestore.upload(os.path.join(current_working_dir, f),
f)
except IOError:
pass
# Make sure results meet minimum classification and save the results
for key, res in results['results'].items():
if key.endswith(".e"):
datastore.emptyresult.save(key,
{"expiry_ts": res['expiry_ts']})
else:
res['classification'] = Classification.max_classification(
res['classification'], min_classification)
datastore.result.save(key, res)
# Make sure errors meet minimum classification and save the errors
for ekey, err in errors['errors'].items():
datastore.error.save(ekey, err)
finally:
# Perform working dir cleanup
try:
os.remove(path)
except Exception:
pass
try:
shutil.rmtree(current_working_dir, ignore_errors=True)
except Exception:
pass
def create_bundle(sid, working_dir=WORK_DIR):
with forge.get_datastore() as datastore:
temp_bundle_file = f"bundle_{get_random_id()}"
current_working_dir = os.path.join(working_dir, temp_bundle_file)
try:
submission = datastore.submission.get(sid, as_obj=False)
if submission is None:
raise SubmissionNotFound(
"Can't find submission %s, skipping." % sid)
else:
target_file = os.path.join(working_dir,
f"{temp_bundle_file}.tgz")
try:
os.makedirs(current_working_dir)
except Exception as e:
if isinstance(PermissionError, e):
raise
pass
# Create file information data
file_tree = datastore.get_or_create_file_tree(
submission, config.submission.max_extraction_depth)['tree']
flatten_tree = list(
set(
recursive_flatten_tree(file_tree) +
[r[:64] for r in submission.get("results", [])]))
file_infos, _ = get_file_infos(copy(flatten_tree), datastore)
# Add bundling metadata
if 'bundle.source' not in submission['metadata']:
submission['metadata']['bundle.source'] = config.ui.fqdn
if Classification.enforce and 'bundle.classification' not in submission[
'metadata']:
submission['metadata'][
'bundle.classification'] = submission['classification']
data = {
'submission':
submission,
'files': {
"list": flatten_tree,
"tree": file_tree,
"infos": file_infos
},
'results':
get_results(submission.get("results", []), file_infos,
datastore),
'errors':
get_errors(submission.get("errors", []), datastore)
}
# Save result files
with open(os.path.join(current_working_dir, "results.json"),
"w") as fp:
json.dump(data, fp)
# Download all related files
with forge.get_filestore() as filestore:
for sha256 in flatten_tree:
try:
filestore.download(
sha256,
os.path.join(current_working_dir, sha256))
except FileStoreException:
pass
# Create the bundle
subprocess.check_call("tar czf %s *" % target_file,
shell=True,
cwd=current_working_dir)
return target_file
except Exception as e:
raise BundlingException(
"Could not bundle submission '%s'. [%s: %s]" %
(sid, type(e).__name__, str(e)))
finally:
if current_working_dir:
subprocess.check_call(["rm", "-rf", current_working_dir])
"config_name", "servicename", "vm"
]
AUDIT_LOG = logging.getLogger('assemblyline.ui.audit')
LOGGER = logging.getLogger('assemblyline.ui')
if AUDIT:
AUDIT_LOG.setLevel(logging.INFO)
if DEBUG:
if not os.path.exists(config.logging.log_directory):
os.makedirs(config.logging.log_directory)
fh = logging.FileHandler(
os.path.join(config.logging.log_directory, 'alui_audit.log'))
fh.setLevel(logging.INFO)
fh.setFormatter(logging.Formatter(AL_LOG_FORMAT))
AUDIT_LOG.addHandler(fh)
AUDIT_LOG.debug('Audit logger ready!')
LOGGER.debug('Logger ready!')
# End of prepare logger
#################################################################
#################################################################
# Global instances
STORAGE = forge.get_datastore(archive_access=True)
# End global
#################################################################
def create_bundle(sid, working_dir=WORK_DIR, use_alert=False):
with forge.get_datastore() as datastore:
temp_bundle_file = f"bundle_{get_random_id()}"
current_working_dir = os.path.join(working_dir, temp_bundle_file)
target_file = os.path.join(working_dir, f"{temp_bundle_file}.cart")
tgz_file = os.path.join(working_dir, f"{temp_bundle_file}.tgz")
try:
if use_alert:
alert = datastore.alert.get(sid, as_obj=False)
if alert is None:
raise AlertNotFound("Can't find alert %s, skipping." % sid)
sid = alert['sid']
else:
alert = None
submission = datastore.submission.get(sid, as_obj=False)
if submission is None and alert is None:
raise SubmissionNotFound(
"Can't find submission %s, skipping." % sid)
else:
try:
os.makedirs(current_working_dir)
except PermissionError:
raise
except Exception:
pass
data = {}
if submission:
# Create file information data
file_tree = datastore.get_or_create_file_tree(
submission,
config.submission.max_extraction_depth)['tree']
flatten_tree = list(
set(
recursive_flatten_tree(file_tree) +
[r[:64] for r in submission.get("results", [])]))
file_infos, _ = get_file_infos(copy(flatten_tree),
datastore)
# Add bundling metadata
if 'bundle.source' not in submission['metadata']:
submission['metadata'][
'bundle.source'] = config.ui.fqdn
if 'bundle.created' not in submission['metadata']:
submission['metadata']['bundle.created'] = now_as_iso()
if Classification.enforce and 'bundle.classification' not in submission[
'metadata']:
submission['metadata'][
'bundle.classification'] = submission[
'classification']
data.update({
'submission':
submission,
'files': {
"list": flatten_tree,
"tree": file_tree,
"infos": file_infos
},
'results':
get_results(submission.get("results", []), file_infos,
datastore),
'errors':
get_errors(submission.get("errors", []), datastore)
})
# Download all related files
with forge.get_filestore() as filestore:
for sha256 in flatten_tree:
try:
filestore.download(
sha256,
os.path.join(current_working_dir, sha256))
except FileStoreException:
pass
if alert:
if 'bundle.source' not in alert['metadata']:
alert['metadata']['bundle.source'] = config.ui.fqdn
if 'bundle.created' not in alert['metadata']:
alert['metadata']['bundle.created'] = now_as_iso()
if Classification.enforce and 'bundle.classification' not in alert[
'metadata']:
alert['metadata']['bundle.classification'] = alert[
'classification']
data['alert'] = alert
# Save result files
with open(os.path.join(current_working_dir, "results.json"),
"w") as fp:
json.dump(data, fp)
# Create the bundle
subprocess.check_call("tar czf %s *" % tgz_file,
shell=True,
cwd=current_working_dir)
with open(target_file, 'wb') as oh:
with open(tgz_file, 'rb') as ih:
pack_stream(ih, oh, {
'al': {
"type": BUNDLE_TYPE
},
'name': f"{sid}.tgz"
})
return target_file
except (SubmissionNotFound, AlertNotFound):
raise
except Exception as e:
raise BundlingException(
"Could not bundle submission '%s'. [%s: %s]" %
(sid, type(e).__name__, str(e)))
finally:
if os.path.exists(current_working_dir):
subprocess.check_call(["rm", "-rf", current_working_dir])
if os.path.exists(tgz_file):
os.unlink(tgz_file)
def init():
global DATASTORE
DATASTORE = forge.get_datastore(archive_access=True)
signal.signal(signal.SIGINT, signal.SIG_IGN)
def backup_worker(worker_id: str, instance_id: str, working_dir: str):
datastore = forge.get_datastore(archive_access=True)
worker_queue: NamedQueue[dict[str,
Any]] = NamedQueue(f"r-worker-{instance_id}",
ttl=1800)
done_queue: NamedQueue[dict[str,
Any]] = NamedQueue(f"r-done-{instance_id}",
ttl=1800)
hash_queue: Hash[str] = Hash(f"r-hash-{instance_id}")
stopping = False
with open(os.path.join(working_dir, "backup.part%s" % worker_id),
"w+") as backup_file:
while True:
data = worker_queue.pop(timeout=1)
if data is None:
if stopping:
break
continue
if data.get('stop', False):
if not stopping:
stopping = True
else:
time.sleep(round(random.uniform(0.050, 0.250), 3))
worker_queue.push(data)
continue
missing = False
success = True
try:
to_write = datastore.get_collection(data['bucket_name']).get(
data['key'], as_obj=False)
if to_write:
if data.get('follow_keys', False):
for bucket, bucket_key, getter in FOLLOW_KEYS.get(
data['bucket_name'], []):
for key in getter(to_write.get(bucket_key, None)):
hash_key = "%s_%s" % (bucket, key)
if not hash_queue.exists(hash_key):
hash_queue.add(hash_key, "True")
worker_queue.push({
"bucket_name": bucket,
"key": key,
"follow_keys": True
})
backup_file.write(
json.dumps((data['bucket_name'], data['key'],
to_write)) + "\n")
else:
missing = True
except Exception:
success = False
done_queue.push({
"success": success,
"missing": missing,
"bucket_name": data['bucket_name'],
"key": data['key']
})
done_queue.push({"stopped": True})
def import_bundle(path,
working_dir=WORK_DIR,
min_classification=Classification.UNRESTRICTED,
allow_incomplete=False,
rescan_services=None,
exist_ok=False,
cleanup=True,
identify=None):
with forge.get_datastore(archive_access=True) as datastore:
current_working_dir = os.path.join(working_dir, get_random_id())
res_file = os.path.join(current_working_dir, "results.json")
try:
os.makedirs(current_working_dir)
except Exception:
pass
with open(path, 'rb') as original_file:
if is_cart(original_file.read(256)):
original_file.seek(0)
extracted_fd, extracted_path = tempfile.mkstemp()
extracted_file = os.fdopen(extracted_fd, 'wb')
try:
hdr, _ = unpack_stream(original_file, extracted_file)
if hdr.get('al', {}).get('type', 'unknown') != BUNDLE_TYPE:
raise BundlingException(
f"Not a valid CaRTed bundle, should be of type: {BUNDLE_TYPE}"
)
finally:
extracted_file.close()
else:
extracted_path = path
# Extract the bundle
try:
subprocess.check_call(
["tar", "-zxf", extracted_path, "-C", current_working_dir])
except subprocess.CalledProcessError:
raise BundlingException(
"Bundle decompression failed. Not a valid bundle...")
with open(res_file, 'rb') as fh:
data = json.load(fh)
alert = data.get('alert', None)
submission = data.get('submission', None)
try:
if submission:
sid = submission['sid']
# Load results, files and errors
results = data.get('results', None)
files = data.get('files', None)
errors = data.get('errors', None)
# Check if we have all the service results
for res_key in submission['results']:
if results is None or (res_key
not in results['results'].keys()
and not allow_incomplete):
raise IncompleteBundle(
"Incomplete results in bundle. Skipping %s..." %
sid)
# Check if we have all files
for sha256 in list(set([x[:64]
for x in submission['results']])):
if files is None or (sha256 not in files['infos'].keys()
and not allow_incomplete):
raise IncompleteBundle(
"Incomplete files in bundle. Skipping %s..." % sid)
# Check if we all errors
for err_key in submission['errors']:
if errors is None or (err_key
not in errors['errors'].keys()
and not allow_incomplete):
raise IncompleteBundle(
"Incomplete errors in bundle. Skipping %s..." %
sid)
# Check if the submission does not already exist
if not datastore.submission.exists(sid):
# Make sure bundle's submission meets minimum classification and save the submission
submission[
'classification'] = Classification.max_classification(
submission['classification'], min_classification)
submission.setdefault('metadata', {})
submission['metadata']['bundle.loaded'] = now_as_iso()
submission['metadata'].pop('replay', None)
submission.update(
Classification.get_access_control_parts(
submission['classification']))
if not rescan_services:
# Save the submission in the system
datastore.submission.save(sid, submission)
# Make sure files meet minimum classification and save the files
with forge.get_filestore() as filestore:
for f, f_data in files['infos'].items():
f_classification = Classification.max_classification(
f_data['classification'], min_classification)
datastore.save_or_freshen_file(
f,
f_data,
f_data['expiry_ts'],
f_classification,
cl_engine=Classification)
try:
filestore.upload(
os.path.join(current_working_dir, f), f)
except IOError:
pass
# Make sure results meet minimum classification and save the results
for key, res in results['results'].items():
if key.endswith(".e"):
datastore.emptyresult.save(
key, {"expiry_ts": res['expiry_ts']})
else:
res['classification'] = Classification.max_classification(
res['classification'], min_classification)
datastore.result.save(key, res)
# Make sure errors meet minimum classification and save the errors
for ekey, err in errors['errors'].items():
datastore.error.save(ekey, err)
# Start the rescan
if rescan_services and SubmissionClient:
extracted_file_infos = {
k: {
vk: v[vk]
for vk in [
'magic', 'md5', 'mime', 'sha1',
'sha256', 'size', 'type'
]
}
for k, v in files['infos'].items()
if k in files['list']
}
with SubmissionClient(datastore=datastore,
filestore=filestore,
config=config,
identify=identify) as sc:
sc.rescan(submission, results['results'],
extracted_file_infos, files['tree'],
list(errors['errors'].keys()),
rescan_services)
elif not exist_ok:
raise SubmissionAlreadyExist(
"Submission %s already exists." % sid)
# Save alert if present and does not exist
if alert and not datastore.alert.exists(alert['alert_id']):
alert['classification'] = Classification.max_classification(
alert['classification'], min_classification)
alert.setdefault('metadata', {})
alert['metadata']['bundle.loaded'] = now_as_iso()
alert['metadata'].pop('replay', None)
alert['workflows_completed'] = False
datastore.alert.save(alert['alert_id'], alert)
return submission
finally:
if extracted_path != path and os.path.exists(extracted_path):
os.remove(extracted_path)
if cleanup and os.path.exists(path):
os.remove(path)
if os.path.exists(current_working_dir):
shutil.rmtree(current_working_dir, ignore_errors=True)
def core(request, redis, filestore, config):
from assemblyline.common import log as al_log
al_log.init_logging("simulation")
fields = CoreSession()
fields.redis = redis
fields.ds = ds = forge.get_datastore()
fields.config = config
forge.config_singletons[False, None] = fields.config
threads = []
fields.filestore = filestore
threads: List[ServerBase] = [
# Start the ingester components
IngesterInput(datastore=ds, redis=redis, persistent_redis=redis),
IngesterSubmitter(datastore=ds, redis=redis, persistent_redis=redis),
IngesterInternals(datastore=ds, redis=redis, persistent_redis=redis),
# Start the dispatcher
FileDispatchServer(datastore=ds, redis=redis, redis_persist=redis),
SubmissionDispatchServer(datastore=ds,
redis=redis,
redis_persist=redis),
# Start plumber
Plumber(datastore=ds, redis=redis, redis_persist=redis, delay=0.5),
]
stages = get_service_stage_hash(redis)
ingester_input_thread: IngesterInput = threads[0]
fields.ingest = ingester_input_thread
fields.ingest_queue = ingester_input_thread.ingester.ingest_queue
ds.ds.service = MockCollection(Service)
ds.ds.service_delta = MockCollection(Service)
ds.service.save('pre_0', dummy_service('pre', 'EXTRACT'))
ds.service_delta.save('pre', dummy_service('pre', 'EXTRACT'))
stages.set('pre', ServiceStage.Running)
threads.append(MockService('pre', ds, redis, filestore))
fields.pre_service = threads[-1]
ds.service.save('core-a_0', dummy_service('core-a', 'CORE'))
ds.service_delta.save('core-a', dummy_service('core-a', 'CORE'))
stages.set('core-a', ServiceStage.Running)
threads.append(MockService('core-a', ds, redis, filestore))
ds.service.save('core-b_0', dummy_service('core-b', 'CORE'))
ds.service_delta.save('core-b', dummy_service('core-b', 'CORE'))
threads.append(MockService('core-b', ds, redis, filestore))
stages.set('core-b', ServiceStage.Running)
ds.service.save('finish_0', dummy_service('finish', 'POST'))
ds.service_delta.save('finish', dummy_service('finish', 'POST'))
threads.append(MockService('finish', ds, redis, filestore))
stages.set('finish', ServiceStage.Running)
for t in threads:
t.daemon = True
t.start()
def stop_core():
[tr.close() for tr in threads]
[tr.stop() for tr in threads]
[tr.raising_join() for tr in threads]
request.addfinalizer(stop_core)
return fields
def __init__(self, datastore: AssemblylineDatastore = None, config=None):
self.log = logging.getLogger('assemblyline.safelist_client')
self.config = config or forge.CachedObject(forge.get_config)
self.datastore = datastore or forge.get_datastore(self.config)
