merged code to reduce code duplicates

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

2 files changed, 22 insertions, 153 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()
diff --git a/examples/py-flow-undetected-to-pcap/flow-undetected-to-pcap.py b/examples/py-flow-undetected-to-pcap/flow-undetected-to-pcap.py
index 9af862144..1a87314ad 100755
--- a/examples/py-flow-undetected-to-pcap/flow-undetected-to-pcap.py
+++ b/examples/py-flow-undetected-to-pcap/flow-undetected-to-pcap.py
@@ -1,98 +1,18 @@
 #!/usr/bin/env python3
 
 import base64
-import json
-import re
+import os
 import sys
-import socket
-import scapy.all
 
-HOST = '127.0.0.1'
-PORT = 7000
-NETWORK_BUFFER_MIN_SIZE = 5
-NETWORK_BUFFER_MAX_SIZE = 8192
+sys.path.append(os.path.dirname(sys.argv[0]) + '/../../contrib')
+import nDPIsrvd
+from nDPIsrvd import TermColor, nDPIsrvdSocket, PcapPacket
 
 FLOWS = dict()
 
-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 Flow:
-    def __init__(self, flow_id=-1):
-        self.pktdump = None
-        self.was_dumped = False
-        self.was_detected = False
-        self.flow_id = flow_id
-        self.packets = []
-
-    def addPacket(self, pkt):
-        self.packets += [pkt]
-
-    def detected(self):
-        self.was_detected = True
-
-    def fin(self, filename_suffix):
-        if self.was_dumped is True:
-            return
-        if self.was_detected is True:
-            return
-
-        if self.pktdump is None:
-            if self.flow_id == -1:
-                self.pktdump = scapy.all.PcapWriter('packet-{}.pcap'.format(filename_suffix),
-                                                    append=True, sync=True)
-            else:
-                self.pktdump = scapy.all.PcapWriter('flow-{}-{}.pcap'.format(filename_suffix, self.flow_id),
-                                                    append=False, sync=True)
-
-        for packet in self.packets:
-            self.pktdump.write(scapy.all.Raw(packet))
-
-        self.pktdump.close()
-        self.was_dumped = True
-
 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])
 
     global FLOWS
 
@@ -106,7 +26,7 @@ def parse_json_str(json_str):
 
         if event == 'new':
             print('New flow with id {}.'.format(flow_id))
-            FLOWS[flow_id] = Flow(flow_id)
+            FLOWS[flow_id] = PcapPacket(flow_id)
         elif flow_id not in FLOWS:
             print('Ignore flow event with id {} as we did not get any flow-new event.'.format(flow_id))
             return
@@ -143,13 +63,13 @@ def parse_json_str(json_str):
 
         if j['packet_event_name'] == 'packet':
 
-            flow = Flow()
+            flow = PcapPacket()
             flow.addPacket(buffer_decoded)
 
 
 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]))
@@ -158,7 +78,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()