chg: [internal] Add support for orjson

pull/648/head
Jakub Onderka 2024-01-06 12:11:13 +01:00
parent ea2697c5ce
commit cbaa2f85a2
1 changed files with 12 additions and 8 deletions

View File

@ -22,7 +22,6 @@ import os
import signal import signal
import sys import sys
import importlib import importlib
import json
import logging import logging
import fnmatch import fnmatch
import argparse import argparse
@ -30,6 +29,11 @@ import re
import datetime import datetime
import psutil import psutil
try:
import orjson as json
except ImportError:
import json
import tornado.web import tornado.web
import tornado.process import tornado.process
from tornado.ioloop import IOLoop from tornado.ioloop import IOLoop
@ -182,24 +186,24 @@ class QueryModule(tornado.web.RequestHandler):
executor = ThreadPoolExecutor(nb_threads) executor = ThreadPoolExecutor(nb_threads)
@run_on_executor @run_on_executor
def run_request(self, module, jsonpayload): def run_request(self, module, json_payload):
log.debug('MISP QueryModule request {0}'.format(jsonpayload)) log.debug(f'MISP QueryModule request {json_payload}')
response = mhandlers[module].handler(q=jsonpayload) response = mhandlers[module].handler(q=json_payload)
return json.dumps(response) return json.dumps(response)
@tornado.gen.coroutine @tornado.gen.coroutine
def post(self): def post(self):
try: try:
jsonpayload = self.request.body.decode('utf-8') json_payload = self.request.body
dict_payload = json.loads(jsonpayload) dict_payload = json.loads(json_payload)
if dict_payload.get('timeout'): if dict_payload.get('timeout'):
timeout = datetime.timedelta(seconds=int(dict_payload.get('timeout'))) timeout = datetime.timedelta(seconds=int(dict_payload.get('timeout')))
else: else:
timeout = datetime.timedelta(seconds=300) timeout = datetime.timedelta(seconds=300)
response = yield tornado.gen.with_timeout(timeout, self.run_request(dict_payload['module'], jsonpayload)) response = yield tornado.gen.with_timeout(timeout, self.run_request(dict_payload['module'], json_payload))
self.write(response) self.write(response)
except tornado.gen.TimeoutError: except tornado.gen.TimeoutError:
log.warning('Timeout on {} '.format(dict_payload['module'])) log.warning('Timeout on {}'.format(dict_payload['module']))
self.write(json.dumps({'error': 'Timeout.'})) self.write(json.dumps({'error': 'Timeout.'}))
except Exception: except Exception:
self.write(json.dumps({'error': 'Something went wrong, look in the server logs for details'})) self.write(json.dumps({'error': 'Something went wrong, look in the server logs for details'}))