mirror of https://github.com/MISP/misp-modules
Koen Van Impe
6 years ago
parent
ede0dcdd60
commit
adda9562c0
6 changed files with 758 additions and 2 deletions
@ -1,2 +1,2 @@ |
||||
__all__ = ['asn_history', 'circl_passivedns', 'circl_passivessl', 'countrycode', 'cve', 'dns', |
||||
__all__ = ['vmray_submit','asn_history', 'circl_passivedns', 'circl_passivessl', 'countrycode', 'cve', 'dns', |
||||
'eupi', 'ipasn', 'passivetotal', 'sourcecache', 'virustotal', 'whois', 'shodan', 'reversedns', 'wiki'] |
||||
|
||||
@ -0,0 +1,152 @@ |
||||
#!/usr/bin/python3 |
||||
"""Python client library for VMRay REST API""" |
||||
|
||||
import base64 |
||||
import datetime |
||||
import os.path |
||||
import requests |
||||
#import urlparse |
||||
import urllib.parse |
||||
|
||||
from io import IOBase |
||||
from io import BytesIO |
||||
|
||||
# disable nasty certification warning |
||||
# pylint: disable=no-member |
||||
try: |
||||
requests.packages.urllib3.disable_warnings() |
||||
except AttributeError: |
||||
try: |
||||
import urllib3 |
||||
try: |
||||
urllib3.disable_warnings() |
||||
except AttributeError: |
||||
pass |
||||
except ImportError: |
||||
pass |
||||
|
||||
# pylint: disable= |
||||
|
||||
|
||||
class VMRayRESTAPIError(Exception): |
||||
"""Exception class that is used when API returns an error""" |
||||
|
||||
def __init__(self, *args, **kwargs): |
||||
self.status_code = kwargs.pop("status_code", None) |
||||
Exception.__init__(self, *args, **kwargs) |
||||
|
||||
|
||||
def handle_rest_api_result(result): |
||||
"""Handle result of API request (check for errors)""" |
||||
|
||||
if (result.status_code < 200) or (result.status_code > 299): |
||||
try: |
||||
json_result = result.json() |
||||
except ValueError: |
||||
raise VMRayRESTAPIError("API returned error %u: %s" % (result.status_code, result.text), status_code=result.status_code) |
||||
|
||||
raise VMRayRESTAPIError(json_result.get("error_msg", "Unknown error"), status_code=result.status_code) |
||||
|
||||
|
||||
class VMRayRESTAPI(object): |
||||
"""VMRay REST API class""" |
||||
|
||||
def __init__(self, server, api_key, verify_cert=True): |
||||
# split server URL into components |
||||
url_desc = urllib.parse.urlsplit(server) |
||||
|
||||
# assume HTTPS if no scheme is specified |
||||
if url_desc.scheme == "": |
||||
server = "https://" + server |
||||
|
||||
# save variables |
||||
self.server = server |
||||
self.api_key = api_key |
||||
self.verify_cert = verify_cert |
||||
|
||||
def call(self, http_method, api_path, params=None, raw_data=False): |
||||
"""Call VMRay REST API""" |
||||
|
||||
# get function of requests package |
||||
requests_func = getattr(requests, http_method.lower()) |
||||
|
||||
# parse parameters |
||||
req_params = {} |
||||
file_params = {} |
||||
|
||||
if params is not None: |
||||
for key, value in params.items(): |
||||
if isinstance(value, (datetime.date, |
||||
datetime.datetime, |
||||
float, |
||||
int)): |
||||
req_params[key] = str(value) |
||||
elif isinstance(value, str): |
||||
req_params[key] = str(value) |
||||
elif isinstance(value, dict): |
||||
filename = value["filename"] |
||||
sample = value["data"] |
||||
file_params[key] = (filename, sample, "application/octet-stream") |
||||
elif isinstance(value, file) or hasattr(value, "read"): |
||||
filename = os.path.split(value.name)[1] |
||||
# For the following block refer to DEV-1820 |
||||
try: |
||||
filename.decode("ASCII") |
||||
except (UnicodeDecodeError, UnicodeEncodeError): |
||||
b64_key = key + "name_b64enc" |
||||
byte_value = filename.encode("utf-8") |
||||
b64_value = base64.b64encode(byte_value) |
||||
|
||||
filename = "@param=%s" % b64_key |
||||
req_params[b64_key] = b64_value |
||||
file_params[key] = (filename, value, "application/octet-stream") |
||||
else: |
||||
raise VMRayRESTAPIError("Parameter \"%s\" has unknown type \"%s\"" % (key, type(value))) |
||||
|
||||
# construct request |
||||
if file_params: |
||||
files = file_params |
||||
else: |
||||
files = None |
||||
|
||||
# we need to adjust some stuff for POST requests |
||||
if http_method.lower() == "post": |
||||
req_data = req_params |
||||
req_params = None |
||||
else: |
||||
req_data = None |
||||
|
||||
# do request |
||||
result = requests_func(self.server + api_path, data=req_data, params=req_params, headers={"Authorization": "api_key " + self.api_key}, files=files, verify=self.verify_cert, stream=raw_data) |
||||
handle_rest_api_result(result) |
||||
|
||||
if raw_data: |
||||
return result.raw |
||||
|
||||
# parse result |
||||
try: |
||||
json_result = result.json() |
||||
except ValueError: |
||||
raise ValueError("API returned invalid JSON: %s" % (result.text)) |
||||
|
||||
# if there are no cached elements then return the data |
||||
if "continuation_id" not in json_result: |
||||
return json_result.get("data", None) |
||||
|
||||
data = json_result["data"] |
||||
|
||||
# get cached results |
||||
while "continuation_id" in json_result: |
||||
# send request to server |
||||
result = requests.get("%s/rest/continuation/%u" % (self.server, json_result["continuation_id"]), headers={"Authorization": "api_key " + self.api_key}, verify=self.verify_cert) |
||||
handle_rest_api_result(result) |
||||
|
||||
# parse result |
||||
try: |
||||
json_result = result.json() |
||||
except ValueError: |
||||
raise ValueError("API returned invalid JSON: %s" % (result.text)) |
||||
|
||||
data.extend(json_result["data"]) |
||||
|
||||
return data |
||||
@ -0,0 +1,152 @@ |
||||
#!/usr/bin/env python3 |
||||
|
||||
''' |
||||
Submit sample to VMRay. |
||||
|
||||
Submit a sample to VMRay |
||||
|
||||
TODO: |
||||
# Deal with malicious samples (ZIP file, 'infected') |
||||
# Deal with archive submissions |
||||
|
||||
''' |
||||
|
||||
import json |
||||
import re |
||||
import base64 |
||||
|
||||
import sys |
||||
import os |
||||
base_dir = os.path.dirname(__file__) or '.' |
||||
sys.path.append(base_dir) |
||||
from vmray_rest_api import VMRayRESTAPI, VMRayRESTAPIError |
||||
import io |
||||
|
||||
misperrors = {'error': 'Error'} |
||||
mispattributes = {'input': ['attachment'], 'output': ['text', 'sha1', 'sha256', 'md5', 'link']} |
||||
moduleinfo = {'version': '0.1', 'author': 'Koen Van Impe', |
||||
'description': 'Submit a sample to VMRay', |
||||
'module-type': ['expansion']} |
||||
moduleconfig = ['apikey', 'url', 'shareable', 'do_not_reanalyze', 'do_not_include_vmrayjobids'] |
||||
|
||||
|
||||
include_vmrayjobids = False |
||||
|
||||
|
||||
def handler(q=False): |
||||
global include_vmrayjobids |
||||
|
||||
if q is False: |
||||
return False |
||||
request = json.loads(q) |
||||
|
||||
try: |
||||
data = request.get("data") |
||||
attachment = request.get("attachment") |
||||
data = base64.b64decode(data) |
||||
except: |
||||
misperrors['error'] = "Unable to process submited sample data" |
||||
return misperrors |
||||
|
||||
if (request["config"].get("apikey") is None) or (request["config"].get("url") is None): |
||||
misperrors["error"] = "Missing API key or server URL (hint: try cloud.vmray.com)" |
||||
return misperrors |
||||
|
||||
api = VMRayRESTAPI(request["config"].get("url"), request["config"].get("apikey"), False) |
||||
|
||||
shareable = request["config"].get("shareable") |
||||
do_not_reanalyze = request["config"].get("do_not_reanalyze") |
||||
do_not_include_vmrayjobids = request["config"].get("do_not_include_vmrayjobids") |
||||
|
||||
# Do we want the sample to be shared? |
||||
if shareable == "True": |
||||
shareable = True |
||||
else: |
||||
shareable = False |
||||
|
||||
# Always reanalyze the sample? |
||||
if do_not_reanalyze == "True": |
||||
do_not_reanalyze = True |
||||
else: |
||||
do_not_reanalyze = False |
||||
reanalyze = not do_not_reanalyze |
||||
|
||||
# Include the references to VMRay job IDs |
||||
if do_not_include_vmrayjobids == "True": |
||||
do_not_include_vmrayjobids = True |
||||
else: |
||||
do_not_include_vmrayjobids = False |
||||
include_vmrayjobids = not do_not_include_vmrayjobids |
||||
|
||||
if data and attachment: |
||||
args = {} |
||||
args["shareable"] = shareable |
||||
args["sample_file"] = {'data': io.BytesIO( data ) , 'filename': attachment } |
||||
args["reanalyze"] = reanalyze |
||||
|
||||
try: |
||||
vmraydata = vmraySubmit(api, args) |
||||
if vmraydata["errors"]: |
||||
misperrors['error'] = "VMRay: %s" % vmraydata["errors"][0]["error_msg"] |
||||
return misperrors |
||||
else: |
||||
return vmrayProcess(vmraydata) |
||||
except: |
||||
misperrors['error'] = "Problem when calling API." |
||||
return misperrors |
||||
else: |
||||
misperrors['error'] = "No sample data or filename." |
||||
return misperrors |
||||
|
||||
|
||||
def introspection(): |
||||
return mispattributes |
||||
|
||||
|
||||
def version(): |
||||
moduleinfo['config'] = moduleconfig |
||||
return moduleinfo |
||||
|
||||
|
||||
def vmrayProcess(vmraydata): |
||||
''' Process the JSON file returned by vmray''' |
||||
if vmraydata: |
||||
try: |
||||
submissions = vmraydata["submissions"][0] |
||||
jobs = vmraydata["jobs"] |
||||
|
||||
# Result received? |
||||
if submissions and jobs: |
||||
r = {'results': []} |
||||
r["results"].append( {"types": "md5", "values": submissions["submission_sample_md5"]} ) |
||||
r["results"].append( {"types": "sha1", "values": submissions["submission_sample_sha1"]} ) |
||||
r["results"].append( {"types": "sha256", "values": submissions["submission_sample_sha256"]} ) |
||||
r["results"].append( {"types": "text", "values": "VMRay Sample ID: %s" % submissions["submission_sample_id"]} ) |
||||
r["results"].append( {"types": "text", "values": "VMRay Submission ID: %s" % submissions["submission_id"]} ) |
||||
r["results"].append( {"types": "text", "values": "VMRay Submission Sample IP: %s" % submissions["submission_ip_ip"]} ) |
||||
r["results"].append( {"types": "link", "values": submissions["submission_webif_url"]} ) |
||||
|
||||
# Include data from different jobs |
||||
if include_vmrayjobids: |
||||
for job in jobs: |
||||
job_id = job["job_id"] |
||||
job_vm_name = job["job_vm_name"] |
||||
job_configuration_name = job["job_configuration_name"] |
||||
r["results"].append( {"types": "text", "values": "VMRay Job ID %s (%s - %s)" % (job_id, job_vm_name, job_configuration_name) }) |
||||
return r |
||||
else: |
||||
misperrors['error'] = "No valid results returned." |
||||
return misperrors |
||||
except: |
||||
misperrors['error'] = "No valid submission data returned." |
||||
return misperrors |
||||
else: |
||||
misperrors['error'] = "Unable to parse results." |
||||
return misperrors |
||||
|
||||
|
||||
def vmraySubmit(api, args): |
||||
''' Submit the sample to VMRay''' |
||||
vmraydata = api.call("POST", "/rest/sample/submit", args) |
||||
return vmraydata |
||||
|
||||
@ -1 +1 @@ |
||||
__all__ = ['testimport', 'ocr', 'stiximport'] |
||||
__all__ = ['vmray_import','testimport', 'ocr', 'stiximport'] |
||||
|
||||
@ -0,0 +1,300 @@ |
||||
#!/usr/bin/env python3 |
||||
|
||||
''' |
||||
Import VMRay results. |
||||
|
||||
This version supports import from different analyze jobs, starting from one sample |
||||
(the supplied sample_id). |
||||
|
||||
Requires "vmray_rest_api" |
||||
|
||||
TODO: |
||||
# Import one job (analyze_id) |
||||
# Import STIX package (XML version) |
||||
|
||||
''' |
||||
|
||||
import json |
||||
import re |
||||
import sys |
||||
import os |
||||
base_dir = os.path.dirname(__file__) or '.' |
||||
sys.path.append(base_dir) |
||||
from vmray_rest_api import VMRayRESTAPI, VMRayRESTAPIError |
||||
|
||||
misperrors = {'error': 'Error'} |
||||
inputSource = [] |
||||
moduleinfo = {'version': '0.1', 'author': 'Koen Van Impe', |
||||
'description': 'Import VMRay (VTI) results', |
||||
'module-type': ['import']} |
||||
userConfig = { |
||||
'include_textdescr': { |
||||
'type': 'Boolean', |
||||
'message': 'Include textual description' |
||||
}, |
||||
'include_analysisid': { |
||||
'type': 'Boolean', |
||||
'message': 'Include VMRay analysis_id text' |
||||
}, |
||||
'only_network_info': { |
||||
'type': 'Boolean', |
||||
'message': 'Only include network (src-ip, hostname, domain, ...) information' |
||||
}, |
||||
'sample_id': { |
||||
'type': 'Integer', |
||||
'errorMessage': 'Expected a sample ID', |
||||
'message': 'The VMRay sample_id' |
||||
} |
||||
}; |
||||
moduleconfig = ['apikey', 'url'] |
||||
|
||||
include_textdescr = False |
||||
include_analysisid = False |
||||
only_network_info = False |
||||
|
||||
def handler(q=False): |
||||
global include_textdescr |
||||
global include_analysisid |
||||
global only_network_info |
||||
|
||||
if q is False: |
||||
return False |
||||
request = json.loads(q) |
||||
|
||||
include_textdescr = request["config"].get("include_textdescr") |
||||
include_analysisid = request["config"].get("include_analysisid") |
||||
only_network_info = request["config"].get("only_network_info") |
||||
if include_textdescr == "1": |
||||
include_textdescr = True |
||||
else: |
||||
include_textdescr = False |
||||
if include_analysisid == "1": |
||||
include_analysisid = True |
||||
else: |
||||
include_analysisid = False |
||||
if only_network_info == "1": |
||||
only_network_info = True |
||||
else: |
||||
only_network_info = False |
||||
|
||||
sample_id = int(request["config"].get("sample_id")) |
||||
|
||||
if (request["config"].get("apikey") is None) or (request["config"].get("url") is None): |
||||
misperrors["error"] = "Missing API key or server URL (hint: try cloud.vmray.com)" |
||||
return misperrors |
||||
|
||||
if sample_id > 0: |
||||
try: |
||||
api = VMRayRESTAPI(request["config"].get("url"), request["config"].get("apikey"), False) |
||||
vmray_results = {'results': []} |
||||
# Get all information on the sample, returns a set of finished analyze jobs |
||||
data = vmrayGetInfoAnalysis(api, sample_id) |
||||
if data["data"]: |
||||
for analysis in data["data"]: |
||||
analysis_id = analysis["analysis_id"] |
||||
|
||||
if analysis_id > 0: |
||||
# Get the details for an analyze job |
||||
analysis_data = vmrayDownloadAnalysis(api, analysis_id) |
||||
|
||||
if analysis_data: |
||||
p = vmrayVtiPatterns(analysis_data["vti_patterns"]) |
||||
if p: |
||||
if include_analysisid: |
||||
url1 = "https://cloud.vmray.com/user/analysis/view?from_sample_id=%u" % sample_id |
||||
url2 = "&id=%u" % analysis_id |
||||
url3 = "&sub=%2Freport%2Foverview.html" |
||||
p["results"].append({ "values": url1 + url2 + url3, "types": "link" }) |
||||
vmray_results = {'results': vmray_results["results"] + p["results"] } |
||||
|
||||
# Clean up (remove doubles) |
||||
vmray_results = vmrayCleanup(vmray_results) |
||||
return vmray_results |
||||
else: |
||||
misperrors['error'] = "Unable to fetch sample id %u" % (sample_id) |
||||
return misperrors |
||||
except: |
||||
misperrors['error'] = "Unable to access VMRay API" |
||||
return misperrors |
||||
else: |
||||
misperrors['error'] = "Not a valid sample id" |
||||
return misperrors |
||||
|
||||
|
||||
|
||||
def introspection(): |
||||
modulesetup = {} |
||||
try: |
||||
userConfig |
||||
modulesetup['userConfig'] = userConfig |
||||
except NameError: |
||||
pass |
||||
try: |
||||
inputSource |
||||
modulesetup['inputSource'] = inputSource |
||||
except NameError: |
||||
pass |
||||
return modulesetup |
||||
|
||||
|
||||
def version(): |
||||
moduleinfo['config'] = moduleconfig |
||||
return moduleinfo |
||||
|
||||
|
||||
def vmrayGetInfoAnalysis(api, sample_id): |
||||
''' Get information from a sample, returns a set of analyzed reports''' |
||||
|
||||
if sample_id: |
||||
data = api.call("GET", "/rest/analysis/sample/%u" % (sample_id), raw_data=True) |
||||
return json.loads(data.read().decode()) |
||||
else: |
||||
return False |
||||
|
||||
|
||||
def vmrayDownloadAnalysis(api, analysis_id): |
||||
''' Get the details from an analysis''' |
||||
if analysis_id: |
||||
data = api.call("GET", "/rest/analysis/%u/archive/additional/vti_result.json" % (analysis_id), raw_data=True) |
||||
return json.loads(data.read().decode()) |
||||
else: |
||||
return False |
||||
|
||||
|
||||
def vmrayVtiPatterns(vti_patterns): |
||||
''' Match the VTI patterns to MISP data''' |
||||
|
||||
if vti_patterns: |
||||
r = {'results': []} |
||||
y = {'results': []} |
||||
|
||||
for pattern in vti_patterns: |
||||
content = False |
||||
if pattern["category"] == "_network" and pattern["operation"] == "_download_data": |
||||
content = vmrayGeneric(pattern, "url", 1) |
||||
elif pattern["category"] == "_network" and pattern["operation"] == "_connect": |
||||
content = vmrayConnect(pattern) |
||||
|
||||
elif only_network_info == False and pattern["category"] == "_process" and pattern["operation"] == "_alloc_wx_page": |
||||
content = vmrayGeneric(pattern) |
||||
elif only_network_info == False and pattern["category"] == "_process" and pattern["operation"] == "_install_ipc_endpoint": |
||||
content = vmrayGeneric(pattern, "mutex", 1) |
||||
|
||||
elif only_network_info == False and pattern["category"] == "_anti_analysis" and pattern["operation"] == "_delay_execution": |
||||
content = vmrayGeneric(pattern) |
||||
elif only_network_info == False and pattern["category"] == "_anti_analysis" and pattern["operation"] == "_dynamic_api_usage": |
||||
content = vmrayGeneric(pattern) |
||||
|
||||
elif only_network_info == False and pattern["category"] == "_static" and pattern["operation"] == "_drop_pe_file": |
||||
content = vmrayGeneric(pattern, "filename", 1) |
||||
elif only_network_info == False and pattern["category"] == "_static" and pattern["operation"] == "_execute_dropped_pe_file": |
||||
content = vmrayGeneric(pattern, "filename", 1) |
||||
|
||||
elif only_network_info == False and pattern["category"] == "_injection" and pattern["operation"] == "_modify_memory": |
||||
content = vmrayGeneric(pattern) |
||||
elif only_network_info == False and pattern["category"] == "_injection" and pattern["operation"] == "_modify_control_flow": |
||||
content = vmrayGeneric(pattern) |
||||
elif only_network_info == False and pattern["category"] == "_file_system" and pattern["operation"] == "_create_many_files": |
||||
content = vmrayGeneric(pattern) |
||||
|
||||
elif only_network_info == False and pattern["category"] == "_persistence" and pattern["operation"] == "_install_startup_script": |
||||
content = vmrayGeneric(pattern, "regkey", 1) |
||||
elif only_network_info == False and pattern["category"] == "_os" and pattern["operation"] == "_enable_process_privileges": |
||||
content = vmrayGeneric(pattern) |
||||
|
||||
if content: |
||||
r["results"].append( content["attributes"] ) |
||||
r["results"].append( content["text"] ) |
||||
|
||||
# Remove empty results |
||||
r["results"] = [x for x in r["results"] if isinstance(x, dict) and len(x["values"]) != 0] |
||||
for el in r["results"]: |
||||
if not el in y["results"]: |
||||
y["results"].append( el ) |
||||
return y |
||||
else: |
||||
return False |
||||
|
||||
|
||||
def vmrayCleanup(x): |
||||
''' Remove doubles''' |
||||
y = {'results': []} |
||||
|
||||
for el in x["results"]: |
||||
if not el in y["results"]: |
||||
y["results"].append( el ) |
||||
return y |
||||
|
||||
|
||||
def vmraySanitizeInput(s): |
||||
''' Sanitize some input so it gets properly imported in MISP''' |
||||
if s: |
||||
s = s.replace('"','') |
||||
s = re.sub('\\\\', r'\\', s) |
||||
return s |
||||
else: |
||||
return False |
||||
|
||||
|
||||
def vmrayGeneric(el, attr = "", attrpos = 1): |
||||
''' Convert a 'generic' VTI pattern to MISP data''' |
||||
|
||||
r = {"values": []} |
||||
f = {"values": []} |
||||
|
||||
if el: |
||||
content = el["technique_desc"] |
||||
if content: |
||||
if attr: |
||||
content_split = content.split("\"") |
||||
content_split[attrpos] = vmraySanitizeInput(content_split[attrpos]) |
||||
r["values"].append(content_split[attrpos]) |
||||
r["types"] = [attr] |
||||
|
||||
# Adding the value also as text to get the extra description, |
||||
# but this is pretty useless for "url" |
||||
if include_textdescr and attr != "url": |
||||
f["values"].append(vmraySanitizeInput(content)) |
||||
f["types"] = ["text"] |
||||
|
||||
return { "text": f, |
||||
"attributes": r} |
||||
else: |
||||
return False |
||||
else: |
||||
return False |
||||
|
||||
|
||||
def vmrayConnect(el): |
||||
''' Extension of vmrayGeneric , parse network connect data''' |
||||
ipre = re.compile("([0-9]{1,3}.){3}[0-9]{1,3}") |
||||
|
||||
r = {"values": []} |
||||
f = {"values": []} |
||||
|
||||
if el: |
||||
content = el["technique_desc"] |
||||
if content: |
||||
target = content.split("\"") |
||||
port = (target[1].split(":"))[1] |
||||
host = (target[1].split(":"))[0] |
||||
if ipre.match(str(host)): |
||||
r["values"].append(host) |
||||
r["types"] = ["ip-dst"] |
||||
else: |
||||
r["values"].append(host) |
||||
r["types"] = ["domain", "hostname"] |
||||
|
||||
f["values"].append(vmraySanitizeInput(target[1])) |
||||
f["types"] = ["text"] |
||||
|
||||
if include_textdescr: |
||||
f["values"].append(vmraySanitizeInput(content)) |
||||
f["types"] = ["text"] |
||||
|
||||
return { "text": f, |
||||
"attributes": r} |
||||
else: |
||||
return False |
||||
else: |
||||
return False |
||||
@ -0,0 +1,152 @@ |
||||
#!/usr/bin/python3 |
||||
"""Python client library for VMRay REST API""" |
||||
|
||||
import base64 |
||||
import datetime |
||||
import os.path |
||||
import requests |
||||
#import urlparse |
||||
import urllib.parse |
||||
|
||||
from io import IOBase |
||||
from io import BytesIO |
||||
|
||||
# disable nasty certification warning |
||||
# pylint: disable=no-member |
||||
try: |
||||
requests.packages.urllib3.disable_warnings() |
||||
except AttributeError: |
||||
try: |
||||
import urllib3 |
||||
try: |
||||
urllib3.disable_warnings() |
||||
except AttributeError: |
||||
pass |
||||
except ImportError: |
||||
pass |
||||
|
||||
# pylint: disable= |
||||
|
||||
|
||||
class VMRayRESTAPIError(Exception): |
||||
"""Exception class that is used when API returns an error""" |
||||
|
||||
def __init__(self, *args, **kwargs): |
||||
self.status_code = kwargs.pop("status_code", None) |
||||
Exception.__init__(self, *args, **kwargs) |
||||
|
||||
|
||||
def handle_rest_api_result(result): |
||||
"""Handle result of API request (check for errors)""" |
||||
|
||||
if (result.status_code < 200) or (result.status_code > 299): |
||||
try: |
||||
json_result = result.json() |
||||
except ValueError: |
||||
raise VMRayRESTAPIError("API returned error %u: %s" % (result.status_code, result.text), status_code=result.status_code) |
||||
|
||||
raise VMRayRESTAPIError(json_result.get("error_msg", "Unknown error"), status_code=result.status_code) |
||||
|
||||
|
||||
class VMRayRESTAPI(object): |
||||
"""VMRay REST API class""" |
||||
|
||||
def __init__(self, server, api_key, verify_cert=True): |
||||
# split server URL into components |
||||
url_desc = urllib.parse.urlsplit(server) |
||||
|
||||
# assume HTTPS if no scheme is specified |
||||
if url_desc.scheme == "": |
||||
server = "https://" + server |
||||
|
||||
# save variables |
||||
self.server = server |
||||
self.api_key = api_key |
||||
self.verify_cert = verify_cert |
||||
|
||||
def call(self, http_method, api_path, params=None, raw_data=False): |
||||
"""Call VMRay REST API""" |
||||
|
||||
# get function of requests package |
||||
requests_func = getattr(requests, http_method.lower()) |
||||
|
||||
# parse parameters |
||||
req_params = {} |
||||
file_params = {} |
||||
|
||||
if params is not None: |
||||
for key, value in params.items(): |
||||
if isinstance(value, (datetime.date, |
||||
datetime.datetime, |
||||
float, |
||||
int)): |
||||
req_params[key] = str(value) |
||||
elif isinstance(value, str): |
||||
req_params[key] = str(value) |
||||
elif isinstance(value, dict): |
||||
filename = value["filename"] |
||||
sample = value["data"] |
||||
file_params[key] = (filename, sample, "application/octet-stream") |
||||
elif isinstance(value, file) or hasattr(value, "read"): |
||||
filename = os.path.split(value.name)[1] |
||||
# For the following block refer to DEV-1820 |
||||
try: |
||||
filename.decode("ASCII") |
||||
except (UnicodeDecodeError, UnicodeEncodeError): |
||||
b64_key = key + "name_b64enc" |
||||
byte_value = filename.encode("utf-8") |
||||
b64_value = base64.b64encode(byte_value) |
||||
|
||||
filename = "@param=%s" % b64_key |
||||
req_params[b64_key] = b64_value |
||||
file_params[key] = (filename, value, "application/octet-stream") |
||||
else: |
||||
raise VMRayRESTAPIError("Parameter \"%s\" has unknown type \"%s\"" % (key, type(value))) |
||||
|
||||
# construct request |
||||
if file_params: |
||||
files = file_params |
||||
else: |
||||
files = None |
||||
|
||||
# we need to adjust some stuff for POST requests |
||||
if http_method.lower() == "post": |
||||
req_data = req_params |
||||
req_params = None |
||||
else: |
||||
req_data = None |
||||
|
||||
# do request |
||||
result = requests_func(self.server + api_path, data=req_data, params=req_params, headers={"Authorization": "api_key " + self.api_key}, files=files, verify=self.verify_cert, stream=raw_data) |
||||
handle_rest_api_result(result) |
||||
|
||||
if raw_data: |
||||
return result.raw |
||||
|
||||
# parse result |
||||
try: |
||||
json_result = result.json() |
||||
except ValueError: |
||||
raise ValueError("API returned invalid JSON: %s" % (result.text)) |
||||
|
||||
# if there are no cached elements then return the data |
||||
if "continuation_id" not in json_result: |
||||
return json_result.get("data", None) |
||||
|
||||
data = json_result["data"] |
||||
|
||||
# get cached results |
||||
while "continuation_id" in json_result: |
||||
# send request to server |
||||
result = requests.get("%s/rest/continuation/%u" % (self.server, json_result["continuation_id"]), headers={"Authorization": "api_key " + self.api_key}, verify=self.verify_cert) |
||||
handle_rest_api_result(result) |
||||
|
||||
# parse result |
||||
try: |
||||
json_result = result.json() |
||||
except ValueError: |
||||
raise ValueError("API returned invalid JSON: %s" % (result.text)) |
||||
|
||||
data.extend(json_result["data"]) |
||||
|
||||
return data |
||||
Loading…
Reference in new issue