Telethon/telethon/utils/binary_reader.py
2017-05-21 13:02:54 +02:00

169 lines
5.1 KiB
Python

import os
from datetime import datetime
from io import BufferedReader, BytesIO
from struct import unpack
from ..errors import InvalidParameterError, TypeNotFoundError
from ..tl.all_tlobjects import tlobjects
class BinaryReader:
"""
Small utility class to read binary data.
Also creates a "Memory Stream" if necessary
"""
def __init__(self, data=None, stream=None):
if data:
self.stream = BytesIO(data)
elif stream:
self.stream = stream
else:
raise InvalidParameterError(
'Either bytes or a stream must be provided')
self.reader = BufferedReader(self.stream)
# region Reading
# "All numbers are written as little endian." |> Source: 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('<f', self.read(4))[0]
def read_double(self):
"""Reads a real floating point (8 bytes) value"""
return unpack('<d', self.read(8))[0]
def read_large_int(self, bits, signed=True):
"""Reads a n-bits long integer value"""
return int.from_bytes(
self.read(bits // 8), byteorder='little', signed=signed)
def read(self, length):
"""Read the given amount of bytes"""
result = self.reader.read(length)
if len(result) != length:
raise BufferError(
'Trying to read outside the data bounds (no more data left to read)')
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 ValueError('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 None if value == 0 else datetime.fromtimestamp(value)
def tgread_object(self):
"""Reads a Telegram object"""
constructor_id = self.read_int(signed=False)
clazz = 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
# If there was still no luck, give up
raise TypeNotFoundError(constructor_id)
# Create an empty instance of the class and
# fill it with the read attributes
result = clazz.empty()
result.on_response(self)
return result
def tgread_vector(self):
"""Reads a vector (a list) of Telegram objects"""
if 0x1cb5c415 != self.read_int(signed=False):
raise ValueError('Invalid constructor code, vector was expected')
count = self.read_int()
return [self.tgread_object() for _ in range(count)]
# endregion
def close(self):
self.reader.close()
# region Position related
def tell_position(self):
"""Tells the current position on the stream"""
return self.reader.tell()
def set_position(self, position):
"""Sets the current position on the stream"""
self.reader.seek(position)
def seek(self, offset):
"""Seeks the stream position given an offset from the current position. May be negative"""
self.reader.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