Enable custom JSON filter expressions for Python scripts.

Signed-off-by: Toni Uhlig <matzeton@googlemail.com>

1 files changed, 68 insertions, 10 deletions

diff --git a/dependencies/nDPIsrvd.py b/dependencies/nDPIsrvd.py
index 166af5596..bd15a6bde 100644
--- a/dependencies/nDPIsrvd.py
+++ b/dependencies/nDPIsrvd.py
@@ -342,11 +342,43 @@ class SocketTimeout(nDPIsrvdException):
     def __str__(self):
         return 'Socket timeout.'
 
+class JsonFilter():
+    def __init__(self, filter_string):
+        self.filter_string = filter_string
+        self.filter = compile(filter_string, '<string>', 'eval')
+    def evaluate(self, json_dict):
+        if type(json_dict) is not dict:
+            raise nDPIsrvdException('Could not evaluate JSON Filter: expected dictionary, got {}'.format(type(json_dict)))
+        return eval(self.filter, {'json_dict': json_dict})
+
 class nDPIsrvdSocket:
     def __init__(self):
         self.sock_family = None
         self.flow_mgr = FlowManager()
         self.received_bytes = 0
+        self.json_filter = list()
+
+    def addFilter(self, filter_str):
+        self.json_filter.append(JsonFilter(filter_str))
+
+    def evalFilters(self, json_dict):
+        for jf in self.json_filter:
+            try:
+                json_filter_retval = jf.evaluate(json_dict)
+            except Exception as err:
+                print()
+                sys.stderr.write('Error while evaluating expression "{}"\n'.format(jf.filter_string))
+                raise err
+
+            if type(json_filter_retval) != bool:
+                print()
+                sys.stderr.write('Error while evaluating expression "{}"\n'.format(jf.filter_string))
+                raise nDPIsrvdException('JSON Filter returned an invalid type: expected bool, got {}'.format(type(json_filter_retval)))
+
+            if json_filter_retval is False:
+                return False
+
+        return True
 
     def connect(self, addr):
         if type(addr) is tuple:
@@ -363,6 +395,7 @@ class nDPIsrvdSocket:
         self.digitlen = 0
         self.lines = []
         self.failed_lines = []
+        self.filtered_lines = 0
 
     def timeout(self, timeout):
         self.sock.settimeout(timeout)
@@ -432,19 +465,24 @@ class nDPIsrvdSocket:
                 retval = False
                 continue
 
-            try:
-                if callback_json(json_dict, instance, self.flow_mgr.getFlow(instance, json_dict), global_user_data) is not True:
-                    self.failed_lines += [received_line]
-                    retval = False
-            except Exception as e:
-                    self.failed_lines += [received_line]
-                    self.lines = self.lines[1:]
-                    raise(e)
+            current_flow = self.flow_mgr.getFlow(instance, json_dict)
+            filter_eval = self.evalFilters(json_dict)
+            if filter_eval is True:
+                try:
+                    if callback_json(json_dict, instance, current_flow, global_user_data) is not True:
+                        self.failed_lines += [received_line]
+                        retval = False
+                except Exception as e:
+                        self.failed_lines += [received_line]
+                        self.lines = self.lines[1:]
+                        raise(e)
+            else:
+                self.filtered_lines += 1
 
             for _, flow in self.flow_mgr.getFlowsToCleanup(instance, json_dict).items():
                 if callback_flow_cleanup is None:
                     pass
-                elif callback_flow_cleanup(instance, flow, global_user_data) is not True:
+                elif filter_eval is True and callback_flow_cleanup(instance, flow, global_user_data) is not True:
                     self.failed_lines += [received_line]
                     self.lines = self.lines[1:]
                     retval = False
@@ -477,12 +515,15 @@ class nDPIsrvdSocket:
             raise nDPIsrvdException('Failed lines > 0: {}'.format(len(self.failed_lines)))
         return self.flow_mgr.verifyFlows()
 
-def defaultArgumentParser(desc='nDPIsrvd Python Interface',
+def defaultArgumentParser(desc='nDPIsrvd Python Interface', enable_json_filter=False,
                           help_formatter=argparse.ArgumentDefaultsHelpFormatter):
     parser = argparse.ArgumentParser(description=desc, formatter_class=help_formatter)
     parser.add_argument('--host', type=str, help='nDPIsrvd host IP')
     parser.add_argument('--port', type=int, default=DEFAULT_PORT, help='nDPIsrvd TCP port')
     parser.add_argument('--unix', type=str, help='nDPIsrvd unix socket path')
+    if enable_json_filter is True:
+        parser.add_argument('--filter', type=str, action='append',
+                            help='Set a filter string which if evaluates to True will invoke the JSON callback.')
     return parser
 
 def toSeconds(usec):
@@ -515,6 +556,23 @@ def validateAddress(args):
 
     return address
 
+def prepareJsonFilter(args, nsock):
+    # HowTo use JSON Filters:
+    # Add `--filter [FILTER_STRING]` to the Python scripts that support JSON filtering.
+    # Examples:
+    #   ./examples/py-json-stdout/json-stdout.py --filter '"ndpi" in json_dict and "proto" in json_dict["ndpi"]'
+    #   The command above will print only JSONs that have the subobjects json_dict["ndpi"] and json_dict["ndpi"]["proto"] available.
+    #   ./examples/py-flow-info/flow-info.py --filter 'json_dict["source"] == "eth0"' --filter '"flow_event_name" in json_dict and json_dict["flow_event_name"] == "analyse"'
+    #   Multiple JSON filter will be ANDed together.
+    # Note: You may *only* use the global "json_dict" in your expressions.
+    try:
+        json_filter = args.filter
+        if json_filter is not None:
+            for jf in json_filter:
+                nsock.addFilter(jf)
+    except AttributeError:
+        pass
+
 global schema
 schema = {'packet_event_schema' : None, 'error_event_schema' : None, 'daemon_event_schema' : None, 'flow_event_schema' : None}