203 lines
6.6 KiB
Python
203 lines
6.6 KiB
Python
import struct
|
|
from typing import List, Optional
|
|
|
|
|
|
class StreamError(Exception):
|
|
"""
|
|
An exception thrown when something goes wrong with the stream.
|
|
"""
|
|
|
|
|
|
class InputStream:
|
|
"""
|
|
A class that treats a binary blob as a stream of bytes to be emitted.
|
|
Makes stream-like algorithms much easier to implement. All accessor
|
|
functions that read data will advance the current position. It is not
|
|
rewindable.
|
|
"""
|
|
|
|
def __init__(self, data: bytes) -> None:
|
|
"""
|
|
Initialize the object. Given a data blob, will set this as the stream
|
|
and set the location to the beginning of the data blob.
|
|
|
|
Parameters:
|
|
data - A binary blob to read from.
|
|
"""
|
|
self.data = data
|
|
self.pos = 0
|
|
self.left = len(self.data)
|
|
|
|
def read_blob(self, blob_size: int) -> Optional[bytes]:
|
|
"""
|
|
Given a blob size, read the next blob_size bytes as a binary blob.
|
|
|
|
Parameters:
|
|
blob_size - An integer representing the number of bytes to read.
|
|
|
|
Returns:
|
|
a binary string representing blob_size bytes from the current location, or None
|
|
if there wasn't enough bytes to satisfy this request.
|
|
"""
|
|
if blob_size <= 0:
|
|
return None
|
|
if blob_size <= self.left:
|
|
bytedata = self.data[self.pos : (self.pos + blob_size)]
|
|
self.pos += blob_size
|
|
self.left -= blob_size
|
|
return bytedata
|
|
return None
|
|
|
|
def read_byte(self) -> Optional[bytes]:
|
|
"""
|
|
Grab the next byte at the current position. If no byte is available,
|
|
return None.
|
|
|
|
Returns:
|
|
a raw byte
|
|
"""
|
|
return self.read_blob(1)
|
|
|
|
def read_int(self, size: int = 1, is_unsigned: bool = True) -> Optional[int]:
|
|
"""
|
|
Grab the next integer of size 'size' at the current position. If not enough
|
|
bytes are available to decode this integer, return None.
|
|
|
|
Parameters:
|
|
size - Integer representing the integer size to decode. Valid values are
|
|
1, 2 and 4 for char, short and int respectively.
|
|
is_unsigned - An optional boolean specifying whether the integer read should be
|
|
unsigned. Defaults to True.
|
|
|
|
Returns:
|
|
a python integer representing the big-endian decoding of the current
|
|
position
|
|
"""
|
|
if size == 1:
|
|
data = self.read_blob(1)
|
|
if data is None:
|
|
return None
|
|
|
|
if is_unsigned:
|
|
# Fastpath, just use python's own decoder
|
|
return data[0]
|
|
else:
|
|
return struct.unpack(">b", data)[0]
|
|
elif size == 2:
|
|
data = self.read_blob(2)
|
|
if data is None:
|
|
return None
|
|
|
|
if is_unsigned:
|
|
return struct.unpack(">H", data)[0]
|
|
else:
|
|
return struct.unpack(">h", data)[0]
|
|
elif size == 4:
|
|
data = self.read_blob(4)
|
|
if data is None:
|
|
return None
|
|
|
|
if is_unsigned:
|
|
return struct.unpack(">I", data)[0]
|
|
else:
|
|
return struct.unpack(">i", data)[0]
|
|
else:
|
|
raise StreamError(f"Unsupported size {size}")
|
|
|
|
|
|
class OutputStream:
|
|
"""
|
|
A class that treats a binary blob as a stream of bytes to be constructed.
|
|
Makes stream-like algorithms much easier to implement. All accessor
|
|
functions that write data will advance the current position. It is not
|
|
rewindable. When finished writing, access the finished blob by copying from
|
|
data.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
"""
|
|
Initialize the object.
|
|
"""
|
|
self.__data: List[bytes] = []
|
|
self.__data_len = 0
|
|
self.__formatted_data: Optional[bytes] = None
|
|
|
|
@property
|
|
def data(self) -> bytes:
|
|
if self.__formatted_data is None:
|
|
self.__formatted_data = b"".join(self.__data)
|
|
return self.__formatted_data
|
|
|
|
def write_blob(self, blob: bytes) -> int:
|
|
"""
|
|
Write a binary blob of data to the stream
|
|
|
|
Parameters:
|
|
blob - An blob of data to write.
|
|
|
|
Returns:
|
|
the number of bytes written
|
|
"""
|
|
self.__data.append(blob)
|
|
self.__data_len += len(blob)
|
|
return len(blob)
|
|
|
|
def write_byte(self, byte: bytes) -> None:
|
|
"""
|
|
Write a raw byte to the end of the output stream.
|
|
|
|
Parameters:
|
|
A byte that should be appended to the current stream.
|
|
"""
|
|
self.__data.append(byte)
|
|
self.__data_len = self.__data_len + 1
|
|
self.__formatted_data = None
|
|
|
|
def write_int(self, integer: int, size: int = 1, is_unsigned: bool = True) -> None:
|
|
"""
|
|
Write an integer to the end of the output stream.
|
|
|
|
Parameters:
|
|
integer - The integer that should be written to the stream.
|
|
size - The byte size of the integer. Supports 1, 2 and 4 byte
|
|
integer types.
|
|
is_unsigned - Whether the integer should be written unsigned or
|
|
signed. Defaults to True.
|
|
"""
|
|
if size == 1:
|
|
if is_unsigned:
|
|
self.__data.append(struct.pack(">B", integer))
|
|
else:
|
|
self.__data.append(struct.pack(">b", integer))
|
|
self.__data_len = self.__data_len + 1
|
|
elif size == 2:
|
|
if is_unsigned:
|
|
self.__data.append(struct.pack(">H", integer))
|
|
else:
|
|
self.__data.append(struct.pack(">h", integer))
|
|
self.__data_len = self.__data_len + 2
|
|
elif size == 4:
|
|
if is_unsigned:
|
|
self.__data.append(struct.pack(">I", integer))
|
|
else:
|
|
self.__data.append(struct.pack(">i", integer))
|
|
self.__data_len = self.__data_len + 4
|
|
else:
|
|
raise StreamError(f"Unsupported size {size}")
|
|
self.__formatted_data = None
|
|
|
|
def write_pad(self, pad_to: int) -> None:
|
|
"""
|
|
Pad the current stream to a byte boundary specified by pad_to.
|
|
|
|
Parameters:
|
|
pad_to - An integer specifying the byte alignment that should be present
|
|
after padding is complete. Supports 1, 2, 4, 8, 16 or any other power of
|
|
two padding. After calling this, the next write_byte or write_int will
|
|
be placed on a boundary compatible with the pad_to parameter.
|
|
"""
|
|
while (self.__data_len & (pad_to - 1)) != 0:
|
|
self.__data.append(b"\0")
|
|
self.__data_len = self.__data_len + 1
|
|
self.__formatted_data = None
|