import struct
import socket
import humanize

templatecache = {}


def unpack(buf):
    (version, count) = struct.unpack('!HH', buf[0:4])
    if version != 9:
        raise Exception(f"A version {version} packet found.")
    bufpoint = 20
    offset = 0
    totalbytes = 0
    while bufpoint < len(buf):
        (flowsetid, length) = struct.unpack('!HH', buf[bufpoint:bufpoint + 4])
        if not length:
            return
        offset = bufpoint + 4

        if flowsetid > 255:  # We have a data flowset
            while offset < length:
                if flowsetid not in templatecache:
                    offset += length  # Skip until such a template is defined
                else:
                    for fieldid, fieldlength in templatecache.get(flowsetid, []):
                        bufbuf = buf[offset:offset + fieldlength]
                        if not len(bufbuf):
                            break
                        offset += fieldlength
                        if fieldid not in (1, 23):  # ignore everything other than IN_BYTES and OUT_BYTES
                            continue

                        if fieldlength == 1:
                            x = struct.unpack("!B", bufbuf)[0]
                        elif fieldlength == 2:
                            x = struct.unpack("!H", bufbuf)[0]
                        elif fieldlength == 4:
                            x = struct.unpack("!I", bufbuf)[0]
                        elif fieldlength == 8:
                            x = struct.unpack("!Q", bufbuf)[0]
                        else:
                            raise Exception(f"Unsupported fieldlength: {fieldlength}")

                        totalbytes += x

        elif flowsetid <= 255:  # We have a template flowset
            offset = bufpoint + 4
            while offset < length:
                (templateid, fieldcount) = struct.unpack('!HH', buf[offset:offset + 4])
                offset += 4
                for field in range(fieldcount):
                    (fieldid, fieldlength) = struct.unpack("!HH", buf[offset:offset + 4])
                    templatecache.setdefault(templateid, [])
                    templatecache[templateid].append([fieldid, fieldlength])
                    offset += 4
        bufpoint += length
    return totalbytes


if __name__ == "__main__":
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.bind(("0.0.0.0", 9995))
    totalbytes = 0
    while(True):
        buf, addr = s.recvfrom(10000)
        totalbytes += unpack(buf)
        print(f"Total Bytes: {humanize.naturalsize(totalbytes)}")