def test_set_api_key(): raw_api = VxCubeRawApi() raw_api.api_key = "test_key" assert raw_api.api_key == "test_key" assert "Authorization" in raw_api._api_request.headers assert raw_api._api_request.headers["Authorization"] == "api-key test_key"
def test_task_phone_actions_iter(): request = mock.Mock(return_value={ "total_count": 2, "items": ["phone_action23", "phone_action42"] }) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) task = Task(id=42, _raw_api=raw_api) phone_actions = list( task.phone_actions_iter(count_per_request=23, search="lost")) assert len(phone_actions) == 2 assert phone_actions[0] == "phone_action23" assert phone_actions[1] == "phone_action42" request.assert_called_with( method="get", url="http://test/api-2.0/tasks/42/phone_actions", params={}, headers={}, json={ "count": 23, "offset": 0, "search": "lost" })
def test_raw_api_empty(): raw_api = VxCubeRawApi(base_url="http://test", version=23.42) assert isinstance(raw_api._api_request, VxCubeApiRequest) assert raw_api._version == 23.42 assert raw_api._base_url == "http://test/api-23.42/" assert raw_api._api_key is None assert raw_api.api_key is None
def test_task_restart_processing_task(): request = mock.Mock(return_value=None) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) task = Task(id=42, status="processing", _raw_api=raw_api) assert not task.restart() request.assert_not_called()
def test_task_android_cureit(): request = mock.Mock(return_value={"status": "processing", "retries": None}) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) task = Task(id=42, platform_code="android4.3", _raw_api=raw_api) cureit = task.cureit() assert cureit is None request.assert_not_called()
def test_analysis_delete(): request = mock.Mock(return_value=None) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) analysis = Analysis(id=42, _raw_api=raw_api) assert analysis.delete() request.assert_called_with(method="delete", url="http://test/api-2.0/analyses/42", params={}, headers={})
def test_task_download_report(): request = mock.Mock(return_value=None) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) task = Task(id=42, _raw_api=raw_api) task.download_report(output_file="test_output") request.assert_called_with(method="get", url="http://test/api-2.0/tasks/42/report", params={}, headers={}, output_file="test_output")
def test_subscribe_progress_with_connection_close(): values = dict( id=1, _raw_api=VxCubeRawApi(base_url="https://test", version=2.0), tasks=[dict(id=1, status="processing")], ) analysis = Analysis(**values) ws = mock.MagicMock() ws.send.side_effect = WebSocketConnectionClosedException() with mock.patch("websocket.WebSocket", return_value=ws), mock.patch("tortilla.Wrap.get", return_value={"id": 42}): list(analysis.subscribe_progress())
def test_analysis_download_sample(): request = mock.Mock(return_value=None) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) analysis = Analysis(id=42, _raw_api=raw_api) analysis.download_sample(output_file="test_output") request.assert_called_with(method="get", url="http://test/api-2.0/analyses/42/sample", params={}, headers={}, output_file="test_output")
def test_cureit_download_by_task(): request = mock.Mock(return_value=None) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) cureit = CureIt(task_id=42, _raw_api=raw_api) cureit.download(output_file="test_output") request.assert_called_with(method="get", url="http://test/api-2.0/tasks/42/cureit.exe", params={}, headers={}, output_file="test_output")
def test_task_restart(): request = mock.Mock(return_value={"status": "processing"}) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) task = Task(id=42, status="deleted", _raw_api=raw_api) assert task.restart() assert task.status == "processing" request.assert_called_with(method="post", url="http://test/api-2.0/tasks/42/restart", params={}, headers={})
def main(): # Create raw API api = VxCubeRawApi(api_key=API_KEY, version=1.0) # Upload sample with file_wrapper(FILE_PATH) as file: rs = api.create.post(files=file, data={"analyse_time": 5.0, "winxpx86": True}) # Get analysis data analysis_id = rs["url"].split("/")[-1] rs = api.report.get(analysis_id) print(rs) if all(task["status"] != "processing" for task in rs["tasks"]): print("All tasks finished")
def test_task_storage_list(): return_storage = {"files": [], "folders": []} request = mock.Mock(return_value=return_storage) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) task = Task(id=42, _raw_api=raw_api) storage = task.storage_lists() assert storage is return_storage request.assert_called_with( method="get", url="http://test/api-2.0/tasks/42/archive_storage", params={}, headers={})
def test_task_download_from_storage(): return_storage = {"files": [], "folders": []} request = mock.Mock(return_value=return_storage) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) task = Task(id=42, _raw_api=raw_api) task.download_storage_file(path="test_path", output_file="test_output") request.assert_called_with( method="get", url="http://test/api-2.0/tasks/42/archive_storage", params={}, headers={}, json={"path": "test_path"}, output_file="test_output")
def test_cureit_retry_by_task(): request = mock.Mock(return_value={"status": "processing"}) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) cureit = CureIt(task_id=42, status="failed", _raw_api=raw_api) assert cureit.retry() assert cureit.task_id == 42 assert cureit.status == "processing" assert cureit._raw_api is raw_api request.assert_called_with(method="put", url="http://test/api-2.0/tasks/42/cureit", params={}, headers={})
def test_delete_session(): values = dict(api_key="{}-{}-{}-{}-{}".format("a" * 8, "b" * 4, "c" * 4, "d" * 4, "e" * 12), start_date="2018-04-08T15:16:23.420000+00:00") request = mock.Mock(return_value=None) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) session = Session(_raw_api=raw_api, **values) assert session.delete() request.assert_called_with(method="delete", url="http://test/api-2.0/sessions/{}".format( values["api_key"]), params={}, headers={})
def test_task_windows_cureit(): request = mock.Mock(return_value={"status": "processing", "retries": None}) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) task = Task(id=42, platform_code="winxpx86", _raw_api=raw_api) cureit = task.cureit() assert cureit.task_id == 42 assert cureit.status == "processing" assert cureit.retries is None assert cureit._raw_api is raw_api request.assert_called_with(method="get", url="http://test/api-2.0/tasks/42/cureit", params={}, headers={})
def test_analysis_cureit(): request = mock.Mock(return_value={"status": "processing", "retries": None}) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) analysis = Analysis(id=42, _raw_api=raw_api) cureit = analysis.cureit() assert cureit.analysis_id == 42 assert cureit.status == "processing" assert cureit.retries is None assert cureit._raw_api is raw_api request.assert_called_with(method="get", url="http://test/api-2.0/analyses/42/cureit", params={}, headers={})
def test_task_intents(): request = mock.Mock(return_value=["intent23", "intent42"]) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) task = Task(id=42, _raw_api=raw_api) intents = task.intents(23, 42, "lost") assert len(intents) == 2 assert intents[0] == "intent23" assert intents[1] == "intent42" request.assert_called_with(method="get", url="http://test/api-2.0/tasks/42/intents", params={}, headers={}, json={ "count": 23, "offset": 42, "search": "lost" })
def test_subscribe_progress_with_https_scheme(): values = dict( id=1, _raw_api=VxCubeRawApi(base_url="https://test", version=2.0), tasks=[ dict(id=1, status="processing"), dict(id=2, status="processing") ], ) analysis = Analysis(**values) ws = mock.MagicMock() ws.__iter__.return_value = iter(["{\"message\": \"test\"}"]) with mock.patch("websocket.WebSocket", return_value=ws): with mock.patch("tortilla.Wrap.get", return_value={"id": 42}): assert list(analysis.subscribe_progress()) == [{"message": "test"}] assert analysis.id == 42 ws.connect.assert_called_with("wss://test/api-2.0/ws/progress", header={"Authorization": "api-key None"}) ws.send.assert_called_with("{\"analysis_id\": 1}") ws.close.assert_called_once()
def test_analysis_restart(): request = mock.Mock( return_value={"tasks": [{ "id": 23, "status": "processing" }]}) with mock.patch("vxcube_api.raw_api.VxCubeApiRequest.request", new=request): raw_api = VxCubeRawApi(base_url="http://test", version=2.0) analysis = Analysis(id=42, tasks=[{ "id": 23, "status": "failed" }], _raw_api=raw_api) assert analysis.restart() assert analysis.tasks[0].id == 23 assert analysis.tasks[0].status == "processing" request.assert_called_with(method="post", url="http://test/api-2.0/analyses/42/restart", params={}, headers={})
def __init__(self, api_key=None, base_url="https://vxcube.drweb.com/", version=2.0): self._raw_api = VxCubeRawApi(api_key, base_url, version)
def test_raw_api_deprecation_warning(): with pytest.deprecated_call(): VxCubeRawApi(base_url="http://test", version=1.0)
class VxCubeApi(object): def __init__(self, api_key=None, base_url="https://vxcube.drweb.com/", version=2.0): self._raw_api = VxCubeRawApi(api_key, base_url, version) def login(self, login, password, new_key=False): """ Get API key using login and password. :param str login: :param str password: :return None: :raises VxCubeApiException, VxCubeApiHttpException """ logger.debug("Login with %s", login) if self._raw_api.api_key: logger.info("Use login with existing API key") data = filter_data(login=login, password=password, new_key=new_key) response = self._raw_api.login.post(json=data) if not isinstance(response, dict) or "api_key" not in response: logger.error("Unknown server response") raise VxCubeApiException("Incorrect server response") self._raw_api.api_key = response["api_key"] @return_objects(Session, add_raw_api=True) def sessions(self): """ Get a list of open sessions. :return list[Session]: sessions :raises VxCubeApiHttpException """ logger.debug("Get sessions") return self._raw_api.sessions.get() @return_objects(Format) def formats(self): """ Get a list of supported formats. :return list[Format]: formats :raises VxCubeApiHttpException """ logger.debug("Get formats") return self._raw_api.formats.get() @return_objects(Platform) def platforms(self): """ Get a list of supported platforms. :return list[Platform]: platforms :raises VxCubeApiHttpException """ logger.debug("Get platforms") return self._raw_api.platforms.get() @return_objects(License) def license(self): """ Get information about current license. :return License: license :raises VxCubeApiHttpException """ logger.debug("Get license") return self._raw_api.license.get() @return_objects(Sample, add_raw_api=True) def samples(self, sample_id=None, count=None, offset=None, md5=None, sha1=None, sha256=None, format_name=None, format_group_name=None): """ Get sample by ID or get a filtered list of samples. :param int sample_id: :param int count: :param int offset: :param str md5: :param str sha1: :param str sha256: :param str format_name: :param str format_group_name: :return List[Sample] or Sample: :raises VxCubeApiHttpException """ if sample_id: logger.debug("Get sample") return self._raw_api.samples.get(sample_id) logger.debug("Get list of samples") data = filter_data(count=count, offset=offset, md5=md5, sha1=sha1, sha256=sha256, format_name=format_name, format_group_name=format_group_name) return self._raw_api.samples.get(json=data) def samples_iter(self, count_per_request=100, **kwargs): """ Iterator over self.samples. :param count_per_request: :param kwargs: :return: """ logger.debug("Use sample iterator") kwargs.pop("sample_id", None) return iterator(func=self.samples, count_per_request=count_per_request, item_key=None, **kwargs) @return_objects(Sample, add_raw_api=True) def upload_sample(self, file): """ Upload sample to Dr.Web vxCube server. :param str or file-like object file: path or file-like object :return Sample: :raises VxCubeApiHttpException """ logger.debug("Upload sample to server") with file_wrapper(file) as file: fields = {"file": (file.name, file, "application/octet-stream")} enc = MultipartEncoder(fields=fields) headers = {"Content-Type": enc.content_type} res = self._raw_api.samples.post(data=enc, headers=headers) if "samples" in res: res = res["samples"] return res @return_objects(Analysis, add_raw_api=True) def analyses(self, analysis_id=None, count=None, offset=None, format_group_name=None): """ Get analysis by ID or get a filtered list of analyses. :param int analysis_id: :param int count: :param int offset: :param str format_group_name: :return Analysis or list[Analysis]: :raises VxCubeApiHttpException """ if analysis_id: logger.debug("Get analysis") return self._raw_api.analyses.get(analysis_id) logger.debug("Get analysis list") data = filter_data(count=count, offset=offset, format_group_name=format_group_name) return self._raw_api.analyses.get(json=data) def analyses_iter(self, count_per_request=100, **kwargs): """ Iterator over self.analyses. :param count_per_request: :param kwargs: :return: """ logger.debug("Use analysis iterator") kwargs.pop("analysis_id", None) return iterator(func=self.analyses, count_per_request=count_per_request, item_key=None, **kwargs) @return_objects(Analysis, add_raw_api=True) def start_analysis(self, sample_id, platforms, analysis_time=None, format_name=None, custom_cmd=None, generate_cureit=None, drop_size_limit=None, net=None, copylog=False, crypto_api_limit=64, dump_size_limit=64, flex_time=False, forwards=None, get_lib=False, injects_limit=100, monkey_clicker=None, dump_browsers=True, dump_mapped=True, dump_ssdt=True, no_clean=False, optional_count=None, proc_lifetime=None, set_date=None, userbatch=None, dump_processes=True, write_file_limit=512): """ Start sample analysis. :param int sample_id: :param list platforms: :param int analysis_time: :param str format_name: :param str custom_cmd: :param bool generate_cureit: :param int drop_size_limit: :param str net: :param bool copylog: :param int crypto_api_limit: :param int dump_size_limit: :param bool flex_time: :param list forwards: :param bool get_lib: :param int injects_limit: :param bool monkey_clicker: :param bool dump_browsers: :param bool dump_mapped: :param bool dump_ssdt: :param bool no_clean: :param int optional_count: :param int proc_lifetime: :param str set_date: :param str userbatch: :param bool dump_processes: :param int write_file_limit: :return Analyse: :raises VxCubeApiHttpException """ logger.debug("Start analysis") data = filter_data(sample_id=sample_id, platforms=platforms, analysis_time=analysis_time, format_name=format_name, custom_cmd=custom_cmd, generate_cureit=generate_cureit, drop_size_limit=drop_size_limit, net=net, copylog=copylog, crypto_api_limit=crypto_api_limit, dump_size_limit=dump_size_limit, flex_time=flex_time, forwards=forwards, get_lib=get_lib, injects_limit=injects_limit, monkey_clicker=monkey_clicker, dump_browsers=dump_browsers, dump_mapped=dump_mapped, dump_ssdt=dump_ssdt, dump_processes=dump_processes, no_clean=no_clean, optional_count=optional_count, proc_lifetime=proc_lifetime, set_date=set_date, write_file_limit=write_file_limit) if userbatch: with file_wrapper(userbatch) as file: return self._raw_api.analyses.post(files={ "userbatch": (file.name, file, "application/octet-stream") }, data=data) return self._raw_api.analyses.post(json=data) @return_objects(Task, add_raw_api=True) def task(self, task_id): """ Get task data. :param int task_id: :return Task: """ logger.debug("Get task") return self._raw_api.tasks(task_id).get()