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