""" This module contains the BinaryReader utility class. """ import os import time from datetime import datetime, timezone, timedelta from io import BytesIO from struct import unpack from ..errors import TypeNotFoundError from .. import _tl from ..types import _core _EPOCH_NAIVE = datetime(*time.gmtime(0)[:6]) _EPOCH = _EPOCH_NAIVE.replace(tzinfo=timezone.utc) class BinaryReader: """ Small utility class to read binary data. """ def __init__(self, data): self.stream = BytesIO(data) self._last = None # Should come in handy to spot -404 errors # region Reading # "All numbers are written as little endian." # https://core.telegram.org/mtproto def read_byte(self): """Reads a single byte value.""" return self.read(1)[0] def read_int(self, signed=True): """Reads an integer (4 bytes) value.""" return int.from_bytes(self.read(4), byteorder='little', signed=signed) def read_long(self, signed=True): """Reads a long integer (8 bytes) value.""" return int.from_bytes(self.read(8), byteorder='little', signed=signed) def read_float(self): """Reads a real floating point (4 bytes) value.""" return unpack('= 0) and (len(result) != length): raise BufferError( 'No more data left to read (need {}, got {}: {}); last read {}' .format(length, len(result), repr(result), repr(self._last)) ) self._last = result return result def get_bytes(self): """Gets the byte array representing the current buffer as a whole.""" return self.stream.getvalue() # endregion # region Telegram custom reading def tgread_bytes(self): """ Reads a Telegram-encoded byte array, without the need of specifying its length. """ first_byte = self.read_byte() if first_byte == 254: length = self.read_byte() | (self.read_byte() << 8) | ( self.read_byte() << 16) padding = length % 4 else: length = first_byte padding = (length + 1) % 4 data = self.read(length) if padding > 0: padding = 4 - padding self.read(padding) return data def tgread_string(self): """Reads a Telegram-encoded string.""" return str(self.tgread_bytes(), encoding='utf-8', errors='replace') def tgread_bool(self): """Reads a Telegram boolean value.""" value = self.read_int(signed=False) if value == 0x997275b5: # boolTrue return True elif value == 0xbc799737: # boolFalse return False else: raise RuntimeError('Invalid boolean code {}'.format(hex(value))) def tgread_date(self): """Reads and converts Unix time (used by Telegram) into a Python datetime object. """ value = self.read_int() return _EPOCH + timedelta(seconds=value) def tgread_object(self): """Reads a Telegram object.""" constructor_id = self.read_int(signed=False) clazz = _tl.tlobjects.get(constructor_id, None) if clazz is None: # The class was None, but there's still a # chance of it being a manually parsed value like bool! value = constructor_id if value == 0x997275b5: # boolTrue return True elif value == 0xbc799737: # boolFalse return False elif value == 0x1cb5c415: # Vector return [self.tgread_object() for _ in range(self.read_int())] clazz = _core.core_objects.get(constructor_id, None) if clazz is None: # If there was still no luck, give up self.seek(-4) # Go back pos = self.tell_position() error = TypeNotFoundError(constructor_id, self.read()) self.set_position(pos) raise error return clazz.from_reader(self) def tgread_vector(self): """Reads a vector (a list) of Telegram objects.""" if 0x1cb5c415 != self.read_int(signed=False): raise RuntimeError('Invalid constructor code, vector was expected') count = self.read_int() return [self.tgread_object() for _ in range(count)] # endregion def close(self): """Closes the reader, freeing the BytesIO stream.""" self.stream.close() # region Position related def tell_position(self): """Tells the current position on the stream.""" return self.stream.tell() def set_position(self, position): """Sets the current position on the stream.""" self.stream.seek(position) def seek(self, offset): """ Seeks the stream position given an offset from the current position. The offset may be negative. """ self.stream.seek(offset, os.SEEK_CUR) # endregion # region with block def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self.close() # endregion