Add timeout for the modules, cleanup.

pull/52/head
Raphaël Vinot 2016-08-25 17:36:28 +02:00
parent 1034f73479
commit c69fae087c
3 changed files with 26 additions and 10 deletions

View File

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

1
tests/body_timeout.json Normal file
View File

@ -0,0 +1 @@
{"module": "dns", "hostname": "www.circl.lu", "timeout": 1, "config" : {"nameserver":"8.8.8.8"}}

View File

@ -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: