diff --git a/misp_modules/__init__.py b/misp_modules/__init__.py index 81c92b15..a2562c4e 100644 --- a/misp_modules/__init__.py +++ b/misp_modules/__init__.py @@ -28,8 +28,10 @@ import logging import fnmatch import argparse import re +import datetime import tornado.web +import tornado.process from tornado.ioloop import IOLoop from tornado.concurrent import run_on_executor from concurrent.futures import ThreadPoolExecutor @@ -155,8 +157,6 @@ def load_package_modules(): class ListModules(tornado.web.RequestHandler): def get(self): - global mhandlers - global loaded_modules ret = [] for module in loaded_modules: x = {} @@ -170,26 +170,38 @@ class ListModules(tornado.web.RequestHandler): class QueryModule(tornado.web.RequestHandler): - try: - # Python 3.5, use as many threads as possible - executor = ThreadPoolExecutor() - except: - executor = ThreadPoolExecutor(5) + + # Default value in Python 3.5 + # https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor + nb_threads = tornado.process.cpu_count() * 5 + executor = ThreadPoolExecutor(nb_threads) @run_on_executor def run_request(self, jsonpayload): x = json.loads(jsonpayload) log.debug('MISP QueryModule request {0}'.format(jsonpayload)) - return mhandlers[x['module']].handler(q=jsonpayload) + response = mhandlers[x['module']].handler(q=jsonpayload) + return json.dumps(response) @tornado.gen.coroutine def post(self): try: jsonpayload = self.request.body.decode('utf-8') - response = yield self.run_request(jsonpayload) - self.write(json.dumps(response)) + dict_payload = json.loads(jsonpayload) + if dict_payload.get('timeout'): + timeout = datetime.timedelta(seconds=int(dict_payload.get('timeout'))) + else: + timeout = datetime.timedelta(seconds=30) + response = yield tornado.gen.with_timeout(timeout, self.run_request(jsonpayload)) + self.write(response) + except tornado.gen.TimeoutError: + log.warning('Timeout on {} '.format(dict_payload['module'])) + self.write(json.dumps({'error': 'Timeout.'})) except Exception: + self.write(json.dumps({'error': 'Something went wrong, look in the server logs for details'})) log.exception('Something went wrong:') + finally: + self.finish() def main(): diff --git a/tests/body_timeout.json b/tests/body_timeout.json new file mode 100644 index 00000000..f524c673 --- /dev/null +++ b/tests/body_timeout.json @@ -0,0 +1 @@ +{"module": "dns", "hostname": "www.circl.lu", "timeout": 1, "config" : {"nameserver":"8.8.8.8"}} diff --git a/tests/test.py b/tests/test.py index 1314479a..80b445fd 100644 --- a/tests/test.py +++ b/tests/test.py @@ -26,6 +26,9 @@ class TestModules(unittest.TestCase): with open('tests/body.json', 'r') as f: response = requests.post(self.url + "query", data=f.read()) print(response.json()) + with open('tests/body_timeout.json', 'r') as f: + response = requests.post(self.url + "query", data=f.read()) + print(response.json()) def test_stix(self): with open("tests/stix.xml", "r") as f: