|
# copied from https://github.com/tulir/mautrix-telegram/blob/master/mautrix_telegram/util/parallel_file_transfer.py |
|
# Copyright (C) 2021 Tulir Asokan |
|
import asyncio |
|
import hashlib |
|
import inspect |
|
import logging |
|
import math |
|
import os |
|
from collections import defaultdict |
|
from typing import Optional, List, AsyncGenerator, Union, Awaitable, DefaultDict, Tuple, BinaryIO |
|
|
|
from telethon import utils, helpers, TelegramClient |
|
from telethon.crypto import AuthKey |
|
from telethon.network import MTProtoSender |
|
from telethon.tl.alltlobjects import LAYER |
|
from telethon.tl.functions import InvokeWithLayerRequest |
|
from telethon.tl.functions.auth import ExportAuthorizationRequest, ImportAuthorizationRequest |
|
from telethon.tl.functions.upload import (GetFileRequest, SaveFilePartRequest, |
|
SaveBigFilePartRequest) |
|
from telethon.tl.types import (Document, InputFileLocation, InputDocumentFileLocation, |
|
InputPhotoFileLocation, InputPeerPhotoFileLocation, TypeInputFile, |
|
InputFileBig, InputFile) |
|
|
|
try: |
|
from mautrix.crypto.attachments import async_encrypt_attachment |
|
except ImportError: |
|
async_encrypt_attachment = None |
|
|
|
log: logging.Logger = logging.getLogger("telethon") |
|
|
|
TypeLocation = Union[Document, InputDocumentFileLocation, InputPeerPhotoFileLocation, |
|
InputFileLocation, InputPhotoFileLocation] |
|
|
|
|
|
class DownloadSender: |
|
client: TelegramClient |
|
sender: MTProtoSender |
|
request: GetFileRequest |
|
remaining: int |
|
stride: int |
|
|
|
def __init__(self, client: TelegramClient, sender: MTProtoSender, file: TypeLocation, offset: int, limit: int, |
|
stride: int, count: int) -> None: |
|
self.sender = sender |
|
self.client = client |
|
self.request = GetFileRequest(file, offset=offset, limit=limit) |
|
self.stride = stride |
|
self.remaining = count |
|
|
|
async def next(self) -> Optional[bytes]: |
|
if not self.remaining: |
|
return None |
|
result = await self.client._call(self.sender, self.request) |
|
self.remaining -= 1 |
|
self.request.offset += self.stride |
|
return result.bytes |
|
|
|
def disconnect(self) -> Awaitable[None]: |
|
return self.sender.disconnect() |
|
|
|
|
|
class UploadSender: |
|
client: TelegramClient |
|
sender: MTProtoSender |
|
request: Union[SaveFilePartRequest, SaveBigFilePartRequest] |
|
part_count: int |
|
stride: int |
|
previous: Optional[asyncio.Task] |
|
loop: asyncio.AbstractEventLoop |
|
|
|
def __init__(self, client: TelegramClient, sender: MTProtoSender, file_id: int, part_count: int, big: bool, |
|
index: int, |
|
stride: int, loop: asyncio.AbstractEventLoop) -> None: |
|
self.client = client |
|
self.sender = sender |
|
self.part_count = part_count |
|
if big: |
|
self.request = SaveBigFilePartRequest(file_id, index, part_count, b"") |
|
else: |
|
self.request = SaveFilePartRequest(file_id, index, b"") |
|
self.stride = stride |
|
self.previous = None |
|
self.loop = loop |
|
|
|
async def next(self, data: bytes) -> None: |
|
if self.previous: |
|
await self.previous |
|
self.previous = self.loop.create_task(self._next(data)) |
|
|
|
async def _next(self, data: bytes) -> None: |
|
self.request.bytes = data |
|
log.debug(f"Sending file part {self.request.file_part}/{self.part_count}" |
|
f" with {len(data)} bytes") |
|
await self.client._call(self.sender, self.request) |
|
self.request.file_part += self.stride |
|
|
|
async def disconnect(self) -> None: |
|
if self.previous: |
|
await self.previous |
|
return await self.sender.disconnect() |
|
|
|
|
|
class ParallelTransferrer: |
|
client: TelegramClient |
|
loop: asyncio.AbstractEventLoop |
|
dc_id: int |
|
senders: Optional[List[Union[DownloadSender, UploadSender]]] |
|
auth_key: AuthKey |
|
upload_ticker: int |
|
|
|
def __init__(self, client: TelegramClient, dc_id: Optional[int] = None) -> None: |
|
self.client = client |
|
self.loop = self.client.loop |
|
self.dc_id = dc_id or self.client.session.dc_id |
|
self.auth_key = (None if dc_id and self.client.session.dc_id != dc_id |
|
else self.client.session.auth_key) |
|
self.senders = None |
|
self.upload_ticker = 0 |
|
|
|
async def _cleanup(self) -> None: |
|
await asyncio.gather(*[sender.disconnect() for sender in self.senders]) |
|
self.senders = None |
|
|
|
@staticmethod |
|
def _get_connection_count(file_size: int, max_count: int = 20, |
|
full_size: int = 100 * 1024 * 1024) -> int: |
|
if file_size > full_size: |
|
return max_count |
|
return math.ceil((file_size / full_size) * max_count) |
|
|
|
async def _init_download(self, connections: int, file: TypeLocation, part_count: int, |
|
part_size: int) -> None: |
|
minimum, remainder = divmod(part_count, connections) |
|
|
|
def get_part_count() -> int: |
|
nonlocal remainder |
|
if remainder > 0: |
|
remainder -= 1 |
|
return minimum + 1 |
|
return minimum |
|
|
|
# The first cross-DC sender will export+import the authorization, so we always create it |
|
# before creating any other senders. |
|
self.senders = [ |
|
await self._create_download_sender(file, 0, part_size, connections * part_size, |
|
get_part_count()), |
|
*await asyncio.gather( |
|
*[self._create_download_sender(file, i, part_size, connections * part_size, |
|
get_part_count()) |
|
for i in range(1, connections)]) |
|
] |
|
|
|
async def _create_download_sender(self, file: TypeLocation, index: int, part_size: int, |
|
stride: int, |
|
part_count: int) -> DownloadSender: |
|
return DownloadSender(self.client, await self._create_sender(), file, index * part_size, part_size, |
|
stride, part_count) |
|
|
|
async def _init_upload(self, connections: int, file_id: int, part_count: int, big: bool |
|
) -> None: |
|
self.senders = [ |
|
await self._create_upload_sender(file_id, part_count, big, 0, connections), |
|
*await asyncio.gather( |
|
*[self._create_upload_sender(file_id, part_count, big, i, connections) |
|
for i in range(1, connections)]) |
|
] |
|
|
|
async def _create_upload_sender(self, file_id: int, part_count: int, big: bool, index: int, |
|
stride: int) -> UploadSender: |
|
return UploadSender(self.client, await self._create_sender(), file_id, part_count, big, index, stride, |
|
loop=self.loop) |
|
|
|
async def _create_sender(self) -> MTProtoSender: |
|
dc = await self.client._get_dc(self.dc_id) |
|
sender = MTProtoSender(self.auth_key, loggers=self.client._log) |
|
await sender.connect(self.client._connection(dc.ip_address, dc.port, dc.id, |
|
loggers=self.client._log, |
|
proxy=self.client._proxy)) |
|
if not self.auth_key: |
|
log.debug(f"Exporting auth to DC {self.dc_id}") |
|
auth = await self.client(ExportAuthorizationRequest(self.dc_id)) |
|
self.client._init_request.query = ImportAuthorizationRequest(id=auth.id, |
|
bytes=auth.bytes) |
|
req = InvokeWithLayerRequest(LAYER, self.client._init_request) |
|
await sender.send(req) |
|
self.auth_key = sender.auth_key |
|
return sender |
|
|
|
async def init_upload(self, file_id: int, file_size: int, part_size_kb: Optional[float] = None, |
|
connection_count: Optional[int] = None) -> Tuple[int, int, bool]: |
|
connection_count = connection_count or self._get_connection_count(file_size) |
|
part_size = (part_size_kb or utils.get_appropriated_part_size(file_size)) * 1024 |
|
part_count = (file_size + part_size - 1) // part_size |
|
is_large = file_size > 10 * 1024 * 1024 |
|
await self._init_upload(connection_count, file_id, part_count, is_large) |
|
return part_size, part_count, is_large |
|
|
|
async def upload(self, part: bytes) -> None: |
|
await self.senders[self.upload_ticker].next(part) |
|
self.upload_ticker = (self.upload_ticker + 1) % len(self.senders) |
|
|
|
async def finish_upload(self) -> None: |
|
await self._cleanup() |
|
|
|
async def download(self, file: TypeLocation, file_size: int, |
|
part_size_kb: Optional[float] = None, |
|
connection_count: Optional[int] = None) -> AsyncGenerator[bytes, None]: |
|
connection_count = connection_count or self._get_connection_count(file_size) |
|
part_size = (part_size_kb or utils.get_appropriated_part_size(file_size)) * 1024 |
|
part_count = math.ceil(file_size / part_size) |
|
log.debug("Starting parallel download: " |
|
f"{connection_count} {part_size} {part_count} {file!s}") |
|
await self._init_download(connection_count, file, part_count, part_size) |
|
|
|
part = 0 |
|
while part < part_count: |
|
tasks = [] |
|
for sender in self.senders: |
|
tasks.append(self.loop.create_task(sender.next())) |
|
for task in tasks: |
|
data = await task |
|
if not data: |
|
break |
|
yield data |
|
part += 1 |
|
log.debug(f"Part {part} downloaded") |
|
|
|
log.debug("Parallel download finished, cleaning up connections") |
|
await self._cleanup() |
|
|
|
|
|
parallel_transfer_locks: DefaultDict[int, asyncio.Lock] = defaultdict(lambda: asyncio.Lock()) |
|
|
|
|
|
def stream_file(file_to_stream: BinaryIO, chunk_size=1024): |
|
while True: |
|
data_read = file_to_stream.read(chunk_size) |
|
if not data_read: |
|
break |
|
yield data_read |
|
|
|
|
|
async def _internal_transfer_to_telegram(client: TelegramClient, |
|
response: BinaryIO, |
|
progress_callback: callable |
|
) -> Tuple[TypeInputFile, int]: |
|
file_id = helpers.generate_random_long() |
|
file_size = os.path.getsize(response.name) |
|
|
|
hash_md5 = hashlib.md5() |
|
uploader = ParallelTransferrer(client) |
|
part_size, part_count, is_large = await uploader.init_upload(file_id, file_size) |
|
buffer = bytearray() |
|
for data in stream_file(response): |
|
if progress_callback: |
|
r = progress_callback(response.tell(), file_size) |
|
if inspect.isawaitable(r): |
|
await r |
|
if not is_large: |
|
hash_md5.update(data) |
|
if len(buffer) == 0 and len(data) == part_size: |
|
await uploader.upload(data) |
|
continue |
|
new_len = len(buffer) + len(data) |
|
if new_len >= part_size: |
|
cutoff = part_size - len(buffer) |
|
buffer.extend(data[:cutoff]) |
|
await uploader.upload(bytes(buffer)) |
|
buffer.clear() |
|
buffer.extend(data[cutoff:]) |
|
else: |
|
buffer.extend(data) |
|
if len(buffer) > 0: |
|
await uploader.upload(bytes(buffer)) |
|
await uploader.finish_upload() |
|
if is_large: |
|
return InputFileBig(file_id, part_count, "upload"), file_size |
|
else: |
|
return InputFile(file_id, part_count, "upload", hash_md5.hexdigest()), file_size |
|
|
|
|
|
async def download_file(client: TelegramClient, |
|
location: TypeLocation, |
|
out: BinaryIO, |
|
progress_callback: callable = None |
|
) -> BinaryIO: |
|
size = location.size |
|
dc_id, location = utils.get_input_location(location) |
|
# We lock the transfers because telegram has connection count limits |
|
downloader = ParallelTransferrer(client, dc_id) |
|
downloaded = downloader.download(location, size) |
|
async for x in downloaded: |
|
out.write(x) |
|
if progress_callback: |
|
r = progress_callback(out.tell(), size) |
|
if inspect.isawaitable(r): |
|
await r |
|
|
|
return out |
|
|
|
|
|
async def upload_file(client: TelegramClient, |
|
file: BinaryIO, |
|
progress_callback: callable = None, |
|
|
|
) -> TypeInputFile: |
|
res = (await _internal_transfer_to_telegram(client, file, progress_callback))[0] |
|
return res |