From 148167395bb8085d1dd2d59146c5269b274047e7 Mon Sep 17 00:00:00 2001
From: Marcus Schiesser <mail@marcusschiesser.de>
Date: Mon, 18 Jul 2022 10:12:46 +0000
Subject: [PATCH] add decorator for json rest endpoints

---
 setup.py                         |  3 +-
 splunklib/customrest/__init__.py |  6 ++++
 splunklib/customrest/json.py     | 51 ++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 1 deletion(-)
 create mode 100644 splunklib/customrest/__init__.py
 create mode 100644 splunklib/customrest/json.py

diff --git a/setup.py b/setup.py
index 54856a28..169d7915 100755
--- a/setup.py
+++ b/setup.py
@@ -140,7 +140,8 @@ def run(self):
 
     packages = ["splunklib",
                 "splunklib.modularinput",
-                "splunklib.searchcommands"],
+                "splunklib.searchcommands",
+                "splunklib.customrest"],
 
     url="http://github.com/splunk/splunk-sdk-python",
 
diff --git a/splunklib/customrest/__init__.py b/splunklib/customrest/__init__.py
new file mode 100644
index 00000000..4bee5eaa
--- /dev/null
+++ b/splunklib/customrest/__init__.py
@@ -0,0 +1,6 @@
+"""The following imports allow these classes to be imported via
+the splunklib.customrest package like so:
+
+from splunklib.customrest import *
+"""
+from .json import json_handler
\ No newline at end of file
diff --git a/splunklib/customrest/json.py b/splunklib/customrest/json.py
new file mode 100644
index 00000000..c7b86826
--- /dev/null
+++ b/splunklib/customrest/json.py
@@ -0,0 +1,51 @@
+import logging
+import traceback
+import json
+
+from functools import wraps
+
+def json_handler(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        decorated = json_exception_handler(json_payload_extractor(func))
+        return decorated(*args, **kwargs)
+    return wrapper
+
+
+def json_payload_extractor(func):
+    @wraps(func)
+    def wrapper(self, in_string):
+        try:
+            request = json.loads(in_string)
+            kwargs = {'request': request, 'in_string': in_string}
+            if 'payload' in request:
+                # if request contains payload, parse it and add it as payload parameter
+                kwargs['payload'] = json.loads(request['payload'])
+            if 'query' in request:
+                # if request contains query, parse it and add it as query parameter
+                kwargs['query'] = _convert_tuples_to_dict(request['query'])
+            return func(self, **kwargs)
+        except ValueError as e:
+            return {'payload': {'success': 'false', 'result': f'Error parsing JSON: {e}'},
+                    'status': 400
+                    }
+    return wrapper
+
+
+def json_exception_handler(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            logging.error(
+                f'error={repr(e)} traceback={traceback.format_exc()}')
+            return {'payload': {'success': 'false', 'message': f'Error: {repr(e)}'},
+                    'status': 500
+                    }
+    return wrapper
+
+
+def _convert_tuples_to_dict(tuples):
+    return {t[0]: t[1] for t in tuples}
+