From c69fae087cdb24c85f6d7d2c5adb02132f8823e8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= <raphael@vinot.info>
Date: Thu, 25 Aug 2016 17:36:28 +0200
Subject: [PATCH] Add timeout for the modules, cleanup.

---
 misp_modules/__init__.py | 32 ++++++++++++++++++++++----------
 tests/body_timeout.json  |  1 +
 tests/test.py            |  3 +++
 3 files changed, 26 insertions(+), 10 deletions(-)
 create mode 100644 tests/body_timeout.json

diff --git a/misp_modules/__init__.py b/misp_modules/__init__.py
index 81c92b1..a2562c4 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 0000000..f524c67
--- /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 1314479..80b445f 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: