#!/usr/bin/python3
import binascii
import struct
import sys
import argparse
import subprocess
import io
import time
import datetime


"""
@copyright
(c) Copyright ESO 2024
All Rights Reserved
ESO (eso.org) is an Intergovernmental Organisation, and therefore special legal conditions apply.

 Example runs:

 ./mudcat.py --file m5_dpdk.pcap  --cat
 ./mudcat.py --file m5_dpdk.pcap  --report
 ./mudcat.py --interface lo  --tim 20 --report
 ./mudcat.py --help
 ./mudcat.py --interface lo  --time 20 --cat

Multicast groups are listed here:
https://europeansouthernobservatory.sharepoint.com/sites/ELT_Control/_layouts/OneNote.aspx?id=%2Fsites%2FELT_Control%2FShared%20Documents%2FNotebooks%2FELT_Control_ECM_Notebook&wd=target%28System.one%7C013232E1-2639-492B-BFCB-83417BED70B0%2FNapatech%20sniffer%20mappings%7CEEEC23E8-ED9F-4A32-9F03-9D2950AD05A5%2F%29
onenote:https://europeansouthernobservatory.sharepoint.com/sites/ELT_Control/Shared%20Documents/Notebooks/ELT_Control_ECM_Notebook/System.one#Napatech%20sniffer%20mappings&section-id={013232E1-2639-492B-BFCB-83417BED70B0}&page-id={EEEC23E8-ED9F-4A32-9F03-9D2950AD05A5}&end

"""

# MUDPI dissection from UPD payload
def print_mudpi_packet(topicid, sampleid, ts, frameid, numframes, mpaylen, mpayload):
    print(str(topicid).rjust(4)+"\t"+str(sampleid).rjust(10)+"\t"+f'{ts:.9f}'+"\t"+str(frameid)+"/"+str(numframes)+"\t"+str(mpaylen).rjust(4)+"\t"+str(mpayload))

#
# print the summary dictionary returned from read_and_cat()
#
def print_occurences(summary):
    print('src addr\tdst addr\t  port\ttopic id\t   count')
    for key, value in summary.items():
        print(value[0]+'\t'+value[1]+'\t'+str(value[2]).rjust(6)+'\t'+str(value[3]).rjust(8)+'\t'+str(value[4]).rjust(8))


