From 193d7fd0bce42b1a4a0692d3a9334afb96cbf766 Mon Sep 17 00:00:00 2001
From: Jakub Onderka <jakub.onderka@gmail.com>
Date: Sat, 6 Jan 2024 14:58:47 +0100
Subject: [PATCH] new: [internal] Avoid double JSON decoding

---
 misp_modules/__init__.py                           | 14 ++++++++++----
 misp_modules/modules/expansion/circl_passivedns.py | 10 +++-------
 misp_modules/modules/expansion/clamav.py           |  8 +-------
 misp_modules/modules/expansion/virustotal.py       |  8 ++------
 4 files changed, 16 insertions(+), 24 deletions(-)

diff --git a/misp_modules/__init__.py b/misp_modules/__init__.py
index 2f342e3..b628d78 100644
--- a/misp_modules/__init__.py
+++ b/misp_modules/__init__.py
@@ -187,9 +187,14 @@ class QueryModule(tornado.web.RequestHandler):
     executor = ThreadPoolExecutor(nb_threads)
 
     @run_on_executor
-    def run_request(self, module_name, json_payload):
-        log.debug('MISP QueryModule request %s', json_payload)
-        response = mhandlers[module_name].handler(q=json_payload)
+    def run_request(self, module_name, json_payload, dict_payload):
+        log.debug('MISP QueryModule %s request %s', module_name, json_payload)
+        module = mhandlers[module_name]
+        if getattr(module, "dict_handler", None):
+            # New method that avoids double JSON decoding, new modules should define dict_handler
+            response = module.dict_handler(request=dict_payload)
+        else:
+            response = module.handler(q=json_payload)
         return json.dumps(response)
 
     @tornado.gen.coroutine
@@ -201,7 +206,8 @@ class QueryModule(tornado.web.RequestHandler):
                 timeout = datetime.timedelta(seconds=int(dict_payload.get('timeout')))
             else:
                 timeout = datetime.timedelta(seconds=300)
-            response = yield tornado.gen.with_timeout(timeout, self.run_request(dict_payload['module'], json_payload))
+            future = self.run_request(dict_payload['module'], json_payload, dict_payload)
+            response = yield tornado.gen.with_timeout(timeout, future)
             self.write(response)
         except tornado.gen.TimeoutError:
             log.warning('Timeout on {}'.format(dict_payload['module']))
diff --git a/misp_modules/modules/expansion/circl_passivedns.py b/misp_modules/modules/expansion/circl_passivedns.py
index 5f98314..eca78c8 100755
--- a/misp_modules/modules/expansion/circl_passivedns.py
+++ b/misp_modules/modules/expansion/circl_passivedns.py
@@ -1,4 +1,3 @@
-import json
 import pypdns
 from . import check_input_attribute, standard_error_message
 from pymisp import MISPAttribute, MISPEvent, MISPObject
@@ -10,7 +9,7 @@ moduleinfo = {'version': '0.2', 'author': 'Alexandre Dulaunoy',
 moduleconfig = ['username', 'password']
 
 
-class PassiveDNSParser():
+class PassiveDNSParser:
     def __init__(self, attribute, authentication):
         self.misp_event = MISPEvent()
         self.attribute = MISPAttribute()
@@ -21,7 +20,7 @@ class PassiveDNSParser():
     def get_results(self):
         if hasattr(self, 'result'):
             return self.result
-        event = json.loads(self.misp_event.to_json())
+        event = self.misp_event.to_dict()
         results = {key: event[key] for key in ('Attribute', 'Object')}
         return {'results': results}
 
@@ -50,10 +49,7 @@ class PassiveDNSParser():
             self.misp_event.add_object(**pdns_object)
 
 
-def handler(q=False):
-    if q is False:
-        return False
-    request = json.loads(q)
+def dict_handler(request: dict):
     if not request.get('config'):
         return {'error': 'CIRCL Passive DNS authentication is missing.'}
     if not request['config'].get('username') or not request['config'].get('password'):
diff --git a/misp_modules/modules/expansion/clamav.py b/misp_modules/modules/expansion/clamav.py
index 0b789f7..bdff3b5 100644
--- a/misp_modules/modules/expansion/clamav.py
+++ b/misp_modules/modules/expansion/clamav.py
@@ -1,6 +1,5 @@
 import base64
 import io
-import json
 import logging
 import sys
 import zipfile
@@ -58,12 +57,7 @@ def connect_to_clamav(connection_string: str) -> clamd.ClamdNetworkSocket:
         raise Exception("ClamAV connection string is invalid. It must be unix socket path with 'unix://' prefix or IP:PORT.")
 
 
-def handler(q=False):
-    if q is False:
-        return False
-
-    request = json.loads(q)
-
+def dict_handler(request: dict):
     connection_string: str = request["config"].get("connection")
     if not connection_string:
         return {"error": "No ClamAV connection string provided"}
diff --git a/misp_modules/modules/expansion/virustotal.py b/misp_modules/modules/expansion/virustotal.py
index 93d0966..29f0550 100644
--- a/misp_modules/modules/expansion/virustotal.py
+++ b/misp_modules/modules/expansion/virustotal.py
@@ -1,4 +1,3 @@
-import json
 from urllib.parse import urlparse
 import vt
 from . import check_input_attribute, standard_error_message
@@ -45,7 +44,7 @@ class VirusTotalParser:
         self.input_types_mapping[self.attribute.type](self.attribute.value)
 
     def get_result(self) -> dict:
-        event = json.loads(self.misp_event.to_json())
+        event = self.misp_event.to_dict()
         results = {key: event[key] for key in ('Attribute', 'Object') if (key in event and event[key])}
         return {'results': results}
 
@@ -257,10 +256,7 @@ def parse_error(status_code: int) -> str:
     return "VirusTotal may not be accessible."
 
 
-def handler(q=False):
-    if q is False:
-        return False
-    request = json.loads(q)
+def dict_handler(request: dict):
     if not request.get('config') or not request['config'].get('apikey'):
         misperrors['error'] = 'A VirusTotal api key is required for this module.'
         return misperrors