diff --git a/pymisp/api.py b/pymisp/api.py index 363ab65..c1f3470 100644 --- a/pymisp/api.py +++ b/pymisp/api.py @@ -27,6 +27,13 @@ try: except ImportError: HAVE_REQUESTS = False + +try: + from requests_futures.sessions import FuturesSession + ASYNC_OK = True +except ImportError: + ASYNC_OK = False + from . import __version__ from .exceptions import PyMISPError, SearchError, MissingDependency, NoURL, NoKey from .mispevent import MISPEvent, MISPAttribute, EncodeUpdate @@ -80,9 +87,10 @@ class PyMISP(object): :param debug: print all the messages received from the server :param proxies: Proxy dict as describes here: http://docs.python-requests.org/en/master/user/advanced/#proxies :param cert: Client certificate, as described there: http://docs.python-requests.org/en/master/user/advanced/#ssl-cert-verification + :param asynch: Use asynchronous processing where possible """ - def __init__(self, url, key, ssl=True, out_type='json', debug=False, proxies=None, cert=None): + def __init__(self, url, key, ssl=True, out_type='json', debug=False, proxies=None, cert=None, asynch=False): if not url: raise NoURL('Please provide the URL of your MISP instance.') if not key: @@ -93,6 +101,11 @@ class PyMISP(object): self.ssl = ssl self.proxies = proxies self.cert = cert + self.asynch = asynch + if asynch and not ASYNC_OK: + warnings.warn("You turned on Async, but don't have requests_futures installed") + self.asynch = False + self.ressources_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'data') if out_type != 'json': raise PyMISPError('The only output type supported by PyMISP is JSON. If you still rely on XML, use PyMISP v2.4.49') @@ -137,12 +150,15 @@ class PyMISP(object): self.category_type_mapping = self.describe_types['category_type_mappings'] self.sane_default = self.describe_types['sane_defaults'] - def __prepare_session(self, output='json'): + def __prepare_session(self, output='json', async_implemented=False): """Prepare the headers of the session""" if not HAVE_REQUESTS: raise MissingDependency('Missing dependency, install requests (`pip install requests`)') - session = requests.Session() + if self.asynch and async_implemented: + session = FuturesSession() + else: + session = requests.Session() session.verify = self.ssl session.proxies = self.proxies session.cert = self.cert @@ -774,8 +790,11 @@ class PyMISP(object): def search_index(self, published=None, eventid=None, tag=None, datefrom=None, dateuntil=None, eventinfo=None, threatlevel=None, distribution=None, - analysis=None, attribute=None, org=None): + analysis=None, attribute=None, org=None, async_callback=None): """Search only at the index level. Use ! infront of value as NOT, default OR + If using async, give a callback that takes 2 args, session and response: + basic usage is + pymisp.search_index(..., async_callback=lambda ses,resp: print(resp.json())) :param published: Published (0,1) :param eventid: Evend ID(s) | str or list @@ -787,6 +806,7 @@ class PyMISP(object): :param distribution: Distribution level(s) (0,1,2,3) | str or list :param analysis: Analysis level(s) (0,1,2) | str or list :param org: Organisation(s) | str or list + :param async_callback: Function to call when the request returns (if running async) """ allowed = {'published': published, 'eventid': eventid, 'tag': tag, 'Dateuntil': dateuntil, 'Datefrom': datefrom, 'eventinfo': eventinfo, 'threatlevel': threatlevel, @@ -799,6 +819,7 @@ class PyMISP(object): to_post = {} for rule in allowed.keys(): + if allowed.get(rule) is None: continue param = allowed[rule] @@ -809,10 +830,17 @@ class PyMISP(object): if not set(param).issubset(rule_levels[rule]): raise SearchError('Values in your {} are invalid, has to be in {}'.format(rule, ', '.join(str(x) for x in rule_levels[rule]))) to_post[rule] = '|'.join(str(x) for x in param) - session = self.__prepare_session() + session = self.__prepare_session(async_implemented=True) url = urljoin(self.root_url, buildup_url) - response = session.post(url, data=json.dumps(to_post)) - return self._check_response(response) + + if self.asynch: + if not async_callback: + warnings.warn("You haven't provided a callback!") + response = session.get(url, background_callback=async_callback) + + else: + response = session.post(url, data=json.dumps(to_post)) + return self._check_response(response) def search_all(self, value): query = {'value': value, 'searchall': 1}