mirror of
https://github.com/DragonMinded/bemaniutils.git
synced 2026-10-07 06:18:15 +03:00
Initial commit of BEMANI Utilities to GitHub.
This commit is contained in:
Jennifer Taylor
@@ -0,0 +1,779 @@
|
||||
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.add_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('Unknown node type {} at end of document'.format(eod))
|
||||
|
||||
# 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 = '>{}{}'.format(size, enc)
|
||||
else:
|
||||
# The size is built-in
|
||||
ordering.mark_used(size, loc)
|
||||
|
||||
decode_data = body[loc:(loc + size)]
|
||||
decode_value = '>{}'.format(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 = '>{}'.format(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(
|
||||
'Node \'{}\' has invalid value None'.format(value['name']),
|
||||
)
|
||||
|
||||
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(
|
||||
'Node \'{}\' has non-string value!'.format(value['name']),
|
||||
)
|
||||
|
||||
try:
|
||||
valbytes = val.encode(self.encoding) + b'\0'
|
||||
except UnicodeEncodeError:
|
||||
raise BinaryEncodingException(
|
||||
'Node \'{}\' has un-encodable string value \'{}\''.format(value['name'], 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 = '>{}'.format(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 = '>{}'.format(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 = '>{}'.format(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.
|
||||
"""
|
||||
# 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.
|
||||
MAGIC = {
|
||||
0xA04200FF: "ascii",
|
||||
0xA04220DF: "shift-jis-legacy",
|
||||
0xA042609F: "euc-jp",
|
||||
0xA042807F: "shift-jis",
|
||||
0xA042A05F: "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 = struct.unpack(">L", data[0:4])[0]
|
||||
except struct.error:
|
||||
# Couldn't even parse magic
|
||||
return None
|
||||
encoding = BinaryEncoding.MAGIC.get(data_magic)
|
||||
|
||||
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 in BinaryEncoding.MAGIC:
|
||||
if BinaryEncoding.MAGIC[magic] == encoding:
|
||||
encoding_magic = magic
|
||||
break
|
||||
|
||||
if encoding_magic is None:
|
||||
raise BinaryEncodingException("Invalid text encoding {}".format(encoding))
|
||||
|
||||
encoder = BinaryEncoder(tree, self.__sanitize_encoding(encoding))
|
||||
data = encoder.get_data()
|
||||
|
||||
if data is not None:
|
||||
return struct.pack(">L", encoding_magic) + data
|
||||
else:
|
||||
return None
|
||||
Reference in New Issue
Block a user