From 3a96df8add128bea36446f56e83c4640b2036e09 Mon Sep 17 00:00:00 2001 From: profitroll Date: Mon, 26 Jun 2023 12:19:29 +0200 Subject: [PATCH] Added direct PyroClient support --- libbot/pyrogram/classes/__init__.py | 3 + libbot/pyrogram/classes/client.py | 275 ++++++++++++++++++++++++++ libbot/pyrogram/classes/command.py | 9 + libbot/pyrogram/classes/commandset.py | 35 ++++ pyproject.toml | 6 +- 5 files changed, 326 insertions(+), 2 deletions(-) create mode 100644 libbot/pyrogram/classes/__init__.py create mode 100644 libbot/pyrogram/classes/client.py create mode 100644 libbot/pyrogram/classes/command.py create mode 100644 libbot/pyrogram/classes/commandset.py diff --git a/libbot/pyrogram/classes/__init__.py b/libbot/pyrogram/classes/__init__.py new file mode 100644 index 0000000..5b08e07 --- /dev/null +++ b/libbot/pyrogram/classes/__init__.py @@ -0,0 +1,3 @@ +from .client import PyroClient +from .command import PyroCommand +from .commandset import CommandSet diff --git a/libbot/pyrogram/classes/client.py b/libbot/pyrogram/classes/client.py new file mode 100644 index 0000000..2b29842 --- /dev/null +++ b/libbot/pyrogram/classes/client.py @@ -0,0 +1,275 @@ +import logging +from datetime import datetime, timedelta +from os import getpid +from pathlib import Path +from time import time +from typing import List, Union + +try: + import pyrogram + from pyrogram.client import Client + from pyrogram.errors import BadRequest + from pyrogram.handlers.message_handler import MessageHandler + from pyrogram.raw.all import layer + from pyrogram.types import ( + BotCommand, + BotCommandScopeAllChatAdministrators, + BotCommandScopeAllGroupChats, + BotCommandScopeAllPrivateChats, + BotCommandScopeChat, + BotCommandScopeChatAdministrators, + BotCommandScopeChatMember, + BotCommandScopeDefault, + ) +except ImportError as exc: + raise ImportError( + "You need to install libbot[pyrogram] in order to use this class." + ) from exc + +try: + from ujson import dumps, loads +except ImportError: + from json import dumps, loads + +from libbot.i18n import BotLocale +from libbot.i18n.sync import _ +from libbot.pyrogram.classes.command import PyroCommand +from libbot.pyrogram.classes.commandset import CommandSet + +logger = logging.getLogger(__name__) + + +class PyroClient(Client): + def __init__(self): + with open("config.json", "r", encoding="utf-8") as f: + self.config: dict = loads(f.read()) + super().__init__( + name="bot_client", + api_id=self.config["bot"]["api_id"], + api_hash=self.config["bot"]["api_hash"], + bot_token=self.config["bot"]["bot_token"], + # Workers should be commented when using convopyro, otherwise + # handlers land in another event loop and you won't see them + workers=self.config["bot"]["workers"], + plugins=dict(root="plugins", exclude=self.config["disabled_plugins"]), + sleep_threshold=120, + max_concurrent_transmissions=self.config["bot"][ + "max_concurrent_transmissions" + ], + ) + self.owner: int = self.config["bot"]["owner"] + self.commands: List[PyroCommand] = [] + self.scoped_commands: bool = self.config["bot"]["scoped_commands"] + self.start_time: float = 0 + + self.bot_locale: BotLocale = BotLocale(Path(self.config["locations"]["locale"])) + self.default_locale: str = self.bot_locale.default + self.locales: dict = self.bot_locale.locales + + self._ = self.bot_locale._ + self.in_all_locales = self.bot_locale.in_all_locales + self.in_every_locale = self.bot_locale.in_every_locale + + async def start(self, scheduler): + await super().start() + + self.start_time = time() + + logger.info( + "Bot is running with Pyrogram v%s (Layer %s) and has started as @%s on PID %s.", + pyrogram.__version__, + layer, + self.me.username, + getpid(), + ) + + try: + await self.send_message( + chat_id=self.config["reports"]["chat_id"], + text=f"Bot started PID `{getpid()}`", + ) + + scheduler.add_job( + self.register_commands, + trigger="date", + run_date=datetime.now() + timedelta(seconds=5), + kwargs={"command_sets": await self.collect_commands()}, + ) + + scheduler.start() + except BadRequest: + logger.warning("Unable to send message to report chat.") + + async def stop(self): + try: + await self.send_message( + chat_id=self.config["reports"]["chat_id"], + text=f"Bot stopped with PID `{getpid()}`", + ) + except BadRequest: + logger.warning("Unable to send message to report chat.") + await super().stop() + logger.warning("Bot stopped with PID %s.", getpid()) + + async def collect_commands(self) -> Union[List[CommandSet], None]: + """Gather list of the bot's commands + + ### Returns: + * `List[CommandSet]`: List of the commands' sets + """ + command_sets = None + + # If config's bot.scoped_commands is true - more complicated + # scopes system will be used instead of simple global commands + if self.scoped_commands: + scopes = {} + command_sets = [] + + # Iterate through all commands in config + for command, contents in self.config["commands"].items(): + # Iterate through all scopes of a command + for scope in contents["scopes"]: + if dumps(scope) not in scopes: + scopes[dumps(scope)] = {"_": []} + + # Add command to the scope's flattened key in scopes dict + scopes[dumps(scope)]["_"].append( + BotCommand(command, _(command, "commands")) + ) + + for locale, string in ( + self.in_every_locale(command, "commands") + ).items(): + if locale not in scopes[dumps(scope)]: + scopes[dumps(scope)][locale] = [] + + scopes[dumps(scope)][locale].append(BotCommand(command, string)) + + # Iterate through all scopes and its commands + for scope, locales in scopes.items(): + # Make flat key a dict again + scope_dict = loads(scope) + + # Replace "owner" in the bot scope with owner's id + if "chat_id" in scope_dict and scope_dict["chat_id"] == "owner": + scope_dict["chat_id"] = self.owner + + # Create object with the same name and args from the dict + try: + scope_obj = globals()[scope_dict["name"]]( + **{ + key: value + for key, value in scope_dict.items() + if key != "name" + } + ) + except NameError: + logger.error( + "Could not register commands of the scope '%s' due to an invalid scope class provided!", + scope_dict["name"], + ) + continue + except TypeError: + logger.error( + "Could not register commands of the scope '%s' due to an invalid class arguments provided!", + scope_dict["name"], + ) + continue + + # Add set of commands to the list of the command sets + for locale, commands in locales.items(): + if locale == "_": + command_sets.append( + CommandSet(commands, scope=scope_obj, language_code="") + ) + continue + command_sets.append( + CommandSet(commands, scope=scope_obj, language_code=locale) + ) + + logger.info("Registering the following command sets: %s", command_sets) + + else: + # This part here looks into the handlers and looks for commands + # in it, if there are any. Then adds them to self.commands + for handler in self.dispatcher.groups[0]: + if isinstance(handler, MessageHandler): + for entry in [handler.filters.base, handler.filters.other]: + if hasattr(entry, "commands"): + for command in entry.commands: + logger.info("I see a command %s in my filters", command) + self.add_command(command) + + return command_sets + + def add_command( + self, + command: str, + ): + """Add command to the bot's internal commands list + + ### Args: + * command (`str`) + """ + self.commands.append( + PyroCommand( + command, + _(command, "commands"), + ) + ) + logger.info( + "Added command '%s' to the bot's internal commands list", + command, + ) + + async def register_commands( + self, command_sets: Union[List[CommandSet], None] = None + ): + """Register commands stored in bot's 'commands' attribute""" + + if command_sets is None: + commands = [ + BotCommand(command=command.command, description=command.description) + for command in self.commands + ] + + logger.info( + "Registering commands %s with a default scope 'BotCommandScopeDefault'" + ) + + await self.set_bot_commands(commands) + return + + for command_set in command_sets: + logger.info( + "Registering command set with commands %s and scope '%s' (%s)", + command_set.commands, + command_set.scope, + command_set.language_code, + ) + await self.set_bot_commands( + command_set.commands, + command_set.scope, + language_code=command_set.language_code, + ) + + async def remove_commands(self, command_sets: Union[List[CommandSet], None] = None): + """Remove commands stored in bot's 'commands' attribute""" + + if command_sets is None: + logger.info( + "Removing commands with a default scope 'BotCommandScopeDefault'" + ) + await self.delete_bot_commands(BotCommandScopeDefault()) + return + + for command_set in command_sets: + logger.info( + "Removing command set with scope '%s' (%s)", + command_set.scope, + command_set.language_code, + ) + await self.delete_bot_commands( + command_set.scope, + language_code=command_set.language_code, + ) diff --git a/libbot/pyrogram/classes/command.py b/libbot/pyrogram/classes/command.py new file mode 100644 index 0000000..7b01180 --- /dev/null +++ b/libbot/pyrogram/classes/command.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass + + +@dataclass +class PyroCommand: + """Command stored in PyroClient's 'commands' attribute""" + + command: str + description: str diff --git a/libbot/pyrogram/classes/commandset.py b/libbot/pyrogram/classes/commandset.py new file mode 100644 index 0000000..1d9d358 --- /dev/null +++ b/libbot/pyrogram/classes/commandset.py @@ -0,0 +1,35 @@ +from dataclasses import dataclass +from typing import List, Union + +try: + from pyrogram.types import ( + BotCommand, + BotCommandScopeAllChatAdministrators, + BotCommandScopeAllGroupChats, + BotCommandScopeAllPrivateChats, + BotCommandScopeChat, + BotCommandScopeChatAdministrators, + BotCommandScopeChatMember, + BotCommandScopeDefault, + ) +except ImportError as exc: + raise ImportError( + "You need to install libbot[pyrogram] in order to use this class." + ) from exc + + +@dataclass +class CommandSet: + """Command stored in PyroClient's 'commands' attribute""" + + commands: List[BotCommand] + scope: Union[ + BotCommandScopeDefault, + BotCommandScopeAllPrivateChats, + BotCommandScopeAllGroupChats, + BotCommandScopeAllChatAdministrators, + BotCommandScopeChat, + BotCommandScopeChatAdministrators, + BotCommandScopeChatMember, + ] = BotCommandScopeDefault + language_code: str = "" diff --git a/pyproject.toml b/pyproject.toml index f474b45..3893e58 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,10 +37,12 @@ Tracker = "https://git.end-play.xyz/profitroll/LibBotUniversal/issues" [tool.setuptools] packages = [ "libbot", - "libbot.i18n", "libbot.sync", - "libbot.i18n.classes", + "libbot.pyrogram", + "libbot.pyrogram.classes", + "libbot.i18n", "libbot.i18n.sync", + "libbot.i18n.classes", ] [tool.setuptools_scm]