#!/usr/bin/python3

import struct
import dpkt
import bitstring
import sys

LBP16_PORT = 27181

class Command:
    def __init__(self, buffer=b'\x00\x00'):
        if len(buffer) < 2:
            raise EOFError("Insufficient data to parse command")

        command_int, = struct.unpack("<H", buffer[:2])

        bits = bitstring.BitArray(uint=command_int, length=16)
        # Protocol fields
        self.write_bit = bits[0:1].bool
        self.address_bit = bits[1:2].bool
        self.info_bit = bits[2:3].bool
        self.memoryspace = bits[3:6].uint
        self.transfersize = bits[6:8].uint
        self.increment_bit = bits[8:9].bool
        self.transfercount = bits[9:].uint
        # Helper fields
        self.transferbytes = 2**self.transfersize * self.transfercount
        self.commandsize = 2 + int(self.address_bit)*2 + int(self.write_bit)*self.transferbytes

        consumed_bytes = 2
        if self.address_bit:
            if len(buffer) - consumed_bytes < 2:
                raise EOFError("Insufficient data to parse address")
            self.address, = struct.unpack("<H", buffer[consumed_bytes:consumed_bytes+2])
            consumed_bytes = consumed_bytes + 2
        if self.write_bit:
            if len(buffer) - consumed_bytes < self.transferbytes:
                raise EOFError("Insufficient data to write data")
            self.data = buffer[consumed_bytes:consumed_bytes+self.transferbytes]
            consumed_bytes = consumed_bytes + self.transferbytes
        else:
            self.data = None
        
        self.comment = None

    def tostring(self):
        parts = []
        if self.write_bit:
            parts.append("WRITE")
        else:
            parts.append("READ ")
        if self.info_bit:
            parts.append("INFO  ")
        else:
            parts.append("MEMORY")
        parts.append("{}".format(self.memoryspace))
        if self.address_bit:
            parts.append("ADDRESS 0x{:04x}".format(self.address))
        if self.increment_bit:
            parts.append("INC  ")
        else:
            parts.append("NOINC")
        
        parts.append("SIZE {}x{}".format(2**self.transfersize, self.transfercount))

        if self.data is not None:
            parts.append("DATA 0x{}".format(self.data.hex()))
        
        if self.comment is not None:
            parts.append(self.comment)

        return " ".join(parts)

   
def lbp16_handle_transaction(request_packets, response_packets):
    (response_packet_num, response_udp) = response_packets.pop(0)
    response_offset = 0
    for (packet_num, request_udp) in request_packets:
        request_offset = 0
        while len(request_udp.data) - request_offset > 2:
            cmd = Command(request_udp.data[request_offset:])
            request_offset = request_offset + cmd.commandsize
            if cmd.write_bit == False:
                # Get response data
                while len(response_udp) - response_offset < cmd.transferbytes:
                    if len(response_udp) - response_offset != 0:
                        raise ValueError("Fragmented packet response length")
                    (response_packet_num, response_udp) = response_packets.pop(0)
                    response_offset = 0
                
                data = response_udp.data[response_offset:response_offset+cmd.transferbytes]
                cmd.data = data
                cmd.comment = "from (#{} offset {})".format(response_packet_num, response_offset)
                response_offset = response_offset+cmd.transferbytes
                
            print("#{:6d} {}".format(packet_num, cmd.tostring()))


if __name__ == "__main__":
    filename = "capture-onlyfew2.pcap"
    try:
        filename = sys.argv[1]
        print(filename)
    except:
        pass

    pcapfile = open(filename, "rb")
    pcap = dpkt.pcap.Reader(pcapfile)

    request_bytes = b''
    response_bytes = b''
    request_packets = []
    response_packets = []
    packet_num = 0
    for _, packet_raw in pcap:
        packet_num = packet_num + 1

        # Discard non UDP packets
        packet_eth = dpkt.ethernet.Ethernet(packet_raw)
        if not isinstance(packet_eth.data, dpkt.ip.IP):
            continue
        packet_ip = packet_eth.data
        if not isinstance(packet_ip.data, dpkt.udp.UDP):
            continue
        packet_udp = packet_ip.data


        if packet_udp.dport == LBP16_PORT:
            if len(response_packets) > 0:
                lbp16_handle_transaction(request_packets, response_packets)
                request_packets = []
                response_packets = []
            request_packets.append((packet_num, packet_udp))
        elif packet_udp.sport == LBP16_PORT:
            response_packets.append((packet_num, packet_udp))