def test_dispatch_extracted(clean_redis, clean_datastore):
redis = clean_redis
ds = clean_datastore
# def service_queue(name): return get_service_queue(name, redis)
# Setup the fake datastore
file_hash = get_random_hash(64)
second_file_hash = get_random_hash(64)
for fh in [file_hash, second_file_hash]:
obj = random_model_obj(models.file.File)
obj.sha256 = fh
ds.file.save(fh, obj)
# Inject the fake submission
submission = random_model_obj(models.submission.Submission)
submission.files = [dict(name='./file', sha256=file_hash)]
sid = submission.sid = 'first-submission'
disp = Dispatcher(ds, redis, redis)
disp.running = ToggleTrue()
client = DispatchClient(ds, redis, redis)
client.dispatcher_data_age = time.time()
client.dispatcher_data.append(disp.instance_id)
# Launch the submission
client.dispatch_submission(submission)
disp.pull_submissions()
disp.service_worker(disp.process_queue_index(sid))
# Finish one service extracting a file
job = client.request_work('0', 'extract', '0')
assert job.fileinfo.sha256 == file_hash
assert job.filename == './file'
new_result: Result = random_minimal_obj(Result)
new_result.sha256 = file_hash
new_result.response.service_name = 'extract'
new_result.response.extracted = [
dict(sha256=second_file_hash,
name='second-*',
description='abc',
classification='U')
]
client.service_finished(sid, 'extracted-done', new_result)
# process the result
disp.pull_service_results()
disp.service_worker(disp.process_queue_index(sid))
disp.service_worker(disp.process_queue_index(sid))
#
job = client.request_work('0', 'extract', '0')
assert job.fileinfo.sha256 == second_file_hash
assert job.filename == 'second-*'
def test_simple(clean_redis, clean_datastore):
ds = clean_datastore
redis = clean_redis
def service_queue(name):
return get_service_queue(name, redis)
file = random_model_obj(File)
file_hash = file.sha256
file.type = 'unknown'
ds.file.save(file_hash, file)
sub: Submission = random_model_obj(models.submission.Submission)
sub.sid = sid = 'first-submission'
sub.params.ignore_cache = False
sub.params.max_extracted = 5
sub.params.classification = get_classification().UNRESTRICTED
sub.params.initial_data = json.dumps({'cats': 'big'})
sub.files = [dict(sha256=file_hash, name='file')]
disp = Dispatcher(ds, redis, redis)
disp.running = ToggleTrue()
client = DispatchClient(ds, redis, redis)
client.dispatcher_data_age = time.time()
client.dispatcher_data.append(disp.instance_id)
# Submit a problem, and check that it gets added to the dispatch hash
# and the right service queues
logger.info('==== first dispatch')
# task = SubmissionTask(sub.as_primitives(), 'some-completion-queue')
client.dispatch_submission(sub)
disp.pull_submissions()
disp.service_worker(disp.process_queue_index(sid))
task = disp.tasks.get(sid)
assert task.queue_keys[(file_hash, 'extract')] is not None
assert task.queue_keys[(file_hash, 'wrench')] is not None
assert service_queue('extract').length() == 1
assert service_queue('wrench').length() == 1
# Making the same call again will queue it up again
logger.info('==== second dispatch')
disp.dispatch_file(task, file_hash)
assert task.queue_keys[(file_hash, 'extract')] is not None
assert task.queue_keys[(file_hash, 'wrench')] is not None
assert service_queue('extract').length() == 1 # the queue doesn't pile up
assert service_queue('wrench').length() == 1
logger.info('==== third dispatch')
job = client.request_work('0', 'extract', '0')
assert job.temporary_submission_data == [{'name': 'cats', 'value': 'big'}]
client.service_failed(sid, 'abc123', make_error(file_hash, 'extract'))
# Deliberately do in the wrong order to make sure that works
disp.pull_service_results()
disp.service_worker(disp.process_queue_index(sid))
assert task.queue_keys[(file_hash, 'extract')] is not None
assert task.queue_keys[(file_hash, 'wrench')] is not None
assert service_queue('extract').length() == 1
# Mark extract as finished, wrench as failed
logger.info('==== fourth dispatch')
client.request_work('0', 'extract', '0')
client.request_work('0', 'wrench', '0')
client.service_finished(sid, 'extract-result',
make_result(file_hash, 'extract'))
client.service_failed(sid, 'wrench-error',
make_error(file_hash, 'wrench', False))
for _ in range(2):
disp.pull_service_results()
disp.service_worker(disp.process_queue_index(sid))
assert wait_error(task, file_hash, 'wrench')
assert wait_result(task, file_hash, 'extract')
assert service_queue('av-a').length() == 1
assert service_queue('av-b').length() == 1
assert service_queue('frankenstrings').length() == 1
# Have the AVs fail, frankenstrings finishes
logger.info('==== fifth dispatch')
client.request_work('0', 'av-a', '0')
client.request_work('0', 'av-b', '0')
client.request_work('0', 'frankenstrings', '0')
client.service_failed(sid, 'av-a-error',
make_error(file_hash, 'av-a', False))
client.service_failed(sid, 'av-b-error',
make_error(file_hash, 'av-b', False))
client.service_finished(sid, 'f-result',
make_result(file_hash, 'frankenstrings'))
for _ in range(3):
disp.pull_service_results()
disp.service_worker(disp.process_queue_index(sid))
assert wait_result(task, file_hash, 'frankenstrings')
assert wait_error(task, file_hash, 'av-a')
assert wait_error(task, file_hash, 'av-b')
assert service_queue('xerox').length() == 1
# Finish the xerox service and check if the submission completion got checked
logger.info('==== sixth dispatch')
client.request_work('0', 'xerox', '0')
client.service_finished(sid, 'xerox-result-key',
make_result(file_hash, 'xerox'))
disp.pull_service_results()
disp.service_worker(disp.process_queue_index(sid))
disp.save_submission()
assert wait_result(task, file_hash, 'xerox')
assert disp.tasks.get(sid) is None