def test_start_and_wait_noagent(self, mwf, mrsa, msas, mrs):
self.create_cwd()
task = Task()
id = task.add_service(owner="1", tags="service,mitm", timeout=120)
task.load_from_db(id)
# Mock resultserver obj so we can check if add_task was called
resulserver_obj = mock.MagicMock()
mrs.return_value = resulserver_obj
manager = self.get_manager(task)
manager.machine.options = "noagent"
manager.init(self.db)
manager.machinery = mock.MagicMock()
manager.route = mock.MagicMock()
manager.aux = mock.MagicMock()
result = manager.start_and_wait()
# Check if all required methods were called successfully
msas.assert_has_calls([mock.call("starting"), mock.call("running")])
resulserver_obj.add_task.assert_called_once_with(
task.db_task, manager.machine)
manager.aux.start.assert_called_once()
manager.machinery.start.assert_called_once_with(
"machine1", task.db_task)
manager.route.route_network.assert_called_once()
manager.machine_lock.release.assert_called_once()
mrsa.assert_called_once_with("starting")
mwf.assert_called_once()
assert result
def test_start_and_wait_url(self, mrsa, msas, mrs):
self.create_cwd()
task = Task()
id = task.add_url("http://example.com/42")
task.load_from_db(id)
# Mock resultserver obj so we can check if add_task was called
resulserver_obj = mock.MagicMock()
mrs.return_value = resulserver_obj
manager = self.get_manager(task)
manager.init(self.db)
manager.machinery = mock.MagicMock()
manager.route = mock.MagicMock()
manager.aux = mock.MagicMock()
manager.guest_manager = mock.MagicMock()
# Set status manually, because the method used is mocked
manager.analysis.status = "starting"
result = manager.start_and_wait()
# Check if all required methods were called successfully
msas.assert_has_calls([mock.call("starting"), mock.call("running")])
resulserver_obj.add_task.assert_called_once_with(
task.db_task, manager.machine)
manager.aux.start.assert_called_once()
manager.machinery.start.assert_called_once_with(
"machine1", task.db_task)
manager.route.route_network.assert_called_once()
manager.machine_lock.release.assert_called_once()
mrsa.assert_called_once_with("starting")
manager.guest_manager.start_analysis.assert_called_once()
manager.guest_manager.wait_for_completion.assert_called_once()
assert result
def test_stop_and_wait_dump_mem(self, msas, mrs):
self.create_cwd()
task = Task()
id = task.add_path(__file__, memory=True)
task.load_from_db(id)
# Mock resultserver obj so we can check if del_task was called
resulserver_obj = mock.MagicMock()
mrs.return_value = resulserver_obj
manager = self.get_manager(task)
manager.init(self.db)
manager.machinery = mock.MagicMock()
manager.route = mock.MagicMock()
manager.aux = mock.MagicMock()
manager.stop_and_wait()
# Check if all required methods were called successfully
msas.assert_called_once_with("stopping")
manager.aux.stop.assert_called_once()
manager.machinery.dump_memory.assert_called_once_with(
"machine1", cwd("storage", "analyses", str(task.id), "memory.dmp"))
manager.machinery.stop.assert_called_once_with("machine1")
resulserver_obj.del_task.assert_called_once_with(
task.db_task, manager.machine)
manager.route.unroute_network.assert_called_once()
def test_build_options(self):
task = Task()
id = task.add_path(__file__, options={"free": "yes"})
task.load_from_db(id)
a = abstracts.AnalysisManager(
FakeMachine(), mock.MagicMock(), mock.MagicMock()
)
a.set_task(task)
expected = {
"clock": task.clock,
"enforce_timeout": False,
"id": task.id,
"package": "",
"target": None,
"terminate_processes": False,
"ip": "192.168.56.1",
"port": 4242,
"timeout": 120,
"options": "free=yes"
}
assert a.options == {}
a.build_options()
assert a.options == expected
a.build_options({
"file_name": "doge.py",
"options": {"doges": "many"},
"category": "file"
})
assert a.options["options"] == "doges=many,free=yes"
assert a.options["file_name"] == "doge.py"
assert a.options["category"] == "file"
def test_add_reboot_binary_removed(self):
id, sample = self.add_task()
task = Task()
task.load_from_db(id)
task.create_empty()
os.remove(task.targets[0].copied_binary)
newid = task.add_reboot(id)
assert newid is None
def test_set_target(self):
self.create_cwd()
task = Task()
id = task.add_path(__file__)
task.load_from_db(id)
manager = self.get_manager()
manager.set_target(task.targets)
assert manager.target == task.targets[0]
def test_set_status(self):
id, sample = self.add_task()
task = Task()
task.load_from_db(id)
task.set_status("reported")
assert task.status == "reported"
assert task["status"] == "reported"
def test_dir_exists(self):
id, sample = self.add_task()
task = Task()
task.load_from_db(id)
assert not task.dir_exists()
os.mkdir(cwd(analysis=id))
assert task.dir_exists()
def test_set_target_empty(self):
self.create_cwd()
task = Task()
id = task.add_path(__file__)
task.load_from_db(id)
task.task_dict["targets"] = []
manager = self.get_manager()
manager.set_target(task.targets)
assert isinstance(manager.target, Target)
def test_set_target(self):
task = Task()
id = task.add_path(__file__)
task.load_from_db(id)
a = abstracts.AnalysisManager(
FakeMachine(), mock.MagicMock(), mock.MagicMock()
)
with pytest.raises(NotImplementedError):
a.set_target(task.targets)
def test_get_attribute(self):
id, sample = self.add_task()
task = Task()
task.load_from_db(id)
path = cwd(analysis=id)
assert task.id == id
assert task.path == path
assert task.category == "file"
assert task.target == sample
def test_get_item(self):
id, sample = self.add_task()
task = Task()
task.load_from_db(id)
assert task["id"] == id
assert task["category"] == "file"
assert task["target"] == sample
assert task["machine"] is None
assert len(task["targets"]) == 1
def test_set_latest(self):
id, sample = self.add_task()
task = Task()
task.load_from_db(id)
task.create_dirs()
sym_latest = cwd("storage", "analyses", "latest")
task.set_latest()
assert os.path.realpath(sym_latest) == task.path
def test_wait_finish(self, mts):
task = Task()
id = task.add_url("http://example.com/42")
task.load_from_db(id)
a = abstracts.AnalysisManager(
FakeMachine(), mock.MagicMock(), mock.MagicMock()
)
a.set_task(task)
a.analysis.status = "stoppped"
a.wait_finish()
mts.assert_not_called()
def test_refresh(self):
id, sample = self.add_task()
task = Task()
task.load_from_db(id)
self.db.set_machine(id, "machine1")
assert task.machine is None
assert task["machine"] is None
task.refresh()
assert task.machine == "machine1"
assert task["machine"] == "machine1"
def test_process(self, mp, ms, mr):
id, sample = self.add_task()
task = Task()
task.load_from_db(id)
mp.return_value = {"x":"x"}
task.process()
mp.assert_called_once()
ms.assert_called_once()
mr.assert_called_once()
def test_add_massurl(self):
urls = ["http://example%s.com" % n for n in range(500)]
id = submit_task.add_massurl(urls)
task = Task()
task.load_from_db(id)
assert id is not None
assert os.path.exists(cwd(analysis=id))
assert task.path == cwd(analysis=id)
assert len(task.targets) == 500
assert task.type == "massurl"
def test_set_task(self):
self.create_cwd()
task = Task()
id = task.add_path(__file__)
task.load_from_db(id)
manager = self.get_manager()
manager.set_task(task)
assert manager.task == task
assert manager.analysis is not None
assert manager.name == "task_%s_Regular" % task.id
def test_requirement_str(self):
id, sample = self.add_task(
tags=["doge"], platform="DogeOS", machine="Doge1"
)
id = submit_task.add_path(
self.get_file(), tags=["doge"], platform="DogeOS", machine="Doge1"
)
task = Task()
task.load_from_db(id)
req_str = task.requirements_str(task.db_task)
assert req_str == "machine=Doge1 platform=DogeOS tags=doge, "
def test_is_reported(self):
id, sample = self.add_task()
task = Task()
task.load_from_db(id)
task.create_dirs()
assert not task.is_reported()
reports = os.path.join(task.path, "reports")
os.mkdir(reports)
with open(os.path.join(reports, "report.json"),
"wb") as fw:
fw.write(os.urandom(64))
assert task.is_reported()
def test_set_task(self):
task = Task()
id = task.add_path(__file__)
task.load_from_db(id)
a = abstracts.AnalysisManager(
FakeMachine(), mock.MagicMock(), mock.MagicMock()
)
a.set_task(task)
assert a.task == task
assert isinstance(a.analysis, Analysis)
assert a.name == "task_%s_AnalysisManager" % task.id
def test_init_non_file(self, mb):
self.create_cwd()
task = Task()
id = task.add_url("http://example.com/42")
task.load_from_db(id)
manager = self.get_manager(task)
result = manager.init(self.db)
mb.assert_called_once()
assert result
assert isinstance(manager.guest_manager, GuestManager)
assert isinstance(manager.aux, RunAuxiliary)
assert os.path.isfile(os.path.join(task.path, "task.json"))
def test_reschedule_file(self):
id, sample = self.add_task()
task = Task()
task.load_from_db(id)
newid = task.reschedule(priority=3)
oldtask = self.db.view_task(id)
newtask = self.db.view_task(newid)
assert newid is not None
assert oldtask.status == "recovered"
assert newtask.targets[0].category == "file"
assert newtask.targets[0].target == sample
assert newtask.priority == 3
def get_manager(self, task=None):
if task is None:
task = Task()
fd, fpath = tempfile.mkstemp()
os.write(fd, b"\x00" * 32)
os.close(fd)
newname = os.path.join(os.path.dirname(fpath), "testanalysis.exe")
os.rename(fpath, newname)
id = task.add_path(newname)
task.load_from_db(id)
manager = Regular(FakeMachine(), mock.MagicMock(), mock.MagicMock())
manager.set_task(task)
manager.set_target(task.targets)
return manager
def resubmit(request, task_id):
task = Task()
if not task.load_from_db(task_id):
return view_error(request, "No Task was found with this ID")
if task.category == "url":
# TODO This most certainly needs to be improved.
submit_id = submit_manager.pre("strings", [
task.target,
], submit_manager.translate_options_to(task.options))
else:
if not os.path.exists(task.target):
return view_error(
request, "The file you're trying to resubmit "
"no longer exists. Please resubmit it altogether!"
)
# TODO There's a very good chance this won't work properly for
# analyses of type "archive".
submit_id = submit_manager.pre("files", [{
"name": os.path.basename(task.target),
"data": open(task.target, "rb"),
}], submit_manager.translate_options_to(task.options))
return redirect("submission/pre", submit_id=submit_id)
def test_reschedule_url(self):
id, sample = self.add_task(
url="http://example.com/42", category="url"
)
task = Task()
task.load_from_db(id)
newid = task.reschedule(priority=2)
oldtask = self.db.view_task(id)
newtask = self.db.view_task(newid)
assert newid is not None
assert oldtask.status == "recovered"
assert newtask.targets[0].category == "url"
assert newtask.priority == 2
assert newtask.targets[0].target == "http://example.com/42"
def test_create_dirs(self):
id, sample = self.add_task()
task = Task()
task.load_from_db(id)
dirs = ["shots", "logs", "files", "extracted", "buffer", "memory"]
task_path = cwd(analysis=id)
dir_paths = [cwd(task_path, dir) for dir in dirs]
for path in dir_paths:
assert not os.path.exists(path)
assert task.create_dirs()
assert os.path.exists(task_path)
for path in dir_paths:
assert os.path.exists(path)
def test_init_fail(self):
self.create_cwd()
task = Task()
fd, tmpfile = tempfile.mkstemp()
os.write(fd, os.urandom(64))
os.close(fd)
id = task.add_path(tmpfile)
task.load_from_db(id)
manager = self.get_manager(task)
copy_path = cwd("storage", "binaries", File(tmpfile).get_sha256())
# Remove both binaries to make init fail
os.remove(copy_path)
os.remove(tmpfile)
result = manager.init(self.db)
assert not result
def test_load_from_db(self):
id = self.add_task()[0]
task = Task()
assert task.load_from_db(id)
assert task.id == id
assert task.category == "file"
assert task.path == cwd(analysis=id)
def test_request_scheduler_action(self):
task = Task()
id = task.add_url("http://example.com/42")
task.load_from_db(id)
a = abstracts.AnalysisManager(
FakeMachine(), mock.MagicMock(), mock.MagicMock()
)
a.set_task(task)
a.action_lock = mock.MagicMock()
a.action_lock.locked = mock.MagicMock(return_value=False)
a.request_scheduler_action()
a.action_lock.acquire.assert_has_calls([
mock.call(False),
mock.call(True)
])
a.action_lock.release.assert_called_once()
assert a.override_status is None