From 514cf5e13dd614f559eff6716f17edca440da03b Mon Sep 17 00:00:00 2001 From: rking32 Date: Fri, 3 Jul 2020 02:14:03 +0530 Subject: [PATCH] fixes logs, restart, reload, upload and some fixes --- userge/core/client.py | 71 ++++++----- userge/core/methods/messages/send_as_file.py | 14 +- userge/core/types/bound/message.py | 27 ++-- userge/core/types/new/channel_logger.py | 127 +++++++++++-------- userge/plugins/misc/upload.py | 82 ++++++------ userge/plugins/misc/utube.py | 2 +- 6 files changed, 178 insertions(+), 145 deletions(-) diff --git a/userge/core/client.py b/userge/core/client.py index 24925ca89..a1708f353 100644 --- a/userge/core/client.py +++ b/userge/core/client.py @@ -35,35 +35,14 @@ _START_TIME = time.time() -class _UsergeBot(Methods, RawClient): - """ UsergeBot, the bot """ - def __init__(self, **kwargs) -> None: - _LOG.info(_LOG_STR, "Setting UsergeBot Configs") - super().__init__(session_name=":memory:", **kwargs) +async def _complete_init_tasks() -> None: + if not _INIT_TASKS: + return + await asyncio.gather(*_INIT_TASKS) + _INIT_TASKS.clear() -class Userge(Methods, RawClient): - """ Userge, the userbot """ - def __init__(self, **kwargs) -> None: - _LOG.info(_LOG_STR, "Setting Userge Configs") - if not (Config.HU_STRING_SESSION or Config.BOT_TOKEN): - print("Need HU_STRING_SESSION or BOT_TOKEN, Exiting...") - sys.exit() - kwargs = { - 'api_id': Config.API_ID, - 'api_hash': Config.API_HASH, - 'workers': Config.WORKERS - } - if Config.BOT_TOKEN: - if not Config.OWNER_ID: - print("Need OWNER_ID, Exiting...") - sys.exit() - kwargs['bot_token'] = Config.BOT_TOKEN - if Config.HU_STRING_SESSION and Config.BOT_TOKEN: - kwargs['bot'] = _UsergeBot(**kwargs) - kwargs['session_name'] = Config.HU_STRING_SESSION or ":memory:" - super().__init__(**kwargs) - +class _AbstractUserge(Methods, RawClient): @property def uptime(self) -> str: """ returns userge uptime """ @@ -76,15 +55,9 @@ def bot(self) -> '_UsergeBot': raise UsergeBotNotFound("Need BOT_TOKEN ENV!") return self._bot - async def _complete_init_tasks(self) -> None: - if not _INIT_TASKS: - return - await asyncio.gather(*_INIT_TASKS) - _INIT_TASKS.clear() - async def finalize_load(self) -> None: """ finalize the plugins load """ - await asyncio.gather(self._complete_init_tasks(), self.manager.init()) + await asyncio.gather(_complete_init_tasks(), self.manager.init()) async def load_plugin(self, name: str, reload_plugin: bool = False) -> None: """ Load plugin to Userge """ @@ -148,6 +121,36 @@ async def restart(self, update_req: bool = False) -> None: # pylint: disable=ar os.execl(sys.executable, sys.executable, '-m', 'userge') sys.exit() + +class _UsergeBot(_AbstractUserge): + """ UsergeBot, the bot """ + def __init__(self, **kwargs) -> None: + _LOG.info(_LOG_STR, "Setting UsergeBot Configs") + super().__init__(session_name=":memory:", **kwargs) + + +class Userge(_AbstractUserge): + """ Userge, the userbot """ + def __init__(self, **kwargs) -> None: + _LOG.info(_LOG_STR, "Setting Userge Configs") + if not (Config.HU_STRING_SESSION or Config.BOT_TOKEN): + print("Need HU_STRING_SESSION or BOT_TOKEN, Exiting...") + sys.exit() + kwargs = { + 'api_id': Config.API_ID, + 'api_hash': Config.API_HASH, + 'workers': Config.WORKERS + } + if Config.BOT_TOKEN: + if not Config.OWNER_ID: + print("Need OWNER_ID, Exiting...") + sys.exit() + kwargs['bot_token'] = Config.BOT_TOKEN + if Config.HU_STRING_SESSION and Config.BOT_TOKEN: + kwargs['bot'] = _UsergeBot(**kwargs) + kwargs['session_name'] = Config.HU_STRING_SESSION or ":memory:" + super().__init__(**kwargs) + async def start(self) -> None: """ start client and bot """ _LOG.info(_LOG_STR, "Starting Userge") diff --git a/userge/core/methods/messages/send_as_file.py b/userge/core/methods/messages/send_as_file.py index 66e3cf519..c89adae0c 100644 --- a/userge/core/methods/messages/send_as_file.py +++ b/userge/core/methods/messages/send_as_file.py @@ -11,8 +11,7 @@ __all__ = ['SendAsFile'] import os -import asyncio -from typing import Union +from typing import Union, Optional import aiofiles @@ -31,7 +30,7 @@ async def send_as_file(self, filename: str = "output.txt", caption: str = '', log: Union[bool, str] = False, - delete_message: bool = True) -> 'types.bound.Message': + reply_to_message_id: Optional[int] = None) -> 'types.bound.Message': """\nYou can send large outputs as file Example: @@ -63,24 +62,23 @@ async def send_as_file(self, If ``True``, the message will be deleted after sending the file. + reply_to_message_id (``int``, *optional*): + If the message is a reply, ID of the original message. + Returns: On success, the sent Message is returned. """ async with aiofiles.open(filename, "w+", encoding="utf8") as out_file: await out_file.write(text) - reply_to_id = self.reply_to_message.message_id if self.reply_to_message \ - else self.message_id _LOG.debug(_LOG_STR, f"Uploading {filename} To Telegram") msg = await self.send_document(chat_id=chat_id, document=filename, caption=caption, disable_notification=True, - reply_to_message_id=reply_to_id) + reply_to_message_id=reply_to_message_id) os.remove(filename) if log: if isinstance(log, str): self._channel.update(log) await self._channel.fwd_msg(msg) - if delete_message: - asyncio.get_event_loop().create_task(self.delete()) return types.bound.Message(self, msg) diff --git a/userge/core/types/bound/message.py b/userge/core/types/bound/message.py index 54f4d145d..562f456e1 100644 --- a/userge/core/types/bound/message.py +++ b/userge/core/types/bound/message.py @@ -11,6 +11,7 @@ __all__ = ['Message'] import re +import asyncio from typing import List, Dict, Union, Optional, Sequence from pyrogram import InlineKeyboardMarkup, Message as RawMessage @@ -32,7 +33,7 @@ def _msg_to_dict(message: RawMessage) -> Dict[str, object]: kwargs_ = vars(message) del message for key_ in ['_client', '_channel', '_filtered', '_process_canceled', - '_filtered_input_str', '_flags', '_kwargs']: + 'client', '_filtered_input_str', '_flags', '_kwargs']: if key_ in kwargs_: del kwargs_[key_] return kwargs_ @@ -41,7 +42,7 @@ def _msg_to_dict(message: RawMessage) -> Dict[str, object]: class Message(RawMessage): """ Modded Message Class For Userge """ def __init__(self, - client: '_client.Userge', + client: Union['_client.Userge', '_client.UsergeBot'], message: RawMessage, **kwargs: Union[str, bool]) -> None: super().__init__(client=client, **_msg_to_dict(message)) @@ -56,6 +57,11 @@ def __init__(self, self._flags: Dict[str, str] = {} self._kwargs = kwargs + @property + def client(self) -> Union['_client.Userge', '_client.UsergeBot']: + """ returns client """ + return self._client + @property def input_str(self) -> str: """ Returns the input string without command """ @@ -148,12 +154,17 @@ async def send_as_file(self, Returns: On success, the sent Message is returned. """ - return self._client.send_as_file(chat_id=self.chat.id, - text=text, - filename=filename, - caption=caption, - log=log, - delete_message=delete_message) + reply_to_id = self.reply_to_message.message_id if self.reply_to_message \ + else self.message_id + if delete_message: + asyncio.get_event_loop().create_task(self.delete()) + return await self._client.send_as_file(chat_id=self.chat.id, + text=text, + filename=filename, + caption=caption, + log=log, + delete_message=delete_message, + reply_to_message_id=reply_to_id) async def reply(self, text: str, diff --git a/userge/core/types/new/channel_logger.py b/userge/core/types/new/channel_logger.py index 37a6e12d2..049a9a628 100644 --- a/userge/core/types/new/channel_logger.py +++ b/userge/core/types/new/channel_logger.py @@ -14,6 +14,7 @@ from typing import Optional, Tuple from pyrogram import Message as RawMessage +from pyrogram.errors.exceptions.bad_request_400 import ChannelInvalid from userge import logging, Config from userge.utils import SafeDict @@ -92,9 +93,12 @@ async def log(self, text: str) -> Optional[int]: """ _LOG.debug(_LOG_STR, f"logging text : {text} to channel : {Config.LOG_CHANNEL_ID}") if Config.LOG_CHANNEL_ID: - msg = await self._client.send_message(chat_id=Config.LOG_CHANNEL_ID, - text=self._string.format(text.strip())) - return msg.message_id + try: + msg = await self._client.send_message(chat_id=Config.LOG_CHANNEL_ID, + text=self._string.format(text.strip())) + return msg.message_id + except ChannelInvalid: + pass async def fwd_msg(self, message: '_message.Message', @@ -124,15 +128,19 @@ async def fwd_msg(self, _LOG.debug( _LOG_STR, f"forwarding msg : {message} to channel : {Config.LOG_CHANNEL_ID}") if Config.LOG_CHANNEL_ID and isinstance(message, RawMessage): - if message.media: - asyncio.get_event_loop().create_task(self.log("**Forwarding Message...**")) - await self._client.forward_messages(chat_id=Config.LOG_CHANNEL_ID, - from_chat_id=message.chat.id, - message_ids=message.message_id, - as_copy=as_copy, - remove_caption=remove_caption) - else: - await self.log(message.text.html) + try: + if message.media: + asyncio.get_event_loop().create_task(self.log("**Forwarding Message...**")) + await self._client.forward_messages(chat_id=Config.LOG_CHANNEL_ID, + from_chat_id=message.chat.id, + message_ids=message.message_id, + as_copy=as_copy, + remove_caption=remove_caption) + else: + await self.log( + message.text.html if hasattr(message.text, 'html') else message.text) + except ChannelInvalid: + pass async def store(self, message: Optional['_message.Message'], @@ -149,21 +157,25 @@ async def store(self, Returns: message_id on success or None """ - caption = caption or '' - if message and message.caption: - caption = caption + message.caption.html - if message and message.media: - if caption: - caption = self._string.format(caption.strip()) - file_id, file_ref = _get_file_id_and_ref(message) - msg = await self._client.send_cached_media(chat_id=Config.LOG_CHANNEL_ID, - file_id=file_id, - file_ref=file_ref, - caption=caption) - message_id = msg.message_id - else: - message_id = await self.log(caption) - return message_id + if Config.LOG_CHANNEL_ID: + caption = caption or '' + if message and message.caption: + caption = caption + message.caption.html + if message and message.media: + if caption: + caption = self._string.format(caption.strip()) + file_id, file_ref = _get_file_id_and_ref(message) + try: + msg = await self._client.send_cached_media(chat_id=Config.LOG_CHANNEL_ID, + file_id=file_id, + file_ref=file_ref, + caption=caption) + message_id = msg.message_id + except ChannelInvalid: + message_id = None + else: + message_id = await self.log(caption) + return message_id async def forward_stored(self, message_id: int, @@ -192,30 +204,37 @@ async def forward_stored(self, Returns: None """ - message = await self._client.get_messages(chat_id=Config.LOG_CHANNEL_ID, - message_ids=message_id) - caption = '' - if message.caption: - caption = message.caption.html.split('\n\n', maxsplit=1)[-1] - elif message.text: - caption = message.text.html.split('\n\n', maxsplit=1)[-1] - if caption: - u_dict = await self._client.get_user_dict(user_id) - chat = await self._client.get_chat(chat_id) - u_dict.update( - {'chat': chat.title if chat.title else "this group", 'count': chat.members_count}) - caption = caption.format_map(SafeDict(**u_dict)) - if message.media: - file_id, file_ref = _get_file_id_and_ref(message) - msg = await self._client.send_cached_media(chat_id=chat_id, - file_id=file_id, - file_ref=file_ref, - caption=caption, - reply_to_message_id=reply_to_message_id) - else: - msg = await self._client.send_message(chat_id=chat_id, - text=caption, - reply_to_message_id=reply_to_message_id) - if del_in and msg: - await asyncio.sleep(del_in) - await msg.delete() + if Config.LOG_CHANNEL_ID: + try: + message = await self._client.get_messages(chat_id=Config.LOG_CHANNEL_ID, + message_ids=message_id) + caption = '' + if message.caption: + caption = message.caption.html.split('\n\n', maxsplit=1)[-1] + elif message.text: + caption = message.text.html.split('\n\n', maxsplit=1)[-1] + if caption: + u_dict = await self._client.get_user_dict(user_id) + chat = await self._client.get_chat(chat_id) + u_dict.update( + {'chat': chat.title if chat.title else "this group", + 'count': chat.members_count}) + caption = caption.format_map(SafeDict(**u_dict)) + if message.media: + file_id, file_ref = _get_file_id_and_ref(message) + msg = await self._client.send_cached_media( + chat_id=chat_id, + file_id=file_id, + file_ref=file_ref, + caption=caption, + reply_to_message_id=reply_to_message_id) + else: + msg = await self._client.send_message( + chat_id=chat_id, + text=caption, + reply_to_message_id=reply_to_message_id) + if del_in and msg: + await asyncio.sleep(del_in) + await msg.delete() + except ChannelInvalid: + pass diff --git a/userge/plugins/misc/upload.py b/userge/plugins/misc/upload.py index 2ecc06827..db998bbad 100644 --- a/userge/plugins/misc/upload.py +++ b/userge/plugins/misc/upload.py @@ -62,7 +62,7 @@ async def rename_(message: Message): dl_loc = os.path.join(Config.DOWN_PATH, os.path.basename(dl_loc)) new_loc = os.path.join(Config.DOWN_PATH, message.filtered_input_str) os.rename(dl_loc, new_loc) - await upload(Path(new_loc), message.chat.id, message.flags, True) + await upload(message, Path(new_loc), True) else: await message.edit("Please read `.help rename`", del_in=5) @@ -90,8 +90,8 @@ async def convert_(message: Message): else: await message.delete() dl_loc = os.path.join(Config.DOWN_PATH, os.path.basename(dl_loc)) - flags = {} if message.reply_to_message.document else {'d': ''} - await upload(Path(dl_loc), message.chat.id, flags, True) + message.flags = {} if message.reply_to_message.document else {'d': ''} + await upload(message, Path(dl_loc), True) else: await message.edit("Please read `.help convert`", del_in=5) @@ -105,7 +105,6 @@ async def convert_(message: Message): "{tr}upload downloads/test.mp4"]}, del_pre=True) async def uploadtotg(message: Message): """ upload to telegram """ - flags = message.flags path_ = message.filtered_input_str if not path_: await message.edit("invalid input!, check `.help .upload`", del_in=5) @@ -179,40 +178,40 @@ async def uploadtotg(message: Message): await message.edit("wrong syntax\n`.upload [path]`") else: await message.delete() - await explorer(string, message.chat.id, flags, del_path) + await explorer(message, string, del_path) -async def explorer(path: Path, chatid, flags, del_path): +async def explorer(message: Message, path: Path, del_path): if path.is_file(): try: if path.stat().st_size: - await upload(path, chatid, flags, del_path) + await upload(message, path, del_path) except FloodWait as x: time.sleep(x.x) # asyncio sleep ? elif path.is_dir(): for i in sorted(path.iterdir()): - await explorer(i, chatid, flags, del_path) + await explorer(message, i, del_path) -async def upload(path: Path, chat_id: int, flags: dict, del_path: bool = False): - if path.name.endswith((".mkv", ".mp4", ".webm")) and ('d' not in flags): - await vid_upload(chat_id, path, del_path) - elif path.name.endswith((".mp3", ".flac", ".wav", ".m4a")) and ('d' not in flags): - await audio_upload(chat_id, path, del_path) +async def upload(message: Message, path: Path, del_path: bool = False): + if path.name.endswith((".mkv", ".mp4", ".webm")) and ('d' not in message.flags): + await vid_upload(message, path, del_path) + elif path.name.endswith((".mp3", ".flac", ".wav", ".m4a")) and ('d' not in message.flags): + await audio_upload(message, path, del_path) else: - await doc_upload(chat_id, path, del_path) + await doc_upload(message, path, del_path) -async def doc_upload(chat_id, path, del_path: bool): - message: Message = await userge.send_message( - chat_id, f"`Uploading {path.name} ...`") +async def doc_upload(message: Message, path, del_path: bool): + sent: Message = await message.client.send_message( + message.chat.id, f"`Uploading {path.name} ...`") start_t = datetime.now() c_time = time.time() thumb = await get_thumb() - await userge.send_chat_action(chat_id, "upload_document") + await message.client.send_chat_action(message.chat.id, "upload_document") try: - msg = await userge.send_document( - chat_id=chat_id, + msg = await message.client.send_document( + chat_id=message.chat.id, document=str(path), thumb=thumb, caption=path.name, @@ -224,27 +223,28 @@ async def doc_upload(chat_id, path, del_path: bool): ) ) except Exception as u_e: - await message.edit(u_e) + await sent.edit(u_e) raise u_e else: - await finalize(chat_id, message, msg, start_t) + await sent.delete() + await finalize(message, msg, start_t) finally: if os.path.exists(str(path)) and del_path: os.remove(str(path)) -async def vid_upload(chat_id, path, del_path: bool): +async def vid_upload(message: Message, path, del_path: bool): strpath = str(path) thumb = await get_thumb(strpath) metadata = extractMetadata(createParser(strpath)) - message: Message = await userge.send_message( - chat_id, f"`Uploading {path.name} as a video ..`") + sent: Message = await message.client.send_message( + message.chat.id, f"`Uploading {path.name} as a video ..`") start_t = datetime.now() c_time = time.time() - await userge.send_chat_action(chat_id, "upload_video") + await message.client.send_chat_action(message.chat.id, "upload_video") try: - msg = await userge.send_video( - chat_id=chat_id, + msg = await message.client.send_video( + chat_id=message.chat.id, video=strpath, duration=metadata.get("duration").seconds, thumb=thumb, @@ -257,21 +257,22 @@ async def vid_upload(chat_id, path, del_path: bool): ) ) except Exception as u_e: - await message.edit(u_e) + await sent.edit(u_e) raise u_e else: + await sent.delete() await remove_thumb(thumb) - await finalize(chat_id, message, msg, start_t) + await finalize(message, msg, start_t) finally: if os.path.exists(str(path)) and del_path: os.remove(str(path)) -async def audio_upload(chat_id, path, del_path: bool): +async def audio_upload(message: Message, path, del_path: bool): title = None artist = None - message: Message = await userge.send_message( - chat_id, f"`Uploading {path.name} as audio ...`") + sent: Message = await message.client.send_message( + message.chat.id, f"`Uploading {path.name} as audio ...`") strpath = str(path) file_size = humanbytes(os.stat(strpath).st_size) start_t = datetime.now() @@ -282,13 +283,13 @@ async def audio_upload(chat_id, path, del_path: bool): title = metadata.get("title") if metadata.has("artist"): artist = metadata.get("artist") - await userge.send_chat_action(chat_id, "upload_audio") + await message.client.send_chat_action(message.chat.id, "upload_audio") try: audio_caption = "" audio_caption += f"{path.name} " audio_caption += f"[ {file_size} ]" - msg = await userge.send_audio( - chat_id=chat_id, + msg = await message.client.send_audio( + chat_id=message.chat.id, audio=strpath, thumb=thumb, caption=audio_caption, @@ -303,10 +304,11 @@ async def audio_upload(chat_id, path, del_path: bool): ) ) except Exception as u_e: - await message.edit(u_e) + await sent.edit(u_e) raise u_e else: - await finalize(chat_id, message, msg, start_t) + await sent.delete() + await finalize(message, msg, start_t) finally: if os.path.exists(str(path)) and del_path: os.remove(str(path)) @@ -331,9 +333,9 @@ async def remove_thumb(thumb: str) -> None: os.remove(thumb) -async def finalize(chat_id: int, message: Message, msg: Message, start_t): +async def finalize(message: Message, msg: Message, start_t): await CHANNEL.fwd_msg(msg) - await userge.send_chat_action(chat_id, "cancel") + await message.client.send_chat_action(message.chat.id, "cancel") if message.process_is_canceled: await message.edit("`Process Canceled!`", del_in=5) else: diff --git a/userge/plugins/misc/utube.py b/userge/plugins/misc/utube.py index 80636631e..86bbe1e28 100644 --- a/userge/plugins/misc/utube.py +++ b/userge/plugins/misc/utube.py @@ -116,7 +116,7 @@ def __progress(data: dict): _fpath = glob.glob(path.join(Config.DOWN_PATH, str(startTime), '*'))[0] await message.edit(f"**YTDL completed in {round(time() - startTime)} seconds**\n`{_fpath}`") if 't' in message.flags: - await upload(Path(_fpath), message.chat.id, message.flags) + await upload(message, Path(_fpath)) else: await message.edit(str(retcode))