merged code to reduce code duplicates

Signed-off-by: Toni Uhlig <matzeton@googlemail.com>
This commit is contained in:
Toni Uhlig
2020-08-18 16:48:01 +02:00
parent 230e96803e
commit e6c2bc6373
3 changed files with 59 additions and 151 deletions

View File

@@ -1,6 +1,8 @@
#!/usr/bin/env python3
import json
import re
import scapy.all
import socket
DEFAULT_HOST = '127.0.0.1'
@@ -53,3 +55,40 @@ class nDPIsrvdSocket:
self.digitlen = 0
return retval
class PcapPacket:
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 JsonParseBytes(json_bytes):
return json.loads(json_bytes.decode('ascii', errors='replace'), strict=False)

View File

@@ -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
sys.path.append(os.path.dirname(sys.argv[0]) + '/../../contrib')
import nDPIsrvd
from nDPIsrvd import nDPIsrvdSocket, TermColor
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"
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()

View File

@@ -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()