diff --git a/bemani/protocol/binary.py b/bemani/protocol/binary.py index 49d181d..d941b1a 100644 --- a/bemani/protocol/binary.py +++ b/bemani/protocol/binary.py @@ -248,7 +248,7 @@ class BinaryDecoder: A class capable of taking a binary blob and decoding it to a Node tree. """ - def __init__(self, data: bytes, encoding: str) -> None: + def __init__(self, data: bytes, encoding: str, compressed: bool) -> None: """ Initialize the object. @@ -259,6 +259,7 @@ class BinaryDecoder: """ self.stream = InputStream(data) self.encoding = encoding + self.compressed = compressed self.executed = False def __read_node_name(self) -> str: @@ -272,6 +273,31 @@ class BinaryDecoder: length = self.stream.read_int() if length is None: raise BinaryEncodingException("Ran out of data when attempting to read node name length!") + + if not self.compressed: + if length < 0x40: + raise BinaryEncodingException("Node name length under decompressed minimum") + elif length < 0x80: + length -= 0x3f + else: + length_ex = self.stream.read_int() + if length_ex is None: + raise BinaryEncodingException("Ran out of data when attempting to read node name length!") + length = (length << 8) | length_ex + length -= 0x7fbf + + if length > BinaryEncoding.NAME_MAX_DECOMPRESSED: + raise BinaryEncodingException("Node name length over decompressed limit") + + name = self.stream.read_blob(length) + if name is None: + raise BinaryEncodingException("Ran out of data when attempting to read node name!") + + return name.decode(self.encoding) + + if length > BinaryEncoding.NAME_MAX_COMPRESSED: + raise BinaryEncodingException("Node name length over compressed limit") + binary_length = int(((length * 6) + 7) / 8) def int_to_bin(integer: int) -> str: @@ -716,6 +742,9 @@ class BinaryEncoding: DECOMPRESSED_WITH_DATA: Final[int] = 0x45 DECOMPRESSED_WITHOUT_DATA: Final[int] = 0x46 + NAME_MAX_COMPRESSED: Final[int] = 0x24 + NAME_MAX_DECOMPRESSED: Final[int] = 0x1000 + # The string values should match the constants in EAmuseProtocol. # I have no better way to link these than to write this comment, # as otherwise we would have a circular dependency. @@ -732,6 +761,7 @@ class BinaryEncoding: Initialize the encoding object. """ self.encoding: Optional[str] = None + self.compressed: bool = True def __sanitize_encoding(self, enc: str) -> str: """ @@ -770,8 +800,13 @@ class BinaryEncoding: return None if ((~encoding_raw) & 0xFF) != encoding_swapped: return None - if contents not in [BinaryEncoding.COMPRESSED_WITH_DATA, BinaryEncoding.COMPRESSED_WITHOUT_DATA]: - # We don't support uncompressed data. + + self.compressed = contents in [ + BinaryEncoding.COMPRESSED_WITH_DATA, BinaryEncoding.COMPRESSED_WITHOUT_DATA + ] + if not self.compressed and contents not in [ + BinaryEncoding.DECOMPRESSED_WITH_DATA, BinaryEncoding.DECOMPRESSED_WITHOUT_DATA + ]: return None encoding = BinaryEncoding.ENCODINGS.get(encoding_raw) @@ -779,7 +814,9 @@ class BinaryEncoding: if encoding is not None: self.encoding = encoding try: - decoder = BinaryDecoder(data[4:], self.__sanitize_encoding(encoding)) + decoder = BinaryDecoder( + data[4:], self.__sanitize_encoding(encoding), self.compressed + ) return decoder.get_tree() except BinaryEncodingException: if skip_on_exceptions: