| | import asyncio |
| | import logging |
| | from typing import Dict, Union |
| | from FileStream.bot import work_loads |
| | from pyrogram import Client, utils, raw |
| | from .file_properties import get_file_ids |
| | from pyrogram.session import Session, Auth |
| | from pyrogram.errors import AuthBytesInvalid |
| | from pyrogram.file_id import FileId, FileType, ThumbnailSource |
| | from pyrogram.types import Message |
| |
|
| | class ByteStreamer: |
| | def __init__(self, client: Client): |
| | self.clean_timer = 30 * 60 |
| | self.client: Client = client |
| | self.cached_file_ids: Dict[str, FileId] = {} |
| | asyncio.create_task(self.clean_cache()) |
| |
|
| | async def get_file_properties(self, db_id: str, multi_clients) -> FileId: |
| | """ |
| | Returns the properties of a media of a specific message in a FIleId class. |
| | if the properties are cached, then it'll return the cached results. |
| | or it'll generate the properties from the Message ID and cache them. |
| | """ |
| | if not db_id in self.cached_file_ids: |
| | logging.debug("Before Calling generate_file_properties") |
| | await self.generate_file_properties(db_id, multi_clients) |
| | logging.debug(f"Cached file properties for file with ID {db_id}") |
| | return self.cached_file_ids[db_id] |
| | |
| | async def generate_file_properties(self, db_id: str, multi_clients) -> FileId: |
| | """ |
| | Generates the properties of a media file on a specific message. |
| | returns ths properties in a FIleId class. |
| | """ |
| | logging.debug("Before calling get_file_ids") |
| | file_id = await get_file_ids(self.client, db_id, multi_clients, Message) |
| | logging.debug(f"Generated file ID and Unique ID for file with ID {db_id}") |
| | self.cached_file_ids[db_id] = file_id |
| | logging.debug(f"Cached media file with ID {db_id}") |
| | return self.cached_file_ids[db_id] |
| |
|
| | async def generate_media_session(self, client: Client, file_id: FileId) -> Session: |
| | """ |
| | Generates the media session for the DC that contains the media file. |
| | This is required for getting the bytes from Telegram servers. |
| | """ |
| |
|
| | media_session = client.media_sessions.get(file_id.dc_id, None) |
| |
|
| | if media_session is None: |
| | if file_id.dc_id != await client.storage.dc_id(): |
| | media_session = Session( |
| | client, |
| | file_id.dc_id, |
| | await Auth( |
| | client, file_id.dc_id, await client.storage.test_mode() |
| | ).create(), |
| | await client.storage.test_mode(), |
| | is_media=True, |
| | ) |
| | await media_session.start() |
| |
|
| | for _ in range(6): |
| | exported_auth = await client.invoke( |
| | raw.functions.auth.ExportAuthorization(dc_id=file_id.dc_id) |
| | ) |
| |
|
| | try: |
| | await media_session.invoke( |
| | raw.functions.auth.ImportAuthorization( |
| | id=exported_auth.id, bytes=exported_auth.bytes |
| | ) |
| | ) |
| | break |
| | except AuthBytesInvalid: |
| | logging.debug( |
| | f"Invalid authorization bytes for DC {file_id.dc_id}" |
| | ) |
| | continue |
| | else: |
| | await media_session.stop() |
| | raise AuthBytesInvalid |
| | else: |
| | media_session = Session( |
| | client, |
| | file_id.dc_id, |
| | await client.storage.auth_key(), |
| | await client.storage.test_mode(), |
| | is_media=True, |
| | ) |
| | await media_session.start() |
| | logging.debug(f"Created media session for DC {file_id.dc_id}") |
| | client.media_sessions[file_id.dc_id] = media_session |
| | else: |
| | logging.debug(f"Using cached media session for DC {file_id.dc_id}") |
| | return media_session |
| |
|
| |
|
| | @staticmethod |
| | async def get_location(file_id: FileId) -> Union[raw.types.InputPhotoFileLocation, |
| | raw.types.InputDocumentFileLocation, |
| | raw.types.InputPeerPhotoFileLocation,]: |
| | """ |
| | Returns the file location for the media file. |
| | """ |
| | file_type = file_id.file_type |
| |
|
| | if file_type == FileType.CHAT_PHOTO: |
| | if file_id.chat_id > 0: |
| | peer = raw.types.InputPeerUser( |
| | user_id=file_id.chat_id, access_hash=file_id.chat_access_hash |
| | ) |
| | else: |
| | if file_id.chat_access_hash == 0: |
| | peer = raw.types.InputPeerChat(chat_id=-file_id.chat_id) |
| | else: |
| | peer = raw.types.InputPeerChannel( |
| | channel_id=utils.get_channel_id(file_id.chat_id), |
| | access_hash=file_id.chat_access_hash, |
| | ) |
| |
|
| | location = raw.types.InputPeerPhotoFileLocation( |
| | peer=peer, |
| | volume_id=file_id.volume_id, |
| | local_id=file_id.local_id, |
| | big=file_id.thumbnail_source == ThumbnailSource.CHAT_PHOTO_BIG, |
| | ) |
| | elif file_type == FileType.PHOTO: |
| | location = raw.types.InputPhotoFileLocation( |
| | id=file_id.media_id, |
| | access_hash=file_id.access_hash, |
| | file_reference=file_id.file_reference, |
| | thumb_size=file_id.thumbnail_size, |
| | ) |
| | else: |
| | location = raw.types.InputDocumentFileLocation( |
| | id=file_id.media_id, |
| | access_hash=file_id.access_hash, |
| | file_reference=file_id.file_reference, |
| | thumb_size=file_id.thumbnail_size, |
| | ) |
| | return location |
| |
|
| | async def yield_file( |
| | self, |
| | file_id: FileId, |
| | index: int, |
| | offset: int, |
| | first_part_cut: int, |
| | last_part_cut: int, |
| | part_count: int, |
| | chunk_size: int, |
| | ) -> Union[str, None]: |
| | """ |
| | Custom generator that yields the bytes of the media file. |
| | Modded from <https://github.com/eyaadh/megadlbot_oss/blob/master/mega/telegram/utils/custom_download.py#L20> |
| | Thanks to Eyaadh <https://github.com/eyaadh> |
| | """ |
| | client = self.client |
| | work_loads[index] += 1 |
| | logging.debug(f"Starting to yielding file with client {index}.") |
| | media_session = await self.generate_media_session(client, file_id) |
| |
|
| | current_part = 1 |
| |
|
| | location = await self.get_location(file_id) |
| |
|
| | try: |
| | r = await media_session.invoke( |
| | raw.functions.upload.GetFile( |
| | location=location, offset=offset, limit=chunk_size |
| | ), |
| | ) |
| | if isinstance(r, raw.types.upload.File): |
| | while True: |
| | chunk = r.bytes |
| | if not chunk: |
| | break |
| | elif part_count == 1: |
| | yield chunk[first_part_cut:last_part_cut] |
| | elif current_part == 1: |
| | yield chunk[first_part_cut:] |
| | elif current_part == part_count: |
| | yield chunk[:last_part_cut] |
| | else: |
| | yield chunk |
| |
|
| | current_part += 1 |
| | offset += chunk_size |
| |
|
| | if current_part > part_count: |
| | break |
| |
|
| | r = await media_session.invoke( |
| | raw.functions.upload.GetFile( |
| | location=location, offset=offset, limit=chunk_size |
| | ), |
| | ) |
| | except (TimeoutError, AttributeError): |
| | pass |
| | finally: |
| | logging.debug(f"Finished yielding file with {current_part} parts.") |
| | work_loads[index] -= 1 |
| |
|
| | |
| | async def clean_cache(self) -> None: |
| | """ |
| | function to clean the cache to reduce memory usage |
| | """ |
| | while True: |
| | await asyncio.sleep(self.clean_timer) |
| | self.cached_file_ids.clear() |
| | logging.debug("Cleaned the cache") |
| |
|
