a3f247f422
- Remove redundant attribute creation function. - Remove redundant len field. - Add the rest of the missing types to the constructor array. - Reorganize a bit.
796 lines
29 KiB
Python
796 lines
29 KiB
Python
import struct
|
|
from typing import Optional, List, Dict, Any
|
|
|
|
from bemani.protocol.stream import InputStream, OutputStream
|
|
from bemani.protocol.node import Node
|
|
|
|
|
|
class BinaryEncodingException(Exception):
|
|
"""
|
|
Generic exception to be thrown when we encounter an issue decoding a binary stream
|
|
"""
|
|
|
|
|
|
class PackedOrdering:
|
|
"""
|
|
A class that helps us encapsulate Konami's batshit backtracking hole-fill algorithm.
|
|
Everything is aligned on a boundary appropriate for its data size. Strings and arrays are
|
|
forced to be aligned to a 4 byte boundary on account of having an integer length field.
|
|
All of these are padded to 4 bytes in terms of the room they take up in the stream.
|
|
For the things that are 2 byte or 1 byte aligned, we end up sticking them after each other
|
|
in 4 byte increments. That is, to say, if we have a unsigned byte to pack, we reserve 4 bytes
|
|
and stick it in the first byte slot, and if up to three additional bytes come in we will pack
|
|
them after this in sequential order. It would make sense to not pad out strings and arrays and
|
|
store bytes/shorts in these unused locations, but that's not what actually happens. Also note
|
|
that we will never pack bytes after a short or vice versa, even if there is room. This also explains
|
|
the bizarre behavior of not using spare bytes after strings or arrays. I'll emphasize again:
|
|
everything is stored aligned, and in a 4 byte chunk, only similarly-sized objects can be packed. If
|
|
this 4 byte chunk is already partially occupied, we can only add another thing to it if 1) the
|
|
item being added is the same size as the object that exists and 2) the object can be added with
|
|
the correct alignment.
|
|
|
|
A simple example:
|
|
[1: byte] [2: byte] [3: integer]
|
|
Packing would look like this (assuming all locations are a byte):
|
|
1 2 0 0 3 3 3 3
|
|
|
|
An example:
|
|
[1: byte] [2: string, length 3] [3: short] [4: byte]
|
|
Packing would look like this (assuming all locations are a byte):
|
|
1 4 0 0 2 2 2 2 2 2 2 0 3 3 0 0
|
|
"""
|
|
|
|
def __init__(self, size: int, allow_expansion: bool=False) -> None:
|
|
"""
|
|
Initialize with a known size. If this is to be used to create a packing instead of deduce
|
|
a packing, then allow_expansion should be set to true and new holes will be created when
|
|
needed. If this is to be used for decoding a current packing, allow_expansion should be set
|
|
to False to ensure we don't choose locations outside the buffer.
|
|
|
|
Parameters:
|
|
size - Number of bytes to work with as an integer
|
|
allow_expansion - Boolean describing whether to add to the end of the order when needed
|
|
"""
|
|
self.order: List[int] = []
|
|
self.expand = allow_expansion
|
|
|
|
for i in range(size):
|
|
self.order.append(None)
|
|
self.__orderlen = size
|
|
self.__lastbyte = 0
|
|
self.__lastshort = 0
|
|
self.__lastint = 0
|
|
|
|
def __append_empty(self) -> None:
|
|
self.order.append(None)
|
|
self.__orderlen = self.__orderlen + 1
|
|
|
|
def mark_used(self, size: int, offset: int, round_to: int=1) -> None:
|
|
"""
|
|
Mark size bytes at offset as being used. If needed, round to the nearest byte/half/integer.
|
|
|
|
Parameters:
|
|
size - Number of bytes to mark
|
|
offset - Offset into binary chunk to start marking
|
|
round_to - Optional integer specifying how many bytes to round to. Valid values are 1, 2 and 4
|
|
"""
|
|
# Round to nearest value if needed
|
|
while (size & (round_to - 1)) != 0:
|
|
size = size + 1
|
|
|
|
# Expand buffer if needed
|
|
if self.expand:
|
|
while self.__orderlen < (size + offset):
|
|
self.__append_empty()
|
|
|
|
# Mark buffer as used
|
|
for i in range(size):
|
|
self.order[i + offset] = size
|
|
|
|
def get_next_byte(self) -> Optional[int]:
|
|
"""
|
|
Returns an integer location where the next byte will be found/stored, respecting Konami logic.
|
|
Will return None if its not possible to find this integer a spot and we aren't expanding.
|
|
"""
|
|
# If we expand for additions, make sure we've padded to a 4 byte boundary
|
|
if self.expand:
|
|
while (self.__orderlen & 3) != 0:
|
|
self.__append_empty()
|
|
|
|
for i in range(self.__lastbyte, self.__orderlen, 4):
|
|
if self.order[i] is not None:
|
|
# See if this has room for a byte
|
|
for j in range(0, 4):
|
|
if self.order[i + j] == 1:
|
|
# This is okay, we can pack after this
|
|
continue
|
|
elif self.order[i + j] is None:
|
|
# This is open, pack here
|
|
self.__lastbyte = i
|
|
return i + j
|
|
else:
|
|
# This is something else, can't pack here
|
|
break
|
|
else:
|
|
# Couldn't find optimal packing, pack here
|
|
self.__lastbyte = i
|
|
return i
|
|
|
|
if self.expand:
|
|
self.__lastbyte = self.__orderlen
|
|
return self.__orderlen
|
|
else:
|
|
return None
|
|
|
|
def get_next_short(self) -> Optional[int]:
|
|
"""
|
|
Returns an integer location where the next short will be found/stored, respecting Konami logic.
|
|
Will return None if its not possible to find this integer a spot and we aren't expanding.
|
|
"""
|
|
# If we expand for additions, make sure we've padded to a 4 byte boundary
|
|
if self.expand:
|
|
while (self.__orderlen & 3) != 0:
|
|
self.__append_empty()
|
|
|
|
for i in range(self.__lastshort, self.__orderlen, 4):
|
|
if self.order[i] is not None:
|
|
for j in range(0, 4, 2):
|
|
if self.order[i + j] == 2 and self.order[i + j + 1] == 2:
|
|
# This is okay, we can pack after this
|
|
continue
|
|
elif self.order[i + j] is None and self.order[i + j + 1] is None:
|
|
# This is open, pack here
|
|
self.__lastshort = i
|
|
return i + j
|
|
else:
|
|
# This is something else, can't pack here
|
|
break
|
|
else:
|
|
# Couldn't find optimal packing, pack here
|
|
self.__lastshort = i
|
|
return i
|
|
|
|
if self.expand:
|
|
self.__lastshort = self.__orderlen
|
|
return self.__orderlen
|
|
else:
|
|
return None
|
|
|
|
def get_next_int(self) -> Optional[int]:
|
|
"""
|
|
Returns an integer location where the next integer will be found/stored, respecting Konami logic.
|
|
Will return None if its not possible to find this integer a spot and we aren't expanding.
|
|
"""
|
|
# If we expand for additions, make sure we've padded to a 4 byte boundary
|
|
if self.expand:
|
|
while (self.__orderlen & 3) != 0:
|
|
self.__append_empty()
|
|
|
|
for i in range(self.__lastint, self.__orderlen, 4):
|
|
if self.order[i] is not None:
|
|
continue
|
|
if self.order[i + 1] is not None:
|
|
continue
|
|
if self.order[i + 2] is not None:
|
|
continue
|
|
if self.order[i + 3] is not None:
|
|
continue
|
|
|
|
self.__lastint = i
|
|
return i
|
|
|
|
if self.expand:
|
|
self.__lastint = self.__orderlen
|
|
return self.__orderlen
|
|
else:
|
|
return None
|
|
|
|
@staticmethod
|
|
def node_to_body_ordering(node: Node, include_children: bool=True, include_void: bool=False) -> List[Dict[str, Any]]:
|
|
"""
|
|
Walk this node, attributes and children in the correct order to create a node
|
|
ordering for the purpose of mapping Node objects to their actual data
|
|
in a binary packet data chunk. We will use this to unpack data to determine the
|
|
values of nodes, or to create the data that goes with these nodes.
|
|
|
|
Paramters:
|
|
include_children - Whether this ordering should include children. Defaults to True.
|
|
include_void - Whether this ordering should include positions for void nodes. Defaults
|
|
to false.
|
|
|
|
Returns:
|
|
List of dictionary objects:
|
|
- type - 'attribute' or 'value' to specify that this position in the
|
|
node walk is a string attribute or a node value
|
|
- node - This Node object, for the purpose of assignment
|
|
- name - The name of the attribute if type is 'attribute' or the name
|
|
of the node if type is 'value'
|
|
- alignment - The alignment that this particular data object requiers
|
|
"""
|
|
ordering = []
|
|
|
|
# Include the node itself if it has a value or we include voids
|
|
if node.data_length != 0 or include_void:
|
|
alignment = node.data_length
|
|
if alignment is None:
|
|
# Take care of string types
|
|
alignment = 4
|
|
if alignment > 4:
|
|
# Take care of 64 bit integers that are 32 bit aligned
|
|
alignment = 4
|
|
|
|
ordering.append({
|
|
'type': 'value',
|
|
'node': node,
|
|
'name': node.name,
|
|
'alignment': alignment,
|
|
})
|
|
|
|
order = sorted(node.attributes.keys())
|
|
for attr in order:
|
|
ordering.append({
|
|
'type': 'attribute',
|
|
'node': node,
|
|
'name': attr,
|
|
'alignment': 4,
|
|
})
|
|
|
|
if include_children:
|
|
for child in node.children:
|
|
ordering.extend(PackedOrdering.node_to_body_ordering(child))
|
|
|
|
return ordering
|
|
|
|
|
|
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:
|
|
"""
|
|
Initialize the object.
|
|
|
|
Parameters:
|
|
- data - A binary blob of data to be decoded
|
|
- encoding - A string representing the text encoding for string elements. Should be either
|
|
'shift-jis', 'euc-jp' or 'utf-8'
|
|
"""
|
|
self.stream = InputStream(data)
|
|
self.encoding = encoding
|
|
self.executed = False
|
|
|
|
def __read_node_name(self) -> str:
|
|
"""
|
|
Given the current position in the stream, read the 6-bit-byte packed string name of the
|
|
node.
|
|
|
|
Returns:
|
|
A string representing the name in ascii
|
|
"""
|
|
length = self.stream.read_int()
|
|
binary_length = int(((length * 6) + 7) / 8)
|
|
|
|
def int_to_bin(integer: int) -> str:
|
|
val = bin(integer)[2:]
|
|
while len(val) < 8:
|
|
val = '0' + val
|
|
|
|
return val
|
|
|
|
data = ''
|
|
for i in range(binary_length):
|
|
data = data + int_to_bin(self.stream.read_int())
|
|
data_str = [data[i:(i + 6)] for i in range(0, len(data), 6)]
|
|
data_int = [int(val, 2) for val in data_str]
|
|
ret = ''.join([Node.NODE_NAME_CHARS[val] for val in data_int])
|
|
ret = ret[:length]
|
|
return ret
|
|
|
|
def __read_node(self, node_type: int) -> Node:
|
|
"""
|
|
Given an integer node type, read the node's name, possible attributes
|
|
and children. Will return a Node representing this node. Note
|
|
that calling this on the first node should return a tree of all nodes.
|
|
|
|
Returns:
|
|
Node object
|
|
"""
|
|
name = self.__read_node_name()
|
|
node = Node(name=name, type=node_type)
|
|
|
|
while True:
|
|
child_type = self.stream.read_int()
|
|
|
|
if child_type == Node.END_OF_NODE:
|
|
return node
|
|
elif child_type == Node.ATTR_TYPE:
|
|
key = self.__read_node_name()
|
|
node.set_attribute(key)
|
|
else:
|
|
child = self.__read_node(child_type)
|
|
node.add_child(child)
|
|
|
|
def get_tree(self) -> Node:
|
|
"""
|
|
Parse the header and body such that we can return a Node tree
|
|
representing the data passed to us.
|
|
|
|
Returns:
|
|
Node object
|
|
"""
|
|
if self.executed:
|
|
raise Exception("Logic error, should only call this once per instance")
|
|
self.executed = True
|
|
|
|
# Read the header first
|
|
header_length = self.stream.read_int(4)
|
|
|
|
root = self.__read_node(self.stream.read_int())
|
|
|
|
eod = self.stream.read_int()
|
|
if eod != Node.END_OF_DOCUMENT:
|
|
raise BinaryEncodingException(f'Unknown node type {eod} at end of document')
|
|
|
|
# Skip by any padding
|
|
while self.stream.pos < header_length + 4:
|
|
self.stream.read_byte()
|
|
|
|
# Read the body next
|
|
body_length = self.stream.read_int(4)
|
|
|
|
if body_length is not None and body_length > 0:
|
|
# We have a body
|
|
body = self.stream.read_blob(body_length)
|
|
if body is None:
|
|
raise BinaryEncodingException('Body has insufficient data')
|
|
|
|
ordering = PackedOrdering(body_length)
|
|
|
|
values = PackedOrdering.node_to_body_ordering(root)
|
|
|
|
for value in values:
|
|
node = value['node']
|
|
|
|
if value['type'] == 'attribute':
|
|
size = None
|
|
enc = 's'
|
|
dtype = 'str'
|
|
array = False
|
|
composite = False
|
|
else:
|
|
size = node.data_length
|
|
enc = node.data_encoding
|
|
dtype = node.data_type
|
|
array = node.is_array
|
|
composite = node.is_composite
|
|
|
|
if composite and array:
|
|
raise Exception('Logic error, no support for composite arrays!')
|
|
|
|
if not array:
|
|
# Scalar value
|
|
alignment = value['alignment']
|
|
|
|
if alignment == 1:
|
|
loc = ordering.get_next_byte()
|
|
elif alignment == 2:
|
|
loc = ordering.get_next_short()
|
|
elif alignment == 4:
|
|
loc = ordering.get_next_int()
|
|
|
|
if size is None:
|
|
# The size should be read from the first 4 bytes
|
|
size = struct.unpack('>I', body[loc:(loc + 4)])[0]
|
|
ordering.mark_used(size + 4, loc, round_to=4)
|
|
loc = loc + 4
|
|
|
|
decode_data = body[loc:(loc + size)]
|
|
decode_value = f'>{size}{enc}'
|
|
else:
|
|
# The size is built-in
|
|
ordering.mark_used(size, loc)
|
|
|
|
decode_data = body[loc:(loc + size)]
|
|
decode_value = f'>{enc}'
|
|
|
|
if composite:
|
|
val_list = list(struct.unpack(decode_value, decode_data))
|
|
if value['type'] == 'attribute':
|
|
raise Exception('Logic error, shouldn\'t have composite attribute type!')
|
|
node.set_value(val_list)
|
|
continue
|
|
|
|
val = struct.unpack(decode_value, decode_data)[0]
|
|
|
|
if dtype == 'str':
|
|
# Need to convert this from encoding to standard string.
|
|
# Also, need to lob off the trailing null.
|
|
try:
|
|
val = val[:-1].decode(self.encoding)
|
|
except UnicodeDecodeError:
|
|
# Nothing we can do here
|
|
pass
|
|
|
|
if value['type'] == 'attribute':
|
|
node.set_attribute(value['name'], val)
|
|
else:
|
|
node.set_value(val)
|
|
else:
|
|
# Array value
|
|
loc = ordering.get_next_int()
|
|
|
|
# The raw size in bytes
|
|
length = struct.unpack('>I', body[loc:(loc + 4)])[0]
|
|
elems = int(length / size)
|
|
|
|
ordering.mark_used(length + 4, loc, round_to=4)
|
|
loc = loc + 4
|
|
decode_data = body[loc:(loc + length)]
|
|
decode_value = f'>{enc * elems}'
|
|
|
|
val = struct.unpack(decode_value, decode_data)
|
|
node.set_value([v for v in val])
|
|
|
|
return root
|
|
|
|
|
|
class BinaryEncoder:
|
|
"""
|
|
A class capable of taking a Node tree and encoding it into a binary format.
|
|
"""
|
|
|
|
def __init__(self, tree: Node, encoding: str) -> None:
|
|
"""
|
|
Initialize the object.
|
|
|
|
Parameters:
|
|
tree - A binary blob of data to be decoded
|
|
encoding - A string representing the text encoding for string elements. Should be either
|
|
'shift-jis', 'euc-jp' or 'utf-8'
|
|
"""
|
|
self.stream = OutputStream()
|
|
self.encoding = encoding
|
|
self.tree = tree
|
|
self.__body: List[int] = []
|
|
self.__body_len = 0
|
|
self.executed = False
|
|
|
|
# Generate the characer LUT
|
|
self.char_lut: Dict[str, int] = {}
|
|
for i in range(len(Node.NODE_NAME_CHARS)):
|
|
self.char_lut[Node.NODE_NAME_CHARS[i]] = i
|
|
|
|
def __write_node_name(self, name: str) -> None:
|
|
"""
|
|
Given the current position in the stream, write the 6-bit-byte packed string name of the
|
|
node.
|
|
|
|
Parameters:
|
|
name - A string name which should be encoded as a node name
|
|
"""
|
|
def char_to_bin(ch: str) -> str:
|
|
index = self.char_lut[ch]
|
|
val = bin(index)[2:]
|
|
|
|
while len(val) < 6:
|
|
val = '0' + val
|
|
|
|
return val[-6:]
|
|
|
|
# Convert to six bit bytes
|
|
length = len(name)
|
|
data = ''.join([char_to_bin(c) for c in name])
|
|
|
|
# Pad out the rest with zeros
|
|
while (len(data) & 0x7) != 0:
|
|
data = data + '0'
|
|
|
|
# Convert to 8-bit bytes
|
|
data_chunks = [data[i:(i + 8)] for i in range(0, len(data), 8)]
|
|
data_int = [int(val, 2) for val in data_chunks]
|
|
|
|
# Output
|
|
self.stream.write_int(length)
|
|
for val in data_int:
|
|
self.stream.write_int(val)
|
|
|
|
def __write_node(self, node: Node) -> None:
|
|
"""
|
|
Given an integer node type, read the node's name, possible attributes
|
|
and children. Will return a Node representing this node. Note
|
|
that calling this on the first node should return a tree of all nodes.
|
|
|
|
Parameters:
|
|
node - A Node which should be encoded.
|
|
"""
|
|
to_write = PackedOrdering.node_to_body_ordering(node, include_children=False, include_void=True)
|
|
for thing in to_write:
|
|
# First, write the type of this node out
|
|
if thing['type'] == 'value':
|
|
self.stream.write_int(thing['node'].type)
|
|
else:
|
|
self.stream.write_int(Node.ATTR_TYPE)
|
|
# Now, write the name out
|
|
self.__write_node_name(thing['name'])
|
|
|
|
# Now, write out the children
|
|
for child in node.children:
|
|
self.__write_node(child)
|
|
|
|
# Now, write out the end of node marker
|
|
self.stream.write_int(Node.END_OF_NODE)
|
|
|
|
def __add_data(self, data: bytes, length: int, offset: int) -> None:
|
|
"""
|
|
Given some binary data, a length and an offset, add the data to the offset in the
|
|
output body. This function will ensure that any new bytes that aren't copied are
|
|
zero'd out. This includes bytes before the offset as well as any pad bytes after
|
|
the offset + length in order to pad this body to a 4 byte boundary.
|
|
|
|
Parameters:
|
|
data - A blob of binary data which should be copied into the output
|
|
length - Number of characters of data to copy
|
|
offset - Offset into the body to start copying
|
|
"""
|
|
while self.__body_len < (length + offset):
|
|
self.__body.append(0)
|
|
self.__body_len = self.__body_len + 1
|
|
|
|
# Make sure its padded to 4 bytes
|
|
while (self.__body_len & 0x3) != 0:
|
|
self.__body.append(0)
|
|
self.__body_len = self.__body_len + 1
|
|
|
|
for i in range(length):
|
|
self.__body[offset + i] = data[i]
|
|
|
|
def get_data(self) -> bytes:
|
|
"""
|
|
Encode the header and body into binary formrt.
|
|
|
|
Returns:
|
|
Binary blob of data that can be decoded by a game.
|
|
"""
|
|
if self.executed:
|
|
raise Exception("Logic error, should only call this once per instance")
|
|
self.executed = True
|
|
|
|
# Generate the header first
|
|
self.__write_node(self.tree)
|
|
self.stream.write_int(Node.END_OF_DOCUMENT)
|
|
self.stream.write_pad(4)
|
|
|
|
header_length = len(self.stream.data)
|
|
header = self.stream.data[:]
|
|
|
|
# Generate the body
|
|
values = PackedOrdering.node_to_body_ordering(self.tree)
|
|
if len(values) > 0:
|
|
ordering = PackedOrdering(0, allow_expansion=True)
|
|
|
|
for value in values:
|
|
node = value['node']
|
|
|
|
if value['type'] == 'attribute':
|
|
size = None
|
|
enc = 's'
|
|
dtype = 'str'
|
|
array = False
|
|
composite = False
|
|
val = node.attribute(value['name'])
|
|
else:
|
|
size = node.data_length
|
|
enc = node.data_encoding
|
|
dtype = node.data_type
|
|
array = node.is_array
|
|
composite = node.is_composite
|
|
val = node.value
|
|
|
|
if val is None:
|
|
raise BinaryEncodingException(
|
|
f'Node \'{value["name"]}\' has invalid value None',
|
|
)
|
|
|
|
if not array:
|
|
# Scalar value
|
|
alignment = value['alignment']
|
|
|
|
if alignment == 1:
|
|
loc = ordering.get_next_byte()
|
|
elif alignment == 2:
|
|
loc = ordering.get_next_short()
|
|
elif alignment == 4:
|
|
loc = ordering.get_next_int()
|
|
|
|
if dtype == 'str':
|
|
# Need to convert this to encoding from standard string.
|
|
# Also, need to lob off the trailing null.
|
|
if not isinstance(val, str):
|
|
raise BinaryEncodingException(
|
|
f'Node \'{value["name"]}\' has non-string value!',
|
|
)
|
|
|
|
try:
|
|
valbytes = val.encode(self.encoding) + b'\0'
|
|
except UnicodeEncodeError:
|
|
raise BinaryEncodingException(
|
|
f'Node \'{value["name"]}\' has un-encodable string value \'{val}\''
|
|
)
|
|
size = len(valbytes)
|
|
self.__add_data(struct.pack('>I', size) + valbytes, size + 4, loc)
|
|
ordering.mark_used(size + 4, loc, round_to=4)
|
|
|
|
# We took care of this one
|
|
continue
|
|
elif dtype == 'bin':
|
|
# Store raw binary
|
|
size = len(val)
|
|
self.__add_data(struct.pack('>I', size) + val, size + 4, loc)
|
|
ordering.mark_used(size + 4, loc, round_to=4)
|
|
|
|
# We took care of this one
|
|
continue
|
|
elif composite:
|
|
# Array, but not, somewhat silly
|
|
encode_value = f'>{enc}'
|
|
self.__add_data(struct.pack(encode_value, *val), size, loc)
|
|
ordering.mark_used(size, loc)
|
|
|
|
# We took care of this one
|
|
continue
|
|
elif dtype == 'bool':
|
|
val = 1 if val else 0
|
|
|
|
# The size is built-in, emit it
|
|
encode_value = f'>{enc}'
|
|
self.__add_data(struct.pack(encode_value, val), size, loc)
|
|
ordering.mark_used(size, loc)
|
|
else:
|
|
# Array value
|
|
loc = ordering.get_next_int()
|
|
|
|
# The raw size in bytes
|
|
elems = len(val)
|
|
length = elems * size
|
|
|
|
# Write out the header (number of bytes taken up)
|
|
data = struct.pack('>I', length)
|
|
encode_value = f'>{enc}'
|
|
|
|
# Write out data one element at a time
|
|
for v in val:
|
|
if dtype == 'bool':
|
|
data = data + struct.pack(encode_value, 1 if v else 0)
|
|
else:
|
|
data = data + struct.pack(encode_value, v)
|
|
|
|
self.__add_data(data, length + 4, loc)
|
|
ordering.mark_used(length + 4, loc, round_to=4)
|
|
|
|
return b''.join([
|
|
struct.pack('>I', header_length),
|
|
header,
|
|
struct.pack('>I', self.__body_len),
|
|
bytes(self.__body),
|
|
])
|
|
|
|
|
|
class BinaryEncoding:
|
|
"""
|
|
Wrapper class representing a Binary Encoding.
|
|
"""
|
|
MAGIC = 0xA0
|
|
|
|
COMPRESSED_WITH_DATA = 0x42
|
|
COMPRESSED_WITHOUT_DATA = 0x43
|
|
DECOMPRESSED_WITH_DATA = 0x45
|
|
DECOMPRESSED_WITHOUT_DATA = 0x46
|
|
|
|
# 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.
|
|
ENCODINGS = {
|
|
0x00: "ascii",
|
|
0x20: "shift-jis-legacy",
|
|
0x60: "euc-jp",
|
|
0x80: "shift-jis",
|
|
0xA0: "utf-8",
|
|
}
|
|
|
|
def __init__(self) -> None:
|
|
"""
|
|
Initialize the encoding object.
|
|
"""
|
|
self.encoding: Optional[str] = None
|
|
|
|
def __sanitize_encoding(self, enc: str) -> str:
|
|
"""
|
|
Convert an internal encoding value from an externally acceptible value.
|
|
|
|
Parameters:
|
|
enc - The encoding as a string as passed from an outside caller
|
|
|
|
Returns:
|
|
An encoding string suitable for internal use.
|
|
"""
|
|
if enc == "shift-jis-legacy":
|
|
return "shift-jis"
|
|
return enc
|
|
|
|
def decode(self, data: bytes, skip_on_exceptions: bool=False) -> Optional[Node]:
|
|
"""
|
|
Given a data blob, decode the data with the current encoding. Will
|
|
also set the class property value 'encoding' to the encoding used
|
|
on the last decode.
|
|
|
|
Parameters:
|
|
data - Binary blob representing the data to decode
|
|
|
|
Returns:
|
|
Node object representing the root of the decoded tree, or None
|
|
if we couldn't decode the object for some reason.
|
|
"""
|
|
try:
|
|
data_magic, contents, encoding_raw, encoding_swapped = struct.unpack(">BBBB", data[0:4])
|
|
except struct.error:
|
|
# Couldn't even parse magic
|
|
return None
|
|
|
|
if data_magic != BinaryEncoding.MAGIC:
|
|
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.
|
|
return None
|
|
|
|
encoding = BinaryEncoding.ENCODINGS.get(encoding_raw)
|
|
|
|
if encoding is not None:
|
|
self.encoding = encoding
|
|
try:
|
|
decoder = BinaryDecoder(data[4:], self.__sanitize_encoding(encoding))
|
|
return decoder.get_tree()
|
|
except BinaryEncodingException:
|
|
if skip_on_exceptions:
|
|
return None
|
|
else:
|
|
raise
|
|
else:
|
|
return None
|
|
|
|
def encode(self, tree: Node, encoding: str=None) -> bytes:
|
|
"""
|
|
Given a tree of Node objects, encode the data with the current encoding.
|
|
|
|
Parameters:
|
|
tree - Node tree representing the data to encode
|
|
encoding - The text encoding to use. If None, will try to use the encoding from
|
|
the last successful decode
|
|
|
|
Returns:
|
|
Binary blob representing encoded data
|
|
"""
|
|
if encoding is None:
|
|
encoding = self.encoding
|
|
if encoding is None:
|
|
raise BinaryEncodingException('Unknown encoding')
|
|
|
|
encoding_magic = None
|
|
for magic, encstr in BinaryEncoding.ENCODINGS.items():
|
|
if encstr == encoding:
|
|
encoding_magic = magic
|
|
break
|
|
|
|
if encoding_magic is None:
|
|
raise BinaryEncodingException(f"Invalid text encoding {encoding}")
|
|
|
|
encoder = BinaryEncoder(tree, self.__sanitize_encoding(encoding))
|
|
data = encoder.get_data()
|
|
|
|
if data is not None:
|
|
return struct.pack(">BBBB", BinaryEncoding.MAGIC, BinaryEncoding.COMPRESSED_WITH_DATA, encoding_magic, (~encoding_magic & 0xFF)) + data
|
|
else:
|
|
return None
|