mirror of
https://github.com/LonamiWebs/Telethon
synced 2025-08-23 10:27:40 +00:00
94 lines
2.4 KiB
Python
94 lines
2.4 KiB
Python
|
from io import BytesIO, BufferedWriter
|
||
|
|
from struct import pack
|
||
|
|
|
||
|
|
|
||
|
|
class BinaryWriter:
|
||
|
|
"""
|
||
|
|
Small utility class to write binary data.
|
||
|
|
Also creates a "Memory Stream" if necessary
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, stream=None):
|
||
|
|
if not stream:
|
||
|
|
stream = BytesIO()
|
||
|
|
|
||
|
|
self.stream = stream
|
||
|
|
self.writer = BufferedWriter(self.stream)
|
||
|
|
|
||
|
|
# region Writing
|
||
|
|
|
||
|
|
def write_byte(self, byte):
|
||
|
|
self.writer.write(pack('B', byte))
|
||
|
|
|
||
|
|
def write_int(self, integer, signed=True):
|
||
|
|
if not signed:
|
||
|
|
integer &= 0xFFFFFFFF # Ensure it's unsigned (see http://stackoverflow.com/a/30092291/4759433)
|
||
|
|
self.writer.write(pack('I', integer))
|
||
|
|
|
||
|
|
def write_long(self, long, signed=True):
|
||
|
|
if not signed:
|
||
|
|
long &= 0xFFFFFFFFFFFFFFFF
|
||
|
|
self.writer.write(pack('Q', long))
|
||
|
|
|
||
|
|
def write(self, data):
|
||
|
|
self.writer.write(data)
|
||
|
|
|
||
|
|
# endregion
|
||
|
|
|
||
|
|
# region Telegram custom writing
|
||
|
|
|
||
|
|
def tgwrite_bytes(self, data):
|
||
|
|
|
||
|
|
if len(data) < 254:
|
||
|
|
padding = (len(data) + 1) % 4
|
||
|
|
if padding != 0:
|
||
|
|
padding = 4 - padding
|
||
|
|
|
||
|
|
self.write(bytes([len(data)]))
|
||
|
|
self.write(data)
|
||
|
|
|
||
|
|
else:
|
||
|
|
padding = len(data) % 4
|
||
|
|
if padding != 0:
|
||
|
|
padding = 4 - padding
|
||
|
|
|
||
|
|
# TODO ensure that _this_ is right (it appears to be)
|
||
|
|
self.write(bytes([254]))
|
||
|
|
self.write(bytes([len(data) % 256]))
|
||
|
|
self.write(bytes([(len(data) >> 8) % 256]))
|
||
|
|
self.write(bytes([(len(data) >> 16) % 256]))
|
||
|
|
self.write(data)
|
||
|
|
|
||
|
|
""" Original:
|
||
|
|
binaryWriter.Write((byte)254);
|
||
|
|
binaryWriter.Write((byte)(bytes.Length));
|
||
|
|
binaryWriter.Write((byte)(bytes.Length >> 8));
|
||
|
|
binaryWriter.Write((byte)(bytes.Length >> 16));
|
||
|
|
"""
|
||
|
|
|
||
|
|
self.write(bytes(padding))
|
||
|
|
|
||
|
|
def tgwrite_string(self, string):
|
||
|
|
return self.tgwrite_bytes(string.encode('utf-8'))
|
||
|
|
|
||
|
|
# endregion
|
||
|
|
|
||
|
|
def flush(self):
|
||
|
|
self.writer.flush()
|
||
|
|
|
||
|
|
def close(self):
|
||
|
|
self.writer.close()
|
||
|
|
# TODO Do I need to close the underlying stream?
|
||
|
|
|
||
|
|
def get_bytes(self, flush=True):
|
||
|
|
if flush:
|
||
|
|
self.writer.flush()
|
||
|
|
self.stream.getbuffer()
|
||
|
|
|
||
|
|
# with block
|
||
|
|
def __enter__(self):
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
|
|
self.close()
|