#
# runs tshark and processes the output, including introspection of the MUDPI frames
# tshark -r myfile.pcap -l -T fields -e ip.src -e ip.dst -e udp.dstport -e ip.len -e ip.proto -e udp.payload -e frame.time
# if gather_occurences is True, a dictionary is returned, the key are hashkeys, the values are tuples of the following
# ordering: (srcaddr, destaddr, port, topicid, occurrences)
#
def read_and_cat(dev_name, dev_type='file', run_time=5, cat_packets=True, cat_payload=True, gather_occurences=True):

    occurences = {}
    ifflag='-r'
    if dev_type == 'file':
        ifflag='-r'
    elif dev_type == 'nic':
        ifflag='-i'
    print('tshark', ifflag, dev_name, '-l', '-T', 'fields', '-e' 'ip.src', '-e', 'ip.dst', '-e', 'udp.dstport', '-e', 'ip.len', '-e', 'ip.proto', '-e', 'udp.payload', '-e', 'frame.time')
    tshark_process = subprocess.Popen(['tshark', ifflag, dev_name, '-l', '-T', 'fields', '-e' 'ip.src', '-e', 'ip.dst', '-e', 'udp.dstport', '-e', 'ip.len', '-e', 'ip.proto', '-e', 'udp.payload', '-e', 'frame.time'], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
    #tshark_process.communicate()
    line_num=0
    #for line in sys.stdin:
    #for line in iter(tshark_process.stdout.readline, ""):
    end_at = time.time()+run_time
    if (cat_packets):
        print('Row\tsrc addr\tdst addr\tport\tudp sz\tarrival time\t\ttopicid\t sample id\tmudpi timestamp\t\tframe\tmudpi sz\tmudpi payload')
    for line in io.TextIOWrapper(tshark_process.stdout, encoding="utf-8"):
        if run_time > 0 and time.time() > end_at:
            break
        line_num+=1
        # "172.31.239.121  239.128.0.51    30004   86      17 0000000e0000000000000001000063b041d9b80db121374d0001000100180000000000000000000000000000000000000000000000000000b9d0"
        # print(str(line_num)+": "+line)
        fields=line.split('\t')
        if len(fields) != 7:
            # print(str(line_num)+"\tPacket received is not a UDP Packet")
            continue

        srcaddr=fields[0]
        destaddr=fields[1]
        port=fields[2]
        payload_len=fields[3]
        ip_proto_num = fields[4]
        frame_time = fields[6].strip()
        # frame_time format expected to be: "Nov  5, 2024 13:20:06.117717199 UTC"
        # we trim the last 7 characters (UTC and the nanoseconds component, to get microseconds)
        frame_time_seconds = datetime.datetime.strptime(frame_time[:-7], "%b %d, %Y %H:%M:%S.%f")
        if ip_proto_num != '17':
            # this is not a UDP packet
            print('Non-UDP packet: '+str(line_num)+": IP proto: "+line)
            continue
        data_bytes = bytearray.fromhex(fields[5])

        #print(binascii.hexlify(data_bytes))
        if len(data_bytes) < 30:
            print('Unexpected/Fragment warning?: '+str(line_num)+": "+line)
            continue
        topicid=int.from_bytes(data_bytes[0:4], byteorder='big', signed=False)
        sampleid=int.from_bytes(data_bytes[12:16], byteorder='big', signed=False)
        ts=struct.unpack('>d', data_bytes[16:24])[0]
        frameid=int.from_bytes(data_bytes[24:26], byteorder='big', signed=False)
        numframes=int.from_bytes(data_bytes[26:28], byteorder='big', signed=False)
        mpaylen=int.from_bytes(data_bytes[28:30], byteorder='big', signed=False)
        if cat_payload:
            mpayload=data_bytes[32:32+mpaylen]
            mpayload = ''.join(format(x, '02x') for x in mpayload)
        else:
            mpayload = "-"
        if (cat_packets):
            print(str(line_num)+"\t"+srcaddr+"\t"+destaddr+"\t"+port.rjust(6)+"\t"+payload_len.rjust(4)+"\t"+'{:17.6f}'.format(frame_time_seconds.timestamp())+"\t", end='')
            print_mudpi_packet(topicid, sampleid, ts, frameid, numframes, mpaylen, mpayload)
        if (gather_occurences):
            hash_val = hash((srcaddr, destaddr, port, topicid))
            packet_tuple = occurences.get(hash_val)
            if packet_tuple == None:
                # create a new record
                packet_tuple = (srcaddr, destaddr, port, topicid, 1)
            else:
                packet_tuple = (srcaddr, destaddr, port, topicid, packet_tuple[4]+1)
            occurences[hash_val] = packet_tuple
    return occurences

def main():
    """Main function implementation."""
    parser = argparse.ArgumentParser(
        description="Report or dissect MUDPI data",
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument(
        "-r", "--report", dest='repflag', action='store_true', required=False, help="Print summary of tshark output (default if file is specified)"
    )
    parser.add_argument(
        "-c", "--cat", dest='catflag', action='store_true', required=False, help="Dissect each line of tshark output (default if a device is specified)"
    )
    parser.add_argument(
        "-p", "--payload", dest='catpayflag', action='store_true', required=False, help="Include hex listing of MUDPI payload (only in cat mode)"
    )
    parser.add_argument(
        "-i", "--interface", dest='iname', required=False, help="Network device for tshark"
    )
    parser.add_argument(
        "-f", "--file", dest='fname', required=False, help="PCAP format file for tshark"
    )
    parser.add_argument(
        "-t", "--time", dest='runtime', required=False, help="Length of time to run in -c mode, default 5s"
    )

    args = parser.parse_args()
    interface_name = args.iname
    file_name = args.fname
    if args.runtime:
        run_time = float(args.runtime)
    else:
        run_time = 5.0
    catout=False
    mpayout=False
    reportout=False
    if not interface_name and not file_name:
        print("One of -i or -f must be specified")
        return -1
    if file_name:
        run_time = 0  # no max run time if reading a file

    if args.repflag:
        reportout = True
    if args.catflag:
        catout = True
    if args.catpayflag:
        mpayout = True
    if not reportout and not catout:
        if file_name:
            reportout = True
        if interface_name:
            catout = True
    if not catout and mpayout:
        print("Cannot print MUDPI payload in report mode")
        return -1

    # print(interface_name, file_name, reportout, catout)
    summary = {}
    if interface_name:
        summary = read_and_cat(dev_name=interface_name, dev_type="nic", run_time=run_time, cat_packets=catout, cat_payload=mpayout, gather_occurences=reportout)
    else:
        summary = read_and_cat(dev_name=file_name, dev_type="file", run_time=0, cat_packets=catout, cat_payload=mpayout, gather_occurences=reportout)
    if reportout:
        print_occurences(summary)


if __name__ == "__main__":
    sys.exit(main())

