import struct
import binascii
import io
from typing import Union
from neo.IO.MemoryStream import StreamManager
class BinaryWriter(object):
"""A convenience class for writing data from byte streams"""
def __init__(self, stream: Union[bytearray, bytes]) -> None:
"""
Create an instance.
Args:
stream: a stream to operate on.
"""
super(BinaryWriter, self).__init__()
self._stream = StreamManager.GetStream(stream)
def write_bytes(self, value: bytes, unhex: bool = True) -> int:
"""
Write a `bytes` type to the stream.
Args:
value: array of bytes to write to the stream.
unhex: (Default) True. Set to unhexlify the stream. Use when the bytes are not raw bytes; i.e. b'aabb'
Returns:
int: the number of bytes written.
"""
if unhex:
try:
value = binascii.unhexlify(value)
except binascii.Error:
pass
return self._stream.write(value)
def _pack(self, fmt, data) -> int:
"""
Write bytes by packing them according to the provided format `fmt`.
For more information about the `fmt` format see: https://docs.python.org/3/library/struct.html
Args:
fmt (str): format string.
data (object): the data to write to the raw stream.
Returns:
int: the number of bytes written.
"""
return self.write_bytes(struct.pack(fmt, data), unhex=False)
def write_bool(self, value: bool) -> int:
"""
Pack the value as a bool and write 1 byte to the stream.
Args:
value: the boolean value to write.
Returns:
int: the number of bytes written.
"""
return self._pack('?', value)
def write_uint8(self, value):
return self.write_bytes(bytes([value]))
def write_uint16(self, value, endian="<"):
"""
Pack the value as an unsigned integer and write 2 bytes to the stream.
Args:
value:
endian: specify the endianness. (Default) Little endian ('<'). Use '>' for big endian.
Returns:
int: the number of bytes written.
"""
return self._pack('%sH' % endian, value)
def write_uint32(self, value, endian="<") -> int:
"""
Pack the value as a signed integer and write 4 bytes to the stream.
Args:
value:
endian: specify the endianness. (Default) Little endian ('<'). Use '>' for big endian.
Returns:
int: the number of bytes written.
"""
return self._pack('%sI' % endian, value)
def write_uint64(self, value, endian="<") -> int:
"""
Pack the value as an unsigned integer and write 8 bytes to the stream.
Args:
value:
endian (str): specify the endianness. (Default) Little endian ('<'). Use '>' for big endian.
Returns:
int: the number of bytes written.
"""
return self._pack('%sQ' % endian, value)
def write_uint256(self, value, endian="<") -> int:
return self.write_bytes(value._data)
def write_uint160(self, value, endian="<") -> int:
return self.write_bytes(value._data)
def write_var_string(self, value: str, encoding: str = "utf-8") -> int:
"""
Write a string value to the stream.
Read more about variable size encoding here: http://docs.neo.org/en-us/node/network-protocol.html#convention
Args:
value: value to write to the stream.
encoding: string encoding format.
"""
if type(value) is str:
data = value.encode(encoding)
length = len(data)
self.write_var_int(length)
written = self.write_bytes(data)
return written
def write_var_int(self, value: int, endian: str = "<") -> int:
"""
Write an integer value in a space saving way to the stream.
Read more about variable size encoding here: http://docs.neo.org/en-us/node/network-protocol.html#convention
Args:
value:
endian: specify the endianness. (Default) Little endian ('<'). Use '>' for big endian.
Raises:
{TypeError}: if ``value`` is not of type int.
ValueError: if `value` is < 0.
Returns:
int: the number of bytes written.
"""
if not isinstance(value, int):
raise TypeError('%s not int type.' % value)
if value < 0:
raise ValueError('%d too small.' % value)
elif value < 0xfd:
return self.write_bytes(bytes([value]))
elif value <= 0xffff:
self.write_bytes(bytes([0xfd]))
return self.write_uint16(value, endian)
elif value <= 0xFFFFFFFF:
self.write_bytes(bytes([0xfe]))
return self.write_uint32(value, endian)
else:
self.write_bytes(bytes([0xff]))
return self.write_uint64(value, endian)
def write_fixed_string(self, value, length):
"""
Write a string value to the stream.
Args:
value (str): value to write to the stream.
length (int): length of the string to write.
"""
towrite = value.encode('utf-8')
slen = len(towrite)
if slen > length:
raise Exception("string longer than fixed length: %s " % length)
self.write_bytes(towrite)
diff = length - slen
while diff > 0:
self.write_bytes(bytes([0]))
diff -= 1
def write_var_bytes(self, value: int, endian: str = "<") -> int:
self.write_var_int(len(value), endian)
return self.write_bytes(value)
def cleanup(self):
if self._stream:
StreamManager.ReleaseStream(self._stream)
