mirror of
https://github.com/LonamiWebs/Telethon.git
synced 2024-09-21 03:08:49 +03:00
Completely remove BinaryWriter from the project
This commit is contained in:
Lonami Exo
parent
2fd03e5422
commit
71c2a82b3d
|
|
@ -1,9 +1,7 @@
|
|||
"""
|
||||
Several extensions Python is missing, such as a proper class to handle a TCP
|
||||
communication with support for cancelling the operation, and an utility class
|
||||
to work with arbitrary binary data in a more comfortable way (writing ints,
|
||||
strings, bytes, etc.)
|
||||
to read arbitrary binary data in a more comfortable way, with int/strings/etc.
|
||||
"""
|
||||
from .binary_writer import BinaryWriter
|
||||
from .binary_reader import BinaryReader
|
||||
from .tcp_client import TcpClient
|
||||
|
|
|
|||
|
|
@ -1,152 +0,0 @@
|
|||
from io import BufferedWriter, BytesIO, DEFAULT_BUFFER_SIZE
|
||||
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, known_length=None):
|
||||
if not stream:
|
||||
stream = BytesIO()
|
||||
|
||||
if known_length is None:
|
||||
# On some systems, DEFAULT_BUFFER_SIZE defaults to 8192
|
||||
# That's over 16 times as big as necessary for most messages
|
||||
known_length = max(DEFAULT_BUFFER_SIZE, 1024)
|
||||
|
||||
self.writer = BufferedWriter(stream, buffer_size=known_length)
|
||||
self.written_count = 0
|
||||
|
||||
# region Writing
|
||||
|
||||
# "All numbers are written as little endian."
|
||||
# https://core.telegram.org/mtproto
|
||||
def write_byte(self, value):
|
||||
"""Writes a single byte value"""
|
||||
self.writer.write(pack('B', value))
|
||||
self.written_count += 1
|
||||
|
||||
def write_int(self, value, signed=True):
|
||||
"""Writes an integer value (4 bytes), optionally signed"""
|
||||
self.writer.write(
|
||||
int.to_bytes(
|
||||
value, length=4, byteorder='little', signed=signed))
|
||||
self.written_count += 4
|
||||
|
||||
def write_long(self, value, signed=True):
|
||||
"""Writes a long integer value (8 bytes), optionally signed"""
|
||||
self.writer.write(
|
||||
int.to_bytes(
|
||||
value, length=8, byteorder='little', signed=signed))
|
||||
self.written_count += 8
|
||||
|
||||
def write_float(self, value):
|
||||
"""Writes a floating point value (4 bytes)"""
|
||||
self.writer.write(pack('<f', value))
|
||||
self.written_count += 4
|
||||
|
||||
def write_double(self, value):
|
||||
"""Writes a floating point value (8 bytes)"""
|
||||
self.writer.write(pack('<d', value))
|
||||
self.written_count += 8
|
||||
|
||||
def write_large_int(self, value, bits, signed=True):
|
||||
"""Writes a n-bits long integer value"""
|
||||
self.writer.write(
|
||||
int.to_bytes(
|
||||
value, length=bits // 8, byteorder='little', signed=signed))
|
||||
self.written_count += bits // 8
|
||||
|
||||
def write(self, data):
|
||||
"""Writes the given bytes array"""
|
||||
self.writer.write(data)
|
||||
self.written_count += len(data)
|
||||
|
||||
# endregion
|
||||
|
||||
# region Telegram custom writing
|
||||
|
||||
def tgwrite_bytes(self, data):
|
||||
"""Write bytes by using Telegram guidelines"""
|
||||
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
|
||||
|
||||
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)
|
||||
|
||||
self.write(bytes(padding))
|
||||
|
||||
def tgwrite_string(self, string):
|
||||
"""Write a string by using Telegram guidelines"""
|
||||
self.tgwrite_bytes(string.encode('utf-8'))
|
||||
|
||||
def tgwrite_bool(self, boolean):
|
||||
"""Write a boolean value by using Telegram guidelines"""
|
||||
# boolTrue boolFalse
|
||||
self.write_int(0x997275b5 if boolean else 0xbc799737, signed=False)
|
||||
|
||||
def tgwrite_date(self, datetime):
|
||||
"""Converts a Python datetime object into Unix time
|
||||
(used by Telegram) and writes it
|
||||
"""
|
||||
value = 0 if datetime is None else int(datetime.timestamp())
|
||||
self.write_int(value)
|
||||
|
||||
def tgwrite_object(self, tlobject):
|
||||
"""Writes a Telegram object"""
|
||||
self.write(tlobject.to_bytes())
|
||||
|
||||
def tgwrite_vector(self, vector):
|
||||
"""Writes a vector of Telegram objects"""
|
||||
self.write_int(0x1cb5c415, signed=False) # Vector's constructor ID
|
||||
self.write_int(len(vector))
|
||||
for item in vector:
|
||||
self.tgwrite_object(item)
|
||||
|
||||
# endregion
|
||||
|
||||
def flush(self):
|
||||
"""Flush the current stream to "update" changes"""
|
||||
self.writer.flush()
|
||||
|
||||
def close(self):
|
||||
"""Close the current stream"""
|
||||
self.writer.close()
|
||||
|
||||
def get_bytes(self, flush=True):
|
||||
"""Get the current bytes array content from the buffer,
|
||||
optionally flushing first
|
||||
"""
|
||||
if flush:
|
||||
self.writer.flush()
|
||||
return self.writer.raw.getvalue()
|
||||
|
||||
def get_written_bytes_count(self):
|
||||
"""Gets the count of bytes written in the buffer.
|
||||
This may NOT be equal to the stream length if one
|
||||
was provided when initializing the writer
|
||||
"""
|
||||
return self.written_count
|
||||
|
||||
# with block
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.close()
|
||||
|
|
@ -1,90 +1,61 @@
|
|||
import os
|
||||
import unittest
|
||||
from telethon.extensions import BinaryReader, BinaryWriter
|
||||
from telethon.tl import TLObject
|
||||
from telethon.extensions import BinaryReader
|
||||
|
||||
|
||||
class UtilsTests(unittest.TestCase):
|
||||
@staticmethod
|
||||
def test_binary_writer_reader():
|
||||
# Test that we can write and read properly
|
||||
with BinaryWriter() as writer:
|
||||
writer.write_byte(1)
|
||||
writer.write_int(5)
|
||||
writer.write_long(13)
|
||||
writer.write_float(17.0)
|
||||
writer.write_double(25.0)
|
||||
writer.write(bytes([26, 27, 28, 29, 30, 31, 32]))
|
||||
writer.write_large_int(2**127, 128, signed=False)
|
||||
|
||||
data = writer.get_bytes()
|
||||
expected = b'\x01\x05\x00\x00\x00\r\x00\x00\x00\x00\x00\x00\x00\x00\x00\x88A\x00\x00\x00\x00\x00\x00' \
|
||||
b'9@\x1a\x1b\x1c\x1d\x1e\x1f \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x80'
|
||||
|
||||
assert data == expected, 'Retrieved data does not match the expected value'
|
||||
# Test that we can read properly
|
||||
data = b'\x01\x05\x00\x00\x00\r\x00\x00\x00\x00\x00\x00\x00\x00\x00' \
|
||||
b'\x88A\x00\x00\x00\x00\x00\x009@\x1a\x1b\x1c\x1d\x1e\x1f ' \
|
||||
b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00' \
|
||||
b'\x00\x80'
|
||||
|
||||
with BinaryReader(data) as reader:
|
||||
value = reader.read_byte()
|
||||
assert value == 1, 'Example byte should be 1 but is {}'.format(
|
||||
value)
|
||||
assert value == 1, 'Example byte should be 1 but is {}'.format(value)
|
||||
|
||||
value = reader.read_int()
|
||||
assert value == 5, 'Example integer should be 5 but is {}'.format(
|
||||
value)
|
||||
assert value == 5, 'Example integer should be 5 but is {}'.format(value)
|
||||
|
||||
value = reader.read_long()
|
||||
assert value == 13, 'Example long integer should be 13 but is {}'.format(
|
||||
value)
|
||||
assert value == 13, 'Example long integer should be 13 but is {}'.format(value)
|
||||
|
||||
value = reader.read_float()
|
||||
assert value == 17.0, 'Example float should be 17.0 but is {}'.format(
|
||||
value)
|
||||
assert value == 17.0, 'Example float should be 17.0 but is {}'.format(value)
|
||||
|
||||
value = reader.read_double()
|
||||
assert value == 25.0, 'Example double should be 25.0 but is {}'.format(
|
||||
value)
|
||||
assert value == 25.0, 'Example double should be 25.0 but is {}'.format(value)
|
||||
|
||||
value = reader.read(7)
|
||||
assert value == bytes([26, 27, 28, 29, 30, 31, 32]), 'Example bytes should be {} but is {}' \
|
||||
.format(bytes([26, 27, 28, 29, 30, 31, 32]), value)
|
||||
|
||||
value = reader.read_large_int(128, signed=False)
|
||||
assert value == 2**127, 'Example large integer should be {} but is {}'.format(
|
||||
2**127, value)
|
||||
|
||||
# Test Telegram that types are written right
|
||||
with BinaryWriter() as writer:
|
||||
writer.write_int(0x60469778)
|
||||
buffer = writer.get_bytes()
|
||||
valid = b'\x78\x97\x46\x60' # Tested written bytes using C#'s MemoryStream
|
||||
|
||||
assert buffer == valid, 'Written type should be {} but is {}'.format(
|
||||
list(valid), list(buffer))
|
||||
assert value == 2**127, 'Example large integer should be {} but is {}'.format(2**127, value)
|
||||
|
||||
@staticmethod
|
||||
def test_binary_tgwriter_tgreader():
|
||||
small_data = os.urandom(33)
|
||||
small_data_padded = os.urandom(
|
||||
19) # +1 byte for length = 20 (evenly divisible by 4)
|
||||
small_data_padded = os.urandom(19) # +1 byte for length = 20 (%4 = 0)
|
||||
|
||||
large_data = os.urandom(999)
|
||||
large_data_padded = os.urandom(1024)
|
||||
|
||||
data = (small_data, small_data_padded, large_data, large_data_padded)
|
||||
string = 'Testing Telegram strings, this should work properly!'
|
||||
serialized = b''.join(TLObject.serialize_bytes(d) for d in data) + \
|
||||
TLObject.serialize_bytes(string)
|
||||
|
||||
with BinaryWriter() as writer:
|
||||
# First write the data
|
||||
with BinaryReader(serialized) as reader:
|
||||
# And then try reading it without errors (it should be unharmed!)
|
||||
for datum in data:
|
||||
writer.tgwrite_bytes(datum)
|
||||
writer.tgwrite_string(string)
|
||||
value = reader.tgread_bytes()
|
||||
assert value == datum, 'Example bytes should be {} but is {}'.format(
|
||||
datum, value)
|
||||
|
||||
with BinaryReader(writer.get_bytes()) as reader:
|
||||
# And then try reading it without errors (it should be unharmed!)
|
||||
for datum in data:
|
||||
value = reader.tgread_bytes()
|
||||
assert value == datum, 'Example bytes should be {} but is {}'.format(
|
||||
datum, value)
|
||||
|
||||
value = reader.tgread_string()
|
||||
assert value == string, 'Example string should be {} but is {}'.format(
|
||||
string, value)
|
||||
value = reader.tgread_string()
|
||||
assert value == string, 'Example string should be {} but is {}'.format(
|
||||
string, value)
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user