merged code to reduce code duplicates

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

1 files changed, 12 insertions, 63 deletions

diff --git a/examples/py-flow-info/flow-info.py b/examples/py-flow-info/flow-info.py
index 96f2b0346..9704c7c7d 100755
--- a/examples/py-flow-info/flow-info.py
+++ b/examples/py-flow-info/flow-info.py
@@ -1,67 +1,16 @@
 #!/usr/bin/env python3
 
-import json
-import re
+import os
 import sys
-import socket
-
-HOST = '127.0.0.1'
-PORT = 7000
-NETWORK_BUFFER_MIN_SIZE = 5
-NETWORK_BUFFER_MAX_SIZE = 8192
-
-class nDPIsrvdSocket:
-    def __init__(self, sock=None):
-        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-
-    def connect(self, host, port):
-        self.sock.connect((host, port))
-        self.buffer = bytes()
-        self.msglen = 0
-        self.digitlen = 0
-
-    def receive(self):
-        recvd = self.sock.recv(NETWORK_BUFFER_MAX_SIZE - len(self.buffer))
-
-        if len(recvd) == 0:
-            raise RuntimeError('socket connection broken')
-        self.buffer += recvd
-
-        retval = []
-        while self.msglen + self.digitlen < len(self.buffer):
-
-            if self.msglen == 0:
-                starts_with_digits = re.match(r'(^\d+){', self.buffer[:NETWORK_BUFFER_MIN_SIZE].decode(errors='strict'))
-                if starts_with_digits is None:
-                    if len(self.buffer) < NETWORK_BUFFER_MIN_SIZE:
-                        break
-                    raise RuntimeError('Invalid packet received: {}'.format(self.buffer))
-                self.msglen = int(starts_with_digits[1])
-                self.digitlen = len(starts_with_digits[1])
-
-            if len(self.buffer) >= self.msglen + self.digitlen:
-                recvd = self.buffer[self.digitlen:self.msglen + self.digitlen]
-                self.buffer = self.buffer[self.msglen + self.digitlen:]
-                retval += [(recvd,self.msglen,self.digitlen)]
-
-                self.msglen = 0
-                self.digitlen = 0
-
-        return retval
-
-class TermColor:
-    WARNING = '\033[93m'
-    FAIL = '\033[91m'
-    BOLD = '\033[1m'
-    END = '\033[0m'
-    BLINK = "\x1b[5m"
+
+sys.path.append(os.path.dirname(sys.argv[0]) + '/../../contrib')
+import nDPIsrvd
+from nDPIsrvd import nDPIsrvdSocket, TermColor
+
 
 def parse_json_str(json_str):
 
-    try:
-        j = json.loads(json_str[0])
-    except json.decoder.JSONDecodeError as exc:
-        raise RuntimeError('JSON Exception: {}\n\nJSON String: {}\n'.format(str(exc), str(json_str)))
+    j = nDPIsrvd.JsonParseBytes(json_str[0])
 
     if 'flow_event_name' in j:
         event = j['flow_event_name'].lower()
@@ -101,7 +50,7 @@ def parse_json_str(json_str):
                     ndpi_frisk[:-2])
 
         if j['l3_proto'] == 'ip4':
-            print('{:>14}: [{:.>8}] [{}][{:.>5}] [{:.>15}]{} -> [{:.>15}]{} {}'.format(event_str,
+            print('{:>14}: [{:.>6}] [{}][{:.>5}] [{:.>15}]{} -> [{:.>15}]{} {}'.format(event_str,
                   j['flow_id'], j['l3_proto'], j['l4_proto'],
                   j['src_ip'].lower(),
                   '[{:.>5}]'.format(j['src_port']) if 'src_port' in j else '',
@@ -109,7 +58,7 @@ def parse_json_str(json_str):
                   '[{:.>5}]'.format(j['dst_port']) if 'dst_port' in j else '',
                   ndpi_proto_categ))
         elif j['l3_proto'] == 'ip6':
-            print('{:>14}: [{:.>8}] [{}][{:.>5}] [{:.>39}]{} -> [{:.>39}]{} {}'.format(event_str,
+            print('{:>14}: [{:.>6}] [{}][{:.>5}] [{:.>39}]{} -> [{:.>39}]{} {}'.format(event_str,
                   j['flow_id'], j['l3_proto'], j['l4_proto'],
                   j['src_ip'].lower(),
                   '[{:.>5}]'.format(j['src_port']) if 'src_port' in j else '',
@@ -124,8 +73,8 @@ def parse_json_str(json_str):
 
 
 if __name__ == '__main__':
-    host = HOST
-    port = PORT
+    host = nDPIsrvd.DEFAULT_HOST
+    port = nDPIsrvd.DEFAULT_PORT
 
     if len(sys.argv) == 1:
         sys.stderr.write('usage: {} [host] [port]\n'.format(sys.argv[0]))
@@ -134,7 +83,7 @@ if __name__ == '__main__':
     if len(sys.argv) > 2:
         port = int(sys.argv[2])
 
-    sys.stderr.write('Recv buffer size: {}\n'.format(NETWORK_BUFFER_MAX_SIZE))
+    sys.stderr.write('Recv buffer size: {}\n'.format(nDPIsrvd.NETWORK_BUFFER_MAX_SIZE))
     sys.stderr.write('Connecting to {}:{} ..\n'.format(host, port))
 
     nsock = nDPIsrvdSocket()