""" 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._custom 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