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_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_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_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 export(request, task_id): if request.method == "POST": taken_dirs = request.POST.getlist("dirs") taken_files = request.POST.getlist("files") if not taken_dirs and not taken_files: return view_error( request, "Please select at least one directory or file" " to be exported." ) zip = Task.create_zip( task_id=task_id, taken_dirs=taken_dirs, taken_files=taken_files ) if not zip: return view_error(request, "Failed to create zip.") response = HttpResponse( zip.getvalue(), content_type="application/zip" ) response["Content-Disposition"] = ( "attachment; filename=%s.zip" % task_id ) return response report = AnalysisController.get_report(task_id) dirs, files = Task.get_files(task_id) return render_template( request, "analysis/export.html", report=report, dirs=dirs, files=files )
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_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_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_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_reschedule_id(self): id, sample = self.add_task() task = Task() newid = task.reschedule(task_id=id) 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"
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_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_process_task_range_multi_db(self, mt, mp): task1, task2 = Task(), Task() id1, id2 = task1.add_path(__file__), task2.add_path(__file__) process_task_range("%s,%s" % (id1, id2)) assert mp.call_count == 2 mt.assert_has_calls([ mock.call(), mock.call().set_task(mock.ANY), mock.call(), mock.call().set_task(mock.ANY) ])
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_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_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_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_set_task(self): id, sample = self.add_task() db_task = self.db.view_task(id) task = Task() task.set_task(db_task) assert task.id == id assert task.category == "file" assert task.path == cwd(analysis=id) assert task.db_task == db_task assert task.target == sample assert len(task.targets) == 1 assert isinstance(task.targets[0], Target)
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 test_add_service(self): task = Task() id = task.add_service(timeout=60, tags=["officepc"], owner="Doge") task_path = cwd(analysis=id) db_task = self.db.view_task(id) assert id is not None assert os.path.exists(task_path) assert db_task.type == "service" assert db_task.owner == "Doge" assert db_task.timeout == 60 assert db_task.priority == 999 assert db_task.tags[0].name == "officepc" assert db_task.targets == []
def test_load_task_from_dict(self): task_dict = { "id": 42, "category": "file", "target": "/tmp/stuff/doge42.exe", } task = Task() task.load_task_dict(task_dict) assert task.id == 42 assert task.category == "file" assert task.target == "/tmp/stuff/doge42.exe" assert task.path == cwd(analysis=42) assert task.type == "regular"
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 test_add_baseline(self): task = Task() id = task.add_baseline(timeout=60, owner="Doge", machine="machine1") task_path = cwd(analysis=id) db_task = self.db.view_task(id) assert id is not None assert os.path.exists(task_path) assert db_task.type == "baseline" assert db_task.owner == "Doge" assert db_task.timeout == 60 assert db_task.priority == 999 assert db_task.machine == "machine1" assert db_task.memory == False assert db_task.targets == []
def process_tasks(instance, maxcount, timeout): count = 0 endtime = 0 db = Database() if timeout: endtime = int(time.time() + timeout) try: while process_check_stop(count, maxcount, endtime): task_id = db.processing_get_task(instance) # Wait a small while before trying to fetch a new task. if task_id is None: time.sleep(1) continue task = Task(db.view_task(task_id)) log.info("Task #%d: reporting task", task.id) process_task(task) count += 1 except Exception as e: log.exception("Caught unknown exception: %s", e)
def get_analysis_manager(self, db_task, machine): """Searches all available analysis managers for one that supports the type of the given task. Returns an analysis manager. Returns None if no manager supports the type""" managers = cuckoo.analysis.plugins analysis_manager = None for manager in managers: if db_task.type in manager.supports: core_task = Task(db_task) analysis_manager = manager( machine, self.machinery, self.machine_lock ) try: analysis_manager.set_task(core_task) analysis_manager.set_target(core_task.targets) except Exception as e: analysis_manager = None log.exception( "Failure when setting task and target for analysis" " manager '%s'.", manager ) break return analysis_manager
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_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_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)