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()
