def try_run(self):
while self.running:
task = self.dispatch_client.request_work('worker',
self.service_name,
'0',
timeout=1)
if not task:
continue
print(self.service_name, 'has received a job', task.sid)
file = self.filestore.get(task.fileinfo.sha256)
instructions = json.loads(file)
instructions = instructions.get(self.service_name, {})
print(self.service_name, 'following instruction:', instructions)
hits = self.hits[task.fileinfo.sha256] = self.hits.get(
task.fileinfo.sha256, 0) + 1
if instructions.get('semaphore', False):
_global_semaphore.acquire(blocking=True,
timeout=instructions['semaphore'])
continue
if 'drop' in instructions:
if instructions['drop'] >= hits:
self.drops[task.fileinfo.sha256] = self.drops.get(
task.fileinfo.sha256, 0) + 1
continue
if instructions.get('failure', False):
error = Error(instructions['error'])
error.sha256 = task.fileinfo.sha256
self.dispatch_client.service_failed(task.sid,
error=error,
error_key=get_random_id())
continue
result_data = {
'archive_ts': time.time() + 300,
'classification': 'U',
'response': {
'service_version': '0',
'service_tool_version': '0',
'service_name': self.service_name,
},
'result': {},
'sha256': task.fileinfo.sha256,
'expiry_ts': time.time() + 600
}
result_data.update(instructions.get('result', {}))
result_data['response'].update(instructions.get('response', {}))
result = Result(result_data)
result_key = instructions.get('result_key', get_random_id())
self.dispatch_client.service_finished(task.sid, result_key, result)
def test_ingest_sha(datastore, client):
file_id = random_id_from_collection(datastore, 'file')
res = client.ingest(alert=True,
sha256=file_id,
metadata={
"file_id": get_random_id(),
"comment": "test"
},
nq=get_random_id(),
nt=100)
assert res.get('ingest_id', None) is not None
def test_live_namespace(datastore, sio):
wq_data = {'wq_id': get_random_id()}
wq = NamedQueue(wq_data['wq_id'], private=True)
start_msg = {'status_code': 200, 'msg': "Start listening..."}
stop_msg = {
'status_code': 200,
'msg': "All messages received, closing queue..."
}
cachekey_msg = {'status_code': 200, 'msg': get_random_id()}
cachekeyerr_msg = {'status_code': 200, 'msg': get_random_id()}
test_res_array = []
@sio.on('start', namespace='/live_submission')
def on_start(data):
test_res_array.append(('on_start', data == start_msg))
@sio.on('stop', namespace='/live_submission')
def on_stop(data):
test_res_array.append(('on_stop', data == stop_msg))
@sio.on('cachekey', namespace='/live_submission')
def on_cachekey(data):
test_res_array.append(('on_cachekey', data == cachekey_msg))
@sio.on('cachekeyerr', namespace='/live_submission')
def on_stop(data):
test_res_array.append(('on_cachekeyerr', data == cachekeyerr_msg))
try:
sio.emit('listen', wq_data, namespace='/live_submission')
sio.sleep(1)
wq.push({"status": "START"})
wq.push({"status": "OK", "cache_key": cachekey_msg['msg']})
wq.push({"status": "FAIL", "cache_key": cachekeyerr_msg['msg']})
wq.push({"status": "STOP"})
start_time = time.time()
while len(test_res_array) < 4 and time.time() - start_time < 5:
sio.sleep(0.1)
assert len(test_res_array) == 4
for test, result in test_res_array:
if not result:
pytest.fail(f"{test} failed.")
finally:
sio.disconnect()
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 __init__(self,
name,
host=None,
export_interval_secs=None,
counter_type=None,
config=None,
redis=None,
counter_names=None,
timer_names=None,
export_zero=True):
config = config or forge.get_config()
self.channel = forge.get_metrics_sink(redis)
self.export_interval = export_interval_secs or config.core.metrics.export_interval
self.name = name
self.host = host or get_random_id()
self.type = counter_type or name
self.export_zero = export_zero
self.counter_schema = set(counter_names)
self.timer_schema = set(timer_names)
self.counts = Counters({key: 0 for key in self.counter_schema})
self.values = {}
self.lock = threading.Lock()
self.scheduler = None
self.reset()
assert self.channel
assert (self.export_interval > 0)
def generate_conf_key(service_tool_version=None, task=None):
ignore_salt = None
service_config = None
submission_params_str = None
if task is not None:
service_config = json.dumps(sorted(task.service_config.items()))
submission_params = {
"deep_scan": task.deep_scan,
"max_files": task.max_files,
"min_classification": task.min_classification.value
}
submission_params_str = json.dumps(sorted(submission_params.items()))
if task.ignore_cache:
ignore_salt = get_random_id()
if service_tool_version is None and \
service_config is None and \
submission_params_str is None and \
ignore_salt is None:
return "0"
total_str = f"{service_tool_version}_{service_config}_{submission_params_str}_{ignore_salt}".encode(
'utf-8')
partial_md5 = hashlib.md5(
(str(total_str).encode('utf-8'))).hexdigest()[:16]
return baseconv.base62.encode(int(partial_md5, 16))
def __init__(self,
working_dir,
worker_count=50,
spawn_workers=True,
use_threading=False,
logger=None):
self.working_dir = working_dir
self.datastore = forge.get_datastore(archive_access=True)
self.logger = logger
self.plist = []
self.use_threading = use_threading
self.instance_id = get_random_id()
self.worker_queue = NamedQueue(f"r-worker-{self.instance_id}",
ttl=1800)
self.done_queue = NamedQueue(f"r-done-{self.instance_id}", ttl=1800)
self.hash_queue = Hash(f"r-hash-{self.instance_id}")
self.bucket_error = []
self.VALID_BUCKETS = sorted(list(
self.datastore.ds.get_models().keys()))
self.worker_count = worker_count
self.spawn_workers = spawn_workers
self.total_count = 0
self.error_map_count = {}
self.missing_map_count = {}
self.map_count = {}
self.last_time = 0
self.last_count = 0
self.error_count = 0
def test_user_quota_tracker(redis_connection):
if redis_connection:
from assemblyline.remote.datatypes.user_quota_tracker import UserQuotaTracker
max_quota = 3
timeout = 2
name = get_random_id()
uqt = UserQuotaTracker('test-quota', timeout=timeout)
# First 0 to max_quota items should succeed
for _ in range(max_quota):
assert uqt.begin(name, max_quota) is True
# All other items should fail until items timeout
for _ in range(max_quota):
assert uqt.begin(name, max_quota) is False
# if you remove and item only one should be able to go in
uqt.end(name)
assert uqt.begin(name, max_quota) is True
assert uqt.begin(name, max_quota) is False
# if you wait the timeout, all items can go in
time.sleep(timeout + 1)
for _ in range(max_quota):
assert uqt.begin(name, max_quota) is True
def __init__(self,
working_dir: str,
worker_count: int = 50,
spawn_workers: bool = True,
use_threading: bool = False,
logger: logging.Logger = None):
self.working_dir = working_dir
self.datastore = forge.get_datastore(archive_access=True)
self.logger = logger
self.plist: list[Process] = []
self.use_threading = use_threading
self.instance_id = get_random_id()
self.worker_queue: NamedQueue[dict[str, Any]] = NamedQueue(
f"r-worker-{self.instance_id}", ttl=1800)
self.done_queue: NamedQueue[dict[str, Any]] = NamedQueue(
f"r-done-{self.instance_id}", ttl=1800)
self.hash_queue: Hash[str] = Hash(f"r-hash-{self.instance_id}")
self.bucket_error: list[str] = []
self.valid_buckets: list[str] = sorted(
list(self.datastore.ds.get_models().keys()))
self.worker_count = worker_count
self.spawn_workers = spawn_workers
self.total_count = 0
self.error_map_count: dict[str, int] = {}
self.missing_map_count: dict[str, int] = {}
self.map_count: dict[str, int] = {}
self.last_time: float = 0
self.last_count = 0
self.error_count = 0
def put(self, path, content):
path = self.normalize(path)
dirname = os.path.dirname(path)
filename = os.path.basename(path)
tempname = get_random_id()
temppath = _join(dirname, tempname)
finalpath = _join(dirname, filename)
assert (finalpath == path)
self.makedirs(dirname)
fh = None
try:
fh = open(temppath, "wb")
return fh.write(content)
finally:
if fh:
fh.close()
try:
shutil.move(temppath, finalpath)
except shutil.Error:
pass
assert (self.exists(path))
def test_submit_content(datastore, client):
content = get_random_phrase(wmin=15, wmax=50).encode()
fname = "{}.txt".format(get_random_id())
res = client.submit(content=content, fname=fname)
assert res['sid'] is not None
assert res['files'][0]['name'] == fname
assert res == datastore.submission.get(res['sid'], as_obj=False)
def __init__(self, name, timeout, host=None, port=None):
self.uuid = get_random_id()
self.c = get_client(host, port, False)
self.lock_release = '-'.join(('lock', str(timeout), name, 'released'))
self.lock_holder = '-'.join(('lock', str(timeout), name, 'holder'))
self.timeout = timeout
self._acquire = self.c.register_script(lock_acquire_script)
self._release = self.c.register_script(lock_release_script)
def test_submit_sha(datastore, client):
file_id = random_id_from_collection(datastore, 'file')
metadata = {"file_id": get_random_id(), "comment": "test"}
res = client.submit(sha256=file_id, metadata=metadata)
assert res['sid'] is not None
assert res['files'][0]['sha256'] == file_id
assert res['metadata'] == metadata
assert res == datastore.submission.get(res['sid'], as_obj=False)
def create_workflows(ds, log=None):
for _ in range(20):
w_id = get_random_id()
ds.workflow.save(w_id, random_model_obj(Workflow))
if log:
log.info(f'\t{w_id}')
ds.workflow.commit()
def test_watcher(redis_connection):
redis_connection.time = RedisTime()
rds = redis_connection
queue_name = get_random_id()
out_queue = NamedQueue(queue_name, rds)
try:
# Create a server and hijack its running flag and the current time in 'redis'
client = WatcherClient(rds)
server = WatcherServer(rds, rds)
server.running = ToggleTrue()
rds.time.current = 0
assert out_queue.length() == 0
# Send a simple event to occur soon
client.touch(10, 'one-second', queue_name, {'first': 'one'})
server.try_run()
assert out_queue.length() == 0 # Nothing yet
rds.time.current = 12 # Jump forward 12 seconds
server.try_run()
assert out_queue.length() == 1
assert out_queue.pop() == {'first': 'one'}
# Send a simple event to occur soon, then change our mind
client.touch(10, 'one-second', queue_name, {'first': 'one'})
client.touch(20, 'one-second', queue_name, {'first': 'one'})
server.try_run()
assert out_queue.length() == 0 # Nothing yet
# Set events to occur, in inverse order, reuse a key, overwrite content and timeout
client.touch(200, 'one-second', queue_name, {'first': 'last'})
client.touch(100, '100-second', queue_name, {'first': '100'})
client.touch(50, '50-second', queue_name, {'first': '50'})
server.try_run()
assert out_queue.length() == 0 # Nothing yet
for _ in range(15):
rds.time.current += 20
server.try_run()
assert out_queue.length() == 3
assert out_queue.pop() == {'first': '50'}
assert out_queue.pop() == {'first': '100'}
assert out_queue.pop() == {'first': 'last'}
# Send a simple event to occur soon, then stop it
rds.time.current = 0
client.touch(10, 'one-second', queue_name, {'first': 'one'})
server.try_run()
assert out_queue.length() == 0 # Nothing yet
client.clear('one-second')
rds.time.current = 12 # Jump forward 12 seconds
server.try_run()
assert out_queue.length() == 0 # still nothing because it was cleared
finally:
out_queue.delete()
def add_workflow(**kwargs):
"""
Add a workflow to the system
Variables:
None
Arguments:
None
Data Block:
{
"name": "Workflow name", # Name of the workflow
"classification": "", # Max classification for workflow
"label": ['label1'], # Labels for the workflow
"priority": "LOW", # Priority of the workflow
"status": "MALICIOUS", # Status of the workflow
"query": "*:*" # Query to match the data
}
Result example:
{
"success": true # Saving the user info succeded
}
"""
data = request.json
name = data.get('name', None)
query = data.get('query', None)
if not name:
return make_api_response({"success": False}, err="Name field is required", status_code=400)
if not query:
return make_api_response({"success": False}, err="Query field is required", status_code=400)
if not verify_query(query):
return make_api_response({"success": False}, err="Query contains an error", status_code=400)
data.update({
"workflow_id": get_random_id(),
"creator": kwargs['user']['uname'],
"edited_by": kwargs['user']['uname'],
"priority": data['priority'] or None,
"status": data['status'] or None
})
try:
workflow_data = Workflow(data)
except ValueError as e:
return make_api_response({'success': False}, err=str(e), status_code=400)
return make_api_response({"success": STORAGE.workflow.save(workflow_data.workflow_id, workflow_data),
"workflow_id": workflow_data.workflow_id})
def datastore(datastore_connection):
ds = datastore_connection
try:
create_users(ds)
signatures.extend(create_signatures(ds))
ds.signature.commit()
for _ in range(TEST_SIZE):
f = random_model_obj(File)
ds.file.save(f.sha256, f)
file_list.append(f.sha256)
ds.file.commit()
for x in range(TEST_SIZE):
a = random_model_obj(Alert)
a.file.sha256 = file_list[x]
ds.alert.save(a.alert_id, a)
ds.alert.commit()
for x in range(TEST_SIZE):
r = random_model_obj(Result)
r.sha256 = file_list[x]
ds.result.save(r.build_key(), r)
ds.result.commit()
for x in range(TEST_SIZE):
s = random_model_obj(Submission)
for f in s.files:
f.sha256 = file_list[x]
ds.submission.save(s.sid, s)
ds.submission.commit()
for x in range(TEST_SIZE):
h = random_model_obj(Heuristic)
h.heur_id = f"AL_TEST_{x}"
ds.heuristic.save(h.heur_id, h)
ds.heuristic.commit()
for _ in range(TEST_SIZE):
w_id = get_random_id()
w = random_model_obj(Workflow)
ds.workflow.save(w_id, w)
ds.workflow.commit()
yield ds
finally:
ds.alert.wipe()
ds.file.wipe()
ds.result.wipe()
ds.signature.wipe()
ds.submission.wipe()
ds.heuristic.wipe()
ds.workflow.wipe()
wipe_users(ds)
def upload(self, src_path, dst_path):
dst_path = self.normalize(dst_path)
dirname = posixpath.dirname(dst_path)
filename = posixpath.basename(dst_path)
tempname = get_random_id()
temppath = posixpath.join(dirname, tempname)
finalpath = posixpath.join(dirname, filename)
assert (finalpath == dst_path)
self.makedirs(dirname)
self.sftp.put(src_path, temppath)
self.sftp.rename(temppath, finalpath)
assert (self.exists(dst_path))
def test_submit_path(datastore, client):
content = get_random_phrase(wmin=15, wmax=50).encode()
test_path = "/tmp/test_submit_{}.txt".format(get_random_id())
with open(test_path, 'wb') as test_file:
test_file.write(content + b"PATH")
params = {'service_spec': {"extract": {"password": "******"}}}
res = client.submit(path=test_path, params=params)
assert res.get('sid', None) is not None
assert res['files'][0]['name'] == os.path.basename(test_path)
assert res['params']['service_spec'] == params['service_spec']
assert res == datastore.submission.get(res['sid'], as_obj=False)
def reply_queue_name(prefix=None, suffix=None):
if prefix:
components = [prefix]
else:
components = []
components.append(get_random_id())
if suffix:
components.append(str(suffix))
return '-'.join(components)
def client(datastore):
user = datastore.user.get('admin')
random_pass = get_random_password(length=48)
key_name = "key_%s" % get_random_id().lower()
user.apikeys[key_name] = {
"password": bcrypt.hash(random_pass),
"acl": ["R", "W"]
}
datastore.user.save('admin', user)
api_key = "%s:%s" % (key_name, random_pass)
c = get_client(UI_HOST, apikey=('admin', api_key), verify=False)
return c
def test_get_message(datastore, client):
notification_queue = get_random_id()
queue = NamedQueue("nq-%s" % notification_queue,
host=config.core.redis.persistent.host,
port=config.core.redis.persistent.port)
queue.delete()
msg = random_model_obj(Submission).as_primitives()
queue.push(msg)
res = client.ingest.get_message(notification_queue)
assert isinstance(res, dict)
assert 'sid' in res
assert 'results' in res
assert res == msg
def upload(self, src_path, dst_path):
dst_path = self.normalize(dst_path)
if src_path == dst_path:
return
dirname = os.path.dirname(dst_path)
filename = os.path.basename(dst_path)
tempname = get_random_id()
temppath = _join(dirname, tempname)
finalpath = _join(dirname, filename)
assert (finalpath == dst_path)
self.makedirs(dirname)
shutil.copy(src_path, temppath)
shutil.move(temppath, finalpath)
assert (self.exists(dst_path))
def upload(self, src_path, dst_path):
dst_path = self.normalize(dst_path)
dirname = posixpath.dirname(dst_path)
filename = posixpath.basename(dst_path)
tempname = get_random_id()
temppath = posixpath.join(dirname, tempname)
finalpath = posixpath.join(dirname, filename)
assert (finalpath == dst_path)
self.makedirs(dirname)
with open(src_path, 'rb') as localfile:
self.log.debug("Storing: %s", temppath)
self.ftp.storbinary('STOR ' + temppath, localfile)
self.log.debug("Rename: %s -> %s", temppath, finalpath)
self.ftp.rename(temppath, finalpath)
assert (self.exists(dst_path))
def test_ingest_path(datastore, client):
content = get_random_phrase(wmin=15, wmax=50).encode()
test_path = "/tmp/test_ingest_{}".format(get_random_id())
with open(test_path, 'wb') as test_file:
test_file.write(content + b"PATH")
res = client.ingest(
alert=True,
path=test_path,
params={'service_spec': {
"extract": {
"password": "******"
}
}})
assert res.get('ingest_id', None) is not None
def test_get_message_list(datastore, client):
notification_queue = get_random_id()
queue = NamedQueue("nq-%s" % notification_queue,
host=config.core.redis.persistent.host,
port=config.core.redis.persistent.port)
queue.delete()
msg_0 = random_model_obj(Submission).as_primitives()
queue.push(msg_0)
msg_1 = random_model_obj(Submission).as_primitives()
queue.push(msg_1)
res = client.ingest.get_message_list(notification_queue)
assert len(res) == 2
assert res[0] == msg_0
assert res[1] == msg_1
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_download_file_handle(datastore, client):
signature_id = random_id_from_collection(datastore,
'signature',
q="type:yara")
query = "id:{}".format(signature_id)
output = "/tmp/sigs_{}_obj".format(get_random_id())
res = client.signature.download(query=query, output=open(output, 'wb'))
assert res
found = False
with open(output, 'rb') as fh:
if b"yara/sample_rules.yar" in fh.read():
found = True
if not found:
pytest.fail("This is not the signature file that we were expecting.")
def test_uid():
test_data = "test" * 1000
rid = get_random_id()
id_test = get_id_from_data(test_data)
id_test_l = get_id_from_data(test_data, length=LONG)
id_test_m = get_id_from_data(test_data, length=MEDIUM)
id_test_s = get_id_from_data(test_data, length=SHORT)
id_test_t = get_id_from_data(test_data, length=TINY)
assert 23 > len(rid) >= 20
assert 23 > len(id_test) >= 20
assert 44 > len(id_test_l) >= 41
assert 23 > len(id_test_m) >= 20
assert 13 > len(id_test_s) >= 10
assert 8 > len(id_test_t) >= 5
assert id_test == id_test_m
for c_id in [rid, id_test, id_test_l, id_test_m, id_test_s, id_test_t]:
for x in c_id:
assert x in BASE62_ALPHABET
def test_single(clean_redis):
disp = dispatch_hash.DispatchHash('test-disptach-hash', clean_redis)
try:
file_hash = get_random_id()
service = 'service_name'
result_key = 'some-result'
# An empty dispatch hash isn't finished
assert not disp.all_finished()
# If we call dispatch, the time should be set
now = time.time()
disp.dispatch(file_hash, service)
assert abs(now - disp.dispatch_time(file_hash, service)) < 1
assert not disp.finished(file_hash, service)
assert disp.dispatch_count() == 1
assert not disp.all_finished()
# After failing, the time should be reset
disp.fail_recoverable(file_hash, service)
assert disp.dispatch_time(file_hash, service) == 0
assert disp.dispatch_count() == 0
assert not disp.finished(file_hash, service)
assert not disp.all_finished()
# Try dispatching again
now = time.time()
disp.dispatch(file_hash, service)
assert abs(now - disp.dispatch_time(file_hash, service)) < 1
assert not disp.finished(file_hash, service)
assert disp.dispatch_count() == 1
assert not disp.all_finished()
# Success rather than failure
disp.finish(file_hash, service, result_key, 0, "U")
assert disp.dispatch_time(file_hash, service) == 0
assert disp.dispatch_count() == 0
assert disp.finished_count() == 1
assert disp.all_finished()
assert disp.finished(file_hash, service) == dispatch_hash.DispatchRow('result', result_key, 0, False, 'U')
assert disp.all_finished()
finally:
disp.delete()
