mirror of
https://github.com/LonamiWebs/Telethon.git
synced 2025-08-08 12:59:46 +00:00
Slightly reorganise the project structure
This commit is contained in:
Lonami Exo
9
telethon/extensions/__init__.py
Normal file
9
telethon/extensions/__init__.py
Normal file
@@ -0,0 +1,9 @@
|
||||
"""
|
||||
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.)
|
||||
"""
|
||||
from .binary_writer import BinaryWriter
|
||||
from .binary_reader import BinaryReader
|
||||
from .tcp_client import TcpClient
|
||||
168
telethon/extensions/binary_reader.py
Normal file
168
telethon/extensions/binary_reader.py
Normal file
@@ -0,0 +1,168 @@
|
||||
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
|
||||
140
telethon/extensions/binary_writer.py
Normal file
140
telethon/extensions/binary_writer.py
Normal file
@@ -0,0 +1,140 @@
|
||||
from io import BufferedWriter, BytesIO
|
||||
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.writer = BufferedWriter(stream)
|
||||
self.written_count = 0
|
||||
|
||||
# region Writing
|
||||
|
||||
# "All numbers are written as little endian." |> Source: 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), which can or cannot be 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), which can or cannot be 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"""
|
||||
tlobject.on_send(self)
|
||||
|
||||
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()
|
||||
116
telethon/extensions/tcp_client.py
Normal file
116
telethon/extensions/tcp_client.py
Normal file
@@ -0,0 +1,116 @@
|
||||
# Python rough implementation of a C# TCP client
|
||||
import socket
|
||||
import time
|
||||
from datetime import datetime, timedelta
|
||||
from io import BytesIO, BufferedWriter
|
||||
from threading import Event, Lock
|
||||
|
||||
from ..errors import ReadCancelledError
|
||||
|
||||
|
||||
class TcpClient:
|
||||
def __init__(self, proxy=None):
|
||||
self.connected = False
|
||||
self._proxy = proxy
|
||||
self._recreate_socket()
|
||||
|
||||
# Support for multi-threading advantages and safety
|
||||
self.cancelled = Event() # Has the read operation been cancelled?
|
||||
self.delay = 0.1 # Read delay when there was no data available
|
||||
self._lock = Lock()
|
||||
|
||||
def _recreate_socket(self):
|
||||
if self._proxy is None:
|
||||
self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
else:
|
||||
import socks
|
||||
self._socket = socks.socksocket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
if type(self._proxy) is dict:
|
||||
self._socket.set_proxy(**self._proxy)
|
||||
else: # tuple, list, etc.
|
||||
self._socket.set_proxy(*self._proxy)
|
||||
|
||||
def connect(self, ip, port):
|
||||
"""Connects to the specified IP and port number"""
|
||||
if not self.connected:
|
||||
self._socket.connect((ip, port))
|
||||
self._socket.setblocking(False)
|
||||
self.connected = True
|
||||
|
||||
def close(self):
|
||||
"""Closes the connection"""
|
||||
if self.connected:
|
||||
self._socket.shutdown(socket.SHUT_RDWR)
|
||||
self._socket.close()
|
||||
self.connected = False
|
||||
self._recreate_socket()
|
||||
|
||||
def write(self, data):
|
||||
"""Writes (sends) the specified bytes to the connected peer"""
|
||||
|
||||
# Ensure that only one thread can send data at once
|
||||
with self._lock:
|
||||
view = memoryview(data)
|
||||
total_sent, total = 0, len(data)
|
||||
while total_sent < total:
|
||||
sent = self._socket.send(view[total_sent:])
|
||||
if sent == 0:
|
||||
raise ConnectionResetError(
|
||||
'The server has closed the connection.')
|
||||
total_sent += sent
|
||||
|
||||
def read(self, size, timeout=timedelta(seconds=5)):
|
||||
"""Reads (receives) a whole block of 'size bytes
|
||||
from the connected peer.
|
||||
|
||||
A timeout can be specified, which will cancel the operation if
|
||||
no data has been read in the specified time. If data was read
|
||||
and it's waiting for more, the timeout will NOT cancel the
|
||||
operation. Set to None for no timeout
|
||||
"""
|
||||
|
||||
# Ensure that only one thread can receive data at once
|
||||
with self._lock:
|
||||
# Ensure it is not cancelled at first, so we can enter the loop
|
||||
self.cancelled.clear()
|
||||
|
||||
# Set the starting time so we can
|
||||
# calculate whether the timeout should fire
|
||||
start_time = datetime.now() if timeout else None
|
||||
|
||||
with BufferedWriter(BytesIO(), buffer_size=size) as buffer:
|
||||
bytes_left = size
|
||||
while bytes_left != 0:
|
||||
# Only do cancel if no data was read yet
|
||||
# Otherwise, carry on reading and finish
|
||||
if self.cancelled.is_set() and bytes_left == size:
|
||||
raise ReadCancelledError()
|
||||
|
||||
try:
|
||||
partial = self._socket.recv(bytes_left)
|
||||
if len(partial) == 0:
|
||||
raise ConnectionResetError(
|
||||
'The server has closed the connection (recv() returned 0 bytes).')
|
||||
|
||||
buffer.write(partial)
|
||||
bytes_left -= len(partial)
|
||||
|
||||
except BlockingIOError as error:
|
||||
# No data available yet, sleep a bit
|
||||
time.sleep(self.delay)
|
||||
|
||||
# Check if the timeout finished
|
||||
if timeout:
|
||||
time_passed = datetime.now() - start_time
|
||||
if time_passed > timeout:
|
||||
raise TimeoutError(
|
||||
'The read operation exceeded the timeout.') from error
|
||||
|
||||
# If everything went fine, return the read bytes
|
||||
buffer.flush()
|
||||
return buffer.raw.getvalue()
|
||||
|
||||
def cancel_read(self):
|
||||
"""Cancels the read operation IF it hasn't yet
|
||||
started, raising a ReadCancelledError"""
|
||||
self.cancelled.set()
|
||||
Reference in New Issue
Block a user