decoder.py

from collections import OrderedDict
from coc.message.reader import CoCMessageReader
from coc.message.definitions import CoCMessageDefinitions
import json
import math
from coc.hexdump import hexdump



class CoCMessageDecoder:
    lengthTypes = ['BYTE', 'INT', 'RRSINT32']
    _bitfield = None

    def __init__(self, definitions=None):
        if not definitions:
            self._definitions = CoCMessageDefinitions.read()
        self._definitions = definitions

    def decodeFile(self, filepath):
        messageid = None
        length = None
        version = None
        payload = None
        with open(filepath, 'rb') as fh:
            messageid = int.from_bytes(fh.read(2), byteorder="big")
            length = int.from_bytes(fh.read(3), byteorder="big")
            version = int.from_bytes(fh.read(2), byteorder="big")
            payload = fh.read(length)
        if messageid and length and payload:
            return self.decode(messageid, version, payload)
        return None

    def decode(self, messageid, unknown, payload):
        if messageid in self._definitions:
            reader = CoCMessageReader(messageid, unknown, payload)
            if "fields" in self._definitions[messageid]:
                decoded = {
                    "name": self._definitions[messageid]["name"],
                    "fields": self._decode_fields(reader, self._definitions[messageid]["fields"])
                }
            else:
                decoded = {
                    "name": self._definitions[messageid]["name"]
                }

            unused = reader.read()
            if unused:
                self.dump(decoded)
                raise IndexError("Unused {} bytes in buffer remains".
                                 format(len(unused)))
            return decoded
        else:
            raise KeyError("Message definition missing ({}).".format(messageid))

    def _decode_fields(self, reader, fields):
        decoded = OrderedDict()
        for index, field in enumerate(fields):
            if "name" not in field:
                field["name"] = "unknown_{}".format(str(index).zfill(math.floor(math.log10(len(fields)))))

            # default in lengthType to 4 byte int. Used to count the number of members in an array
            if "lengthType" not in field:
                field["lengthType"] = "INT"

            self._lengthTypeCheck(field["lengthType"])

            value = self._decode_field(reader, field["name"], field["type"], field["lengthType"])
            decoded[field["name"]] = value
        return decoded

    def _lengthTypeCheck(self, lengthType):
        if lengthType not in self.lengthTypes:
            raise ValueError("lengthType {} not supported".format(lengthType))

    def _decode_field(self, reader, name, type, lengthType):
        if not len(reader.peek(1)):
            raise IndexError("Read buffer out of data.")
        if type[:1] == "?":
            if self._bitfield:
                self._bitfield = (self._bitfield << 1) % 16
                if not self._bitfield:
                    return None
            else:
                self._bitfield = 1
            if reader.peek_int(1) & self._bitfield:
                type = type[1:]
                reader.read(1)
                self._bitfield = None
            else:
                return None
        elif self._bitfield:
            reader.read(1)
            self._bitfield = None

        found = type.find("[")
        if found >= 0:
            count = type[found + 1:-1]
            type = type[:found]
            if not count:
                if lengthType == "BYTE":
                    count = reader.read_int(1)
                elif lengthType == "INT":
                    count = reader.read_int(4)
                else:
                    count = reader.read_rrsint32()
            decoded = []
            for i in range(int(count)):
                decoded.append(self._decode_field(reader, "{}[{}]".format(name, i), type, lengthType))
            return decoded
        elif type == "BOOLEAN":
            return bool(reader.read_int(1))
        elif type == "BYTE":
            return reader.read_byte()
        elif type == "SCID":
            return reader.read_scid()
        elif type == "SHORT":
            return reader.read_short()
        elif type == "INT":
            return reader.read_int()
        elif type == "INT32":
            return reader.read_int32()
        elif type == "SINT32":
            return reader.read_sint32()
        elif type == "RRSINT32":
            return reader.read_rrsint32()
        elif type == "RRSLONG":
            return reader.read_rrslong()
        elif type == "LONG":
            return reader.read_long()
        elif type == "STRING":
            return reader.read_string()
        elif type == "ZIP_STRING":
            decoded = reader.read_zstring()
            if not decoded:
                return decoded
            try:
                decoded = decoded.decode()
            except UnicodeDecodeError:
                raise ValueError("Failed to decode JSON.")
            else:
                return json.loads(decoded)
        elif type == "IGNORE":
            # consume all the trailing data and ignore
            ignore_count = 0
            while len(reader.peek(1)) > 0:
                ignore_count += 1
                reader.read_byte()
            return "{} bytes ignored".format(ignore_count)
        elif type in self._definitions["component"]:
            decoded = self._decode_fields(reader, self._definitions["component"][type]["fields"])
            if "extensions" in self._definitions["component"][type]:
                if not decoded["id"] in self._definitions["component"][type]["extensions"]:
                    raise NotImplementedError("{}(id={}) has not yet been implemented.".format(type, decoded["id"]))
                decoded["payload"] = self._decode_fields(reader, self._definitions["component"][type]["extensions"][decoded["id"]]["fields"])
            self.dump(decoded)
            return decoded
        else:
            raise NotImplementedError("{} has not yet been implemented.".format(type))

    def dump(self, decoded, hide_unknown=False):
        if "name" not in decoded:
            decoded["name"] = "unknownName"

        if "fields" in decoded:
            print("{}: {}".format(decoded["name"], json.dumps(self.stringify(decoded["fields"], hide_unknown), indent=2)))
        else:
            print("{}: {{}}".format(decoded["name"]))

    def stringify(self, decoded, hide_unknown=False):
        stringified = type(decoded)()
        if type(decoded) is list:
            keys = range(len(decoded))
        else:
            keys = decoded.keys()
        for key in keys:
            if hide_unknown and key[:len("unknown_")] == "unknown_":
                continue
            value = decoded[key]
            if type(value) is bytes:
                try:
                    str = value.decode()
                except UnicodeDecodeError:
                    str = value.hex()
            elif type(value) in {dict, list, OrderedDict}:
                str = self.stringify(value)
            else:
                str = value
            if type(stringified) is list:
                stringified.append(str)
            else:
                stringified[key] = str
        return stringified