"""Pyrogram ext""" import asyncio import os import secrets import struct import time from datetime import datetime from functools import wraps from io import BytesIO, StringIO from mimetypes import MimeTypes from typing import Callable, Iterable, List, Optional, Union import pyrogram from loguru import logger from pyrogram import types from pyrogram.client import Cache from pyrogram.file_id import ( FILE_REFERENCE_FLAG, PHOTO_TYPES, WEB_LOCATION_FLAG, FileType, b64_decode, rle_decode, ) from pyrogram.mime_types import mime_types from module.app import ( Application, CloudDriveUploadStat, DownloadStatus, ForwardStatus, TaskNode, UploadProgressStat, UploadStatus, ) from module.download_stat import get_download_result from module.language import Language, _t from module.send_media_group_v2 import cache_media, send_media_group_v2 from utils.format import ( create_progress_bar, extract_info_from_link, format_byte, truncate_filename, ) from utils.meta_data import MetaData _mimetypes = MimeTypes() _mimetypes.readfp(StringIO(mime_types)) _download_cache = Cache(1024 * 1024 * 1024) def reset_download_cache(): """Reset download cache""" _download_cache.store.clear() def _guess_mime_type(filename: str) -> Optional[str]: """Guess mime type""" return _mimetypes.guess_type(filename)[0] def _guess_extension(mime_type: str) -> Optional[str]: """Guess extension""" return _mimetypes.guess_extension(mime_type) def get_media_obj( message: pyrogram.types.Message, media: str = None, caption: str = None ) -> Union[ types.InputMediaPhoto, types.InputMediaVideo, types.InputMediaAudio, types.InputMediaDocument, types.InputMediaAnimation, ]: """Get media object""" media_type = message.media if media_type == pyrogram.enums.MessageMediaType.PHOTO: return types.InputMediaPhoto(media, caption=caption) if media_type == pyrogram.enums.MessageMediaType.VIDEO: return types.InputMediaVideo( media, caption=caption, width=message.video.width, height=message.video.height, duration=message.video.duration, ) if media_type in [ pyrogram.enums.MessageMediaType.AUDIO, pyrogram.enums.MessageMediaType.VOICE, ]: return types.InputMediaAudio(media, caption=caption) if media_type == pyrogram.enums.MessageMediaType.DOCUMENT: return types.InputMediaDocument(media, caption=caption) if media_type == pyrogram.enums.MessageMediaType.ANIMATION: return types.InputMediaAnimation(media, caption=caption) return None def _get_file_type(file_id: str): """Get file type""" decoded = rle_decode(b64_decode(file_id)) # File id versioning. Major versions lower than 4 don't have a minor version major = decoded[-1] if major < 4: buffer = BytesIO(decoded[:-1]) else: buffer = BytesIO(decoded[:-2]) file_type, _ = struct.unpack(" str: """Get extension""" if not file_id: if dot: return ".unknown" return "unknown" file_type = _get_file_type(file_id) guessed_extension = _guess_extension(mime_type) if file_type in PHOTO_TYPES: extension = "jpg" elif file_type == FileType.VOICE: extension = guessed_extension or "ogg" elif file_type in (FileType.VIDEO, FileType.ANIMATION, FileType.VIDEO_NOTE): extension = guessed_extension or "mp4" elif file_type == FileType.DOCUMENT: extension = guessed_extension or "zip" elif file_type == FileType.STICKER: extension = guessed_extension or "webp" elif file_type == FileType.AUDIO: extension = guessed_extension or "mp3" else: extension = "unknown" if dot: extension = "." + extension return extension async def send_message_by_language( client: pyrogram.client.Client, language: Language, chat_id: Union[int, str], reply_to_message_id: int, language_str: List[str], ): """Record download status""" msg = language_str[language.value - 1] return await client.send_message( chat_id, msg, reply_to_message_id=reply_to_message_id ) async def download_thumbnail( client: pyrogram.Client, temp_path: str, message: pyrogram.types.Message, ): """Downloads the thumbnail of a video message to a temporary file. Args: client: A Pyrogram client instance. temp_path: The path to a temporary directory where the thumbnail file will be stored. message: A Pyrogram Message object representing the video message. Returns: A string representing the path of the thumbnail file, or None if the download failed. Raises: ValueError: If the downloaded thumbnail file size doesn't match the expected file size. """ thumbnail_file = None if message.video.thumbs: message = await fetch_message(client, message) thumbnail = message.video.thumbs[0] if message.video.thumbs else None unique_name = os.path.join( temp_path, "thumbnail", f"thumb-{int(time.time())}-{secrets.token_hex(8)}.jpg", ) max_attempts = 3 for attempt in range(1, max_attempts + 1): try: thumbnail_file = await client.download_media( thumbnail, file_name=unique_name ) if os.path.getsize(thumbnail_file) == thumbnail.file_size: break raise ValueError( f"Thumbnail file size is {os.path.getsize(thumbnail_file)}" f" bytes, actual {thumbnail.file_size}: {thumbnail_file}" ) except Exception as e: if attempt == max_attempts: logger.exception( f"Failed to download thumbnail after {max_attempts}" f" attempts: {e}" ) else: message = await fetch_message(client, message) logger.warning( f"Attempt {attempt} to download thumbnail failed: {e}" ) # Wait 2 seconds before retrying await asyncio.sleep(2) thumbnail = None thumbnail_file = None return thumbnail_file async def upload_telegram_chat( client: pyrogram.Client, upload_user: pyrogram.Client, app: Application, node: TaskNode, message: pyrogram.types.Message, download_status: DownloadStatus, file_name: str = None, ): """Upload telegram chat""" # upload telegram if node.upload_telegram_chat_id: if download_status is DownloadStatus.SkipDownload and message.media: if message.media_group_id: await proc_cache_forward(client, node, message, True) return if download_status is DownloadStatus.SuccessDownload or ( download_status is DownloadStatus.SkipDownload and not message.media ): try: await upload_telegram_chat_message( client, upload_user, app, node, message, file_name, ) except Exception as e: logger.exception(f"Upload file {file_name} error: {e}") finally: if file_name and app.after_upload_telegram_delete: os.remove(file_name) # forward text # FIXME: fix upload text # if ( # download_status is DownloadStatus.SkipDownload # and message.text # and bot # ): # await upload_telegram_chat( # client, app, node.upload_telegram_chat_id, message, file_name # ) async def upload_telegram_chat_message( client: pyrogram.Client, upload_user: pyrogram.Client, app: Application, node: TaskNode, message: pyrogram.types.Message, file_name: str = None, ) -> ForwardStatus: """See upload telegram_chat""" forward_status = ForwardStatus.FailedForward max_attempts = 3 for _ in range(1, max_attempts + 1): try: forward_status = await _upload_telegram_chat_message( client, upload_user, app, node, message, file_name ) break except pyrogram.errors.exceptions.flood_420.FloodWait as wait_err: await asyncio.sleep(wait_err.value * 2) logger.warning( "Upload Message[{}]: FlowWait {}", message.id, wait_err.value ) except Exception as e: logger.exception(f"Upload file {file_name} error: {e}") return ForwardStatus.FailedForward if forward_status != ForwardStatus.CacheForward: node.stat_forward(forward_status) return forward_status # pylint: disable=R0912 async def _upload_signal_message( client: pyrogram.Client, upload_user: pyrogram.Client, app: Application, node: TaskNode, upload_telegram_chat_id: Union[int, str, None], message: pyrogram.types.Message, file_name: Optional[str], caption: Optional[str] = None, ): """ Uploads a video or message to a Telegram chat. Parameters: client (pyrogram.Client): The pyrogram client. upload_telegram_chat_id (Union[int, str]): The ID of the chat to upload to. message (pyrogram.types.Message): The message to upload. file_name (str): The name of the file to upload. """ ui_file_name = file_name if file_name: ui_file_name = ( f"****{os.path.splitext(file_name)[-1]}" if app.hide_file_name else file_name ) if message.video: # Download thumbnail thumbnail_file = await download_thumbnail(client, app.temp_save_path, message) try: # TODO(tangyoha): add more log when upload video more than 2000MB failed # Send video to the destination chat if node.reply_to_message: await node.reply_to_message.reply_video( file_name, caption=caption, message_thread_id=node.topic_id, thumb=thumbnail_file, width=message.video.width, height=message.video.height, duration=message.video.duration, parse_mode=pyrogram.enums.ParseMode.HTML, ) else: await upload_user.send_video( upload_telegram_chat_id, file_name, thumb=thumbnail_file, width=message.video.width, height=message.video.height, duration=message.video.duration, caption=caption, parse_mode=pyrogram.enums.ParseMode.HTML, progress=update_upload_stat, progress_args=( message.id, ui_file_name, time.time(), node, upload_user, ), message_thread_id=node.topic_id, ) except Exception as e: raise e finally: if thumbnail_file: os.remove(str(thumbnail_file)) elif message.photo: if node.reply_to_message: await node.reply_to_message.reply_photo( file_name, caption=caption, message_thread_id=node.topic_id, ) else: await upload_user.send_photo( upload_telegram_chat_id, file_name, caption=caption, progress=update_upload_stat, progress_args=( message.id, ui_file_name, time.time(), node, upload_user, ), message_thread_id=node.topic_id, ) elif message.document: if node.reply_to_message: await node.reply_to_message.reply_document( file_name, caption=caption, message_thread_id=node.topic_id, ) else: await upload_user.send_document( upload_telegram_chat_id, file_name, caption=caption, progress=update_upload_stat, progress_args=( message.id, ui_file_name, time.time(), node, upload_user, ), message_thread_id=node.topic_id, ) elif message.voice: if node.reply_to_message: await node.reply_to_message.reply_voice( file_name, caption=caption, message_thread_id=node.topic_id, ) else: await upload_user.send_voice( upload_telegram_chat_id, file_name, caption=caption, progress=update_upload_stat, progress_args=( message.id, ui_file_name, time.time(), node, upload_user, ), message_thread_id=node.topic_id, ) elif message.video_note: if node.reply_to_message: await node.reply_to_message.reply_video_note( file_name, caption=caption, message_thread_id=node.topic_id, ) else: await upload_user.send_video_note( upload_telegram_chat_id, file_name, caption=caption, progress=update_upload_stat, progress_args=( message.id, ui_file_name, time.time(), node, upload_user, ), message_thread_id=node.topic_id, ) elif message.text: if node.reply_to_message: await node.reply_to_message.reply( message.text, message_thread_id=node.topic_id ) else: await upload_user.send_message( upload_telegram_chat_id, message.text, message_thread_id=node.topic_id, ) async def _upload_telegram_chat_message( client: pyrogram.Client, upload_user: pyrogram.Client, app: Application, node: TaskNode, message: pyrogram.types.Message, file_name: str = None, ): """ Uploads a Telegram chat message to the destination chat. Args: client (pyrogram.Client): The client used to interact with the Telegram API. upload_user (pyrogram.Client): The client used to upload the message. app (Application): The application instance. node (TaskNode): The task node associated with the message. message (pyrogram.types.Message): The Telegram chat message to be uploaded. file_name (str): The name of the file to be uploaded. Returns: None """ await app.forward_limit_call.wait(node) caption = message.caption caption_entities = message.caption_entities # Convert caption and caption_entities to markdown format if caption and caption_entities: caption = pyrogram.parser.Parser.unparse(caption, caption_entities, True) max_caption_length = 4096 if client.me and client.me.is_premium else 1024 # proc caption MEDIA_CAPTION_TOO_LONG if caption and len(caption) > max_caption_length: caption = caption[:max_caption_length] if not message.media_group_id: if not node.has_protected_content: if node.reply_to_message: if message.text: await node.reply_to_message.reply( message.text, message_thread_id=node.topic_id, ) elif message.photo: await node.reply_to_message.reply_photo( message.photo.file_id, caption=caption, message_thread_id=node.topic_id, ) elif message.video: await node.reply_to_message.reply_video( message.video.file_id, caption=caption, message_thread_id=node.topic_id, ) elif message.document: await node.reply_to_message.reply_document( message.document.file_id, caption=caption, message_thread_id=node.topic_id, ) elif message.audio: await node.reply_to_message.reply_audio( message.audio.file_id, caption=caption, message_thread_id=node.topic_id, ) else: # For other types of media, fallback to forward_messages await forward_messages( client, node.upload_telegram_chat_id, node.chat_id, message.id, drop_author=True, topic_id=node.topic_id, caption=caption, ) else: await _upload_signal_message( client, upload_user, app, node, node.upload_telegram_chat_id, message, file_name, caption, ) return ForwardStatus.SuccessForward return await forward_multi_media( client, upload_user, app, node, message, caption, file_name ) # pylint: disable=R0912 async def forward_multi_media( client: pyrogram.Client, _: pyrogram.Client, app: Application, node: TaskNode, message: pyrogram.types.Message, caption: str = None, file_name: str = None, ): """Forward multi media by cache""" caption = message.caption caption_entities = message.caption_entities if not caption: caption = app.get_caption_name(node.chat_id, message.media_group_id) caption_entities = app.get_caption_entities( node.chat_id, message.media_group_id ) # Convert caption and caption_entities to markdown format if caption and caption_entities: caption = pyrogram.parser.Parser.unparse(caption, caption_entities, True) max_caption_length = 4096 if client.me and client.me.is_premium else 1024 # proc caption MEDIA_CAPTION_TOO_LONG if caption and len(caption) > max_caption_length: caption = caption[:max_caption_length] media_obj = get_media_obj(message, file_name, caption) if not node.has_protected_content: media = getattr(message, message.media.value) if not media: return ForwardStatus.SkipForward media_obj.media = media.file_id if media else "" if not node.media_group_ids.get(message.media_group_id): node.media_group_ids[message.media_group_id] = {} if not node.media_group_ids[message.media_group_id]: media_group = await get_media_group_with_retry( client, node.chat_id, message.id, 5 ) if not media_group: logger.error("Get Media Group Error! message id: {}", message.id) return ForwardStatus.FailedForward for it in media_group: node.media_group_ids[message.media_group_id][it.id] = None node.upload_status[message.id] = None if not node.media_group_ids[message.media_group_id][message.id]: node.upload_status[message.id] = UploadStatus.Uploading try: ui_file_name = file_name if file_name: ui_file_name = ( f"****{os.path.splitext(file_name)[-1]}" if app.hide_file_name else file_name ) media_obj.thumb = ( await download_thumbnail(client, app.temp_save_path, message) if message.video else None ) _media = await cache_media( client, node.upload_telegram_chat_id, # type: ignore media_obj, progress=update_upload_stat, progress_args=( message.id, ui_file_name, time.time(), node, client, ), ) except Exception as e: logger.exception(f"{e}") node.upload_status[message.id] = UploadStatus.FailedUpload finally: if file_name and message.video and media_obj.thumb: os.remove(str(media_obj.thumb)) if node.upload_status[message.id] == UploadStatus.FailedUpload: return ForwardStatus.FailedForward node.media_group_ids[message.media_group_id][message.id] = _media node.upload_status[message.id] = UploadStatus.SuccessUpload return await proc_cache_forward(client, node, message, bool(file_name)) async def proc_cache_forward( client: pyrogram.Client, node: TaskNode, message: pyrogram.types.Message, check_download_status: bool, ): """proc other cache forward""" if not node.media_group_ids: return for key in node.media_group_ids[message.media_group_id].keys(): download_status = node.download_status.get(key, DownloadStatus.Downloading) if ( node.skip_msg_id(key) or download_status is DownloadStatus.SkipDownload or download_status is DownloadStatus.FailedDownload ): continue if ( check_download_status and DownloadStatus.Downloading == download_status ) or UploadStatus.Uploading == node.upload_status.get( key, UploadStatus.Uploading ): return ForwardStatus.CacheForward multi_media: List[pyrogram.raw.types.InputSingleMedia] = [] for it in node.media_group_ids[message.media_group_id]: if node.media_group_ids[message.media_group_id][it]: if multi_media: node.media_group_ids[message.media_group_id][it].message = "" multi_media.append(node.media_group_ids[message.media_group_id][it]) forward_status = ForwardStatus.SuccessForward reply_to_message_id = None message_thread_id = node.topic_id business_connection_id = None upload_telegram_chat_id = node.upload_telegram_chat_id if node.reply_to_message: if node.reply_to_message.chat.type != pyrogram.enums.ChatType.PRIVATE: reply_to_message_id = node.reply_to_message.id message_thread_id = node.reply_to_message.message_thread_id business_connection_id = node.reply_to_message.business_connection_id upload_telegram_chat_id = node.reply_to_message.chat.id if not await send_media_group_v2( client, upload_telegram_chat_id, # type: ignore multi_media, message_thread_id=message_thread_id, reply_to_message_id=reply_to_message_id, business_connection_id=business_connection_id, ): forward_status = ForwardStatus.FailedForward node.stat_forward(forward_status, len(multi_media)) node.media_group_ids.pop(message.media_group_id) return ForwardStatus.CacheForward def record_download_status(func): """Record download status""" @wraps(func) async def inner( client: pyrogram.client.Client, message: pyrogram.types.Message, media_types: List[str], file_formats: dict, node: TaskNode, ): if _download_cache[(node.chat_id, message.id)] is DownloadStatus.Downloading: return DownloadStatus.Downloading, None _download_cache[(node.chat_id, message.id)] = DownloadStatus.Downloading status, file_name = await func(client, message, media_types, file_formats, node) _download_cache[(node.chat_id, message.id)] = status return status, file_name return inner async def report_bot_download_status( client: pyrogram.Client, node: TaskNode, download_status: DownloadStatus, download_size: int = 0, ): """ Sends a message with the current status of the download bot. Parameters: client (pyrogram.Client): The client instance. node (TaskNode): The download task node. download_status (DownloadStatus): The current download status. Returns: None """ node.stat(download_status) node.total_download_byte += download_size await report_bot_status(client, node) async def report_bot_forward_status( client: pyrogram.Client, node: TaskNode, status: ForwardStatus, ): """ Sends a message with the current status of the download bot. Parameters: client (pyrogram.Client): The client instance. node (TaskNode): The download task node. status (ForwardStatus): The current forward status. Returns: None """ node.stat_forward(status) await report_bot_status(client, node) async def report_bot_status( client: pyrogram.Client, node: TaskNode, immediate_reply=False, ): """see _report_bot_status""" try: return await _report_bot_status(client, node, immediate_reply) except Exception as e: logger.debug(f"{e}") async def _report_bot_status( client: pyrogram.Client, node: TaskNode, immediate_reply=False, ): """ Sends a message with the current status of the download bot. Parameters: client (pyrogram.Client): The client instance. node (TaskNode): The download task node. immediate_reply(bool): Immediate reply Returns: None """ if not node.reply_message_id or not node.bot: return if immediate_reply or node.can_reply(): if node.upload_telegram_chat_id: node.forward_msg_detail_str = ( f"\nšŸ”„ {_t('Forward')}\n" f"ā”œā”€ šŸ“ {_t('Total')}: {node.total_forward_task}\n" f"ā”œā”€ āœ… {_t('Success')}: {node.success_forward_task}\n" f"ā”œā”€ āŒ {_t('Failed')}: {node.failed_forward_task}\n" f"ā””ā”€ ā© {_t('Skipped')}: {node.skip_forward_task}\n" ) upload_msg_detail_str: str = "" if node.upload_success_count: upload_msg_detail_str = ( f"\nā˜ļø {_t('Upload')}\n" f"ā””ā”€ āœ… {_t('Success')}: {node.upload_success_count}\n" ) for idx, value in node.cloud_drive_upload_stat_dict.items(): if value.transferred == value.total: continue temp_file_name = truncate_filename(os.path.basename(value.file_name), 10) upload_msg_detail_str += ( f" ā”œā”€ šŸ†” {_t('Message ID')}: {idx}\n" f" ā”‚ ā”œā”€ šŸ“ : {temp_file_name}\n" f" ā”‚ ā”œā”€ šŸ“ : {value.total}\n" f" ā”‚ ā”œā”€ ā« : {value.speed}\n" f" ā”‚ ā””ā”€ šŸ“Š : [" f'{create_progress_bar(int(value.percentage.split("%")[0]))}]' f" ({value.percentage})%\n" ) download_result_str = "" download_result = get_download_result() if node.chat_id in download_result: messages = download_result[node.chat_id] for idx, value in messages.items(): task_id = value["task_id"] if task_id != node.task_id or value["down_byte"] == value["total_size"]: continue temp_file_name = truncate_filename( os.path.basename(value["file_name"]), 10 ) progress = int(value["down_byte"] / value["total_size"] * 100) download_result_str += ( f" ā”œā”€ šŸ†” {_t('Message ID')}: {idx}\n" f" ā”‚ ā”œā”€ šŸ“ : {temp_file_name}\n" f" ā”‚ ā”œā”€ šŸ“ : {format_byte(value['total_size'])}\n" f" ā”‚ ā”œā”€ ā¬ : {format_byte(value['download_speed'])}/s\n" f" ā”‚ ā””ā”€ šŸ“Š : [{create_progress_bar(progress)}]" f" ({progress}%)\n" ) if download_result_str: download_result_str = ( f"\nšŸ“„ {_t('Download Progresses')}:\n" + download_result_str ) upload_result_str = "" for idx, value in node.upload_stat_dict.items(): if value.total_size == value.upload_size: continue temp_file_name = truncate_filename(os.path.basename(value.file_name), 10) progress = int(value.upload_size / value.total_size * 100) upload_result_str += ( f" ā”œā”€ šŸ†” {_t('Message ID')}: {idx}\n" f" ā”‚ ā”œā”€ šŸ“ : {temp_file_name}\n" f" ā”‚ ā”œā”€ šŸ“ : {format_byte(value.total_size)}\n" f" ā”‚ ā”œā”€ ā« : {format_byte(value.upload_speed)}/s\n" f" ā”‚ ā””ā”€ šŸ“Š : [{create_progress_bar(progress)}]" f" ({progress}%)\n" ) if upload_result_str: upload_result_str = f"\nšŸ“¤ {_t('Upload Progresses')}:\n" + upload_result_str new_msg_str = ( f"`\n" f"šŸ†” task id: {node.task_id}\n" f"šŸ“„ {_t('Downloading')}: {format_byte(node.total_download_byte)}\n" f"ā”œā”€ šŸ“ {_t('Total')}: {node.total_download_task}\n" f"ā”œā”€ āœ… {_t('Success')}: {node.success_download_task}\n" f"ā”œā”€ āŒ {_t('Failed')}: {node.failed_download_task}\n" f"ā””ā”€ ā© {_t('Skipped')}: {node.skip_download_task}\n" f"{node.forward_msg_detail_str}" f"{upload_msg_detail_str}" f"{upload_result_str}" f"{download_result_str}\n`" ) if new_msg_str != node.last_edit_msg: node.last_edit_msg = new_msg_str await client.edit_message_text( node.from_user_id, node.reply_message_id, new_msg_str, parse_mode=pyrogram.enums.ParseMode.MARKDOWN, ) def set_max_concurrent_transmissions( client: pyrogram.Client, max_concurrent_transmissions: int ): """Set maximum concurrent transmissions""" if getattr(client, "max_concurrent_transmissions", None): client.max_concurrent_transmissions = max_concurrent_transmissions client.save_file_semaphore = asyncio.Semaphore( client.max_concurrent_transmissions ) client.get_file_semaphore = asyncio.Semaphore( client.max_concurrent_transmissions ) async def fetch_message(client: pyrogram.Client, message: pyrogram.types.Message): """ This function retrieves a message from a specified chat using the Pyrogram library. Args: client (pyrogram.Client): A client instance created using Pyrogram. message (pyrogram.types.Message): A message instance returned from Pyrogram. Returns: pyrogram.types.Message: A message object retrieved from the specified chat. """ return await client.get_messages( chat_id=message.chat.id, message_ids=message.id, ) async def retry(func: Callable, args: tuple = (), max_attempts=3, wait_second=15): """ Asynchronously retries the provided function a specified number of times with a specified wait time between retries. :param func: The function to be retried. :param args: The arguments to be passed to the function. :param max_attempts: The maximum number of attempts to retry the function. Defaults to 3. :param wait_second: The wait time in seconds between each retry attempt. Defaults to 15. :return: The result of the function if it succeeds within the maximum number of attempts, otherwise None. """ for _ in range(1, max_attempts + 1): try: return await func(*args) except pyrogram.errors.exceptions.flood_420.FloodWait as wait_err: logger.warning("bad call retry: FlowWait {}", wait_err.value) await asyncio.sleep(wait_err.value) except Exception as e: logger.exception("Error: {}", e) await asyncio.sleep(wait_second) logger.error("Failed after {} attempts", max_attempts) return None async def get_media_group_with_retry( client: pyrogram.Client, chat_id: Union[int, str], message_id: int, max_attempts: int = 3, wait_second: int = 15, ): """ get_media_group_with_retry """ for attempt in range(1, max_attempts + 1): try: return await client.get_media_group(chat_id, message_id) except Exception as e: if attempt == max_attempts: logger.error("Failed Get Media Group[{}]", message_id) return types.List() logger.exception("Get Message[{}]: Error {}", message_id, e) await asyncio.sleep(wait_second) return types.List() async def check_user_permission( client: pyrogram.Client, user_id: Union[int, str], chat_id: Union[int, str] ) -> bool: """ Check if the user has permission to send videos in the group. Args: client (pyrogram.Client): A client instance created using Pyrogram. user_id (Union[int, str]): User Id chat_id (Union[int, str]): Chat Id Returns: if can_send_media_messages return True """ try: member = await client.get_chat_member(chat_id, user_id) return member and ( not member.permissions or member.permissions.can_send_media_messages ) except Exception: # logger.exception(e) pass return False def set_meta_data( meta_data: MetaData, message: pyrogram.types.Message, caption: str = None ): """Get all meta data""" # message meta_data.message_date = getattr(message, "date", None) if caption: meta_data.message_caption = caption else: meta_data.message_caption = getattr(message, "caption", None) or "" meta_data.message_id = getattr(message, "id", None) from_user = getattr(message, "from_user") meta_data.sender_id = from_user.id if from_user else 0 meta_data.sender_name = (from_user.username if from_user else "") or "" meta_data.reply_to_message_id = getattr( message, "reply_to_message_id", 1 ) # 1 for General meta_data.message_thread_id = getattr(message, "message_thread_id", 1) # media for kind in meta_data.AVAILABLE_MEDIA: media_obj = getattr(message, kind, None) if media_obj is not None: meta_data.media_type = kind break else: return meta_data.media_file_name = getattr(media_obj, "file_name", None) or "" meta_data.media_file_size = getattr(media_obj, "file_size", None) meta_data.media_width = getattr(media_obj, "width", None) meta_data.media_height = getattr(media_obj, "height", None) meta_data.media_duration = getattr(media_obj, "duration", None) meta_data.file_extension = get_extension( media_obj.file_id, getattr(media_obj, "mime_type", ""), False ) async def parse_link(client: pyrogram.Client, link_str: str): """Parse link""" link = extract_info_from_link(link_str) if link.comment_id: chat = await client.get_chat(link.group_id) if chat: return chat.linked_chat.id, link.comment_id, link.topic_id return link.group_id, link.post_id, link.topic_id async def update_cloud_upload_stat( transferred: str, total: str, percentage: str, speed: str, eta: str, node: TaskNode, message_id: int, file_name: str, ): """ Update the cloud upload statistics with the given information. Args: transferred (str): The amount of data transferred. total (str): The total size of the file. percentage (str): The percentage of the file uploaded. speed (str): The upload speed. eta (str): The estimated time of arrival for the upload to complete. node (TaskNode): The task node associated with the upload. message_id (int): The ID of the message. file_name (str): The name of the file being uploaded. Returns: None """ node.cloud_drive_upload_stat_dict[message_id] = CloudDriveUploadStat( file_name=file_name, transferred=transferred, total=total, percentage=percentage, speed=speed, eta=eta, ) async def update_upload_stat( upload_size: int, total_size: int, message_id: int, file_name: str, start_time: float, node: TaskNode, client: pyrogram.Client, ): """update_upload_status""" cur_time = time.time() if node.is_stop_transmission: client.stop_transmission() # TODO(tyh): web control upload stop if node.upload_stat_dict.get(message_id): upload_stat = node.upload_stat_dict[message_id] if cur_time - upload_stat.last_stat_time >= 1.0: upload_stat.upload_speed = max( int( (upload_size - upload_stat.upload_size) / (cur_time - upload_stat.last_stat_time) ), 0, ) upload_stat.last_stat_time = cur_time upload_stat.upload_size = upload_size node.upload_stat_dict[message_id] = upload_stat else: upload_stat = UploadProgressStat( file_name=file_name, total_size=total_size, upload_size=upload_size, start_time=start_time, last_stat_time=cur_time, upload_speed=upload_size / (cur_time - start_time), ) node.upload_stat_dict[message_id] = upload_stat # pylint: enable=W0201 class HookSession(pyrogram.session.Session): """Hook Session""" def start_timeout(self: pyrogram.session.Session, start_timeout: int): """ Set the start timeout for the session. Args: start_timeout (int): The start timeout value in seconds. Returns: None """ self.START_TIMEOUT = start_timeout # pylint: disable=all class HookClient(pyrogram.Client): """Hook Client""" # pylint: disable=R0901 START_TIME_OUT = 60 def __init__(self, name: str, **kwargs): if "start_timeout" in kwargs: value = kwargs.get("start_timeout") if value: self.START_TIME_OUT = value kwargs.pop("start_timeout") super().__init__(name, **kwargs) async def connect( self, ) -> bool: """ Connects the client to the server. Returns: bool: True if the client successfully connects to the server, False otherwise. Raises: ConnectionError: If the client is already connected. """ if self.is_connected: # type: ignore raise ConnectionError("Client is already connected") await self.load_session() self.session = HookSession( self, await self.storage.dc_id(), await self.storage.auth_key(), await self.storage.test_mode(), ) self.session.start_timeout(self.START_TIME_OUT) await self.session.start() self.is_connected = True return bool(await self.storage.user_id()) async def start(self): """ Starts the client by performing necessary initialization steps. Returns: The initialized client instance. """ is_authorized = await self.connect() try: if not is_authorized: await self.authorize() if not await self.storage.is_bot() and self.takeout: self.takeout_id = ( await self.invoke( pyrogram.raw.functions.account.InitTakeoutSession() ) ).id logger.warning(f"Takeout session {self.takeout_id} initiated") await self.invoke(pyrogram.raw.functions.updates.GetState()) except (Exception, KeyboardInterrupt): await self.disconnect() raise else: self.me = await self.get_me() await self.initialize() return self # pylint: disable=R0914,R0913 async def forward_messages( client: pyrogram.Client, chat_id: Union[int, str, None], from_chat_id: Union[int, str], message_ids: Union[int, Iterable[int]], disable_notification: bool = None, schedule_date: datetime = None, protect_content: bool = None, drop_author: bool = None, topic_id: int = None, caption: str = None, caption_entities: List[pyrogram.types.MessageEntity] = None, ) -> Union["types.Message", List["types.Message"]]: """Forward messages of any kind.""" is_iterable = not isinstance(message_ids, int) message_ids = list(message_ids) if is_iterable else [message_ids] # type: ignore r = await client.invoke( pyrogram.raw.functions.messages.ForwardMessages( to_peer=await client.resolve_peer(chat_id), from_peer=await client.resolve_peer(from_chat_id), id=message_ids, silent=disable_notification or None, random_id=[client.rnd_id() for _ in message_ids], schedule_date=pyrogram.utils.datetime_to_timestamp(schedule_date), noforwards=protect_content, drop_author=drop_author, top_msg_id=topic_id, entities=caption_entities, ) ) forwarded_messages = [] users = {i.id: i for i in r.users} chats = {i.id: i for i in r.chats} for i in r.updates: if isinstance( i, ( pyrogram.raw.types.UpdateNewMessage, pyrogram.raw.types.UpdateNewChannelMessage, pyrogram.raw.types.UpdateNewScheduledMessage, ), ): forwarded_messages.append( # pylint: disable=W0212 await types.Message._parse(client, i.message, users, chats) ) if caption and not is_iterable and forwarded_messages: await client.edit_message_caption( chat_id, forwarded_messages[0].id, caption=caption ) return types.List(forwarded_messages) if is_iterable else forwarded_messages[0]