import asyncio
import logging
import os
import shlex
import sys
from typing import Dict, List, Optional, Tuple, Union
import aiohttp
import arrow
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.interval import IntervalTrigger
from pyrogram import Client, enums, errors
from pyrogram.enums import ChatType
from pyrogram.session import Session
from pyrogram.storage import Storage
from pyrogram.types import Chat, Message, User
class CustomFormatter(logging.Formatter):
# Colors
black = "\x1b[30m"
red = "\x1b[31m"
green = "\x1b[32m"
yellow = "\x1b[33m"
blue = "\x1b[34m"
gray = "\x1b[38m"
# Styles
reset = "\x1b[0m"
bold = "\x1b[1m"
COLORS = {
logging.DEBUG: gray + bold,
logging.INFO: blue + bold,
logging.WARNING: yellow + bold,
logging.ERROR: red,
logging.CRITICAL: red + bold,
}
def format(self, record):
log_color = self.COLORS[record.levelno]
fmt = "(black){asctime}(reset) (levelcolor){levelname:<8}(reset) (green){name}(reset) {message}"
fmt = fmt.replace("(black)", self.black + self.bold)
fmt = fmt.replace("(reset)", self.reset)
fmt = fmt.replace("(levelcolor)", log_color)
fmt = fmt.replace("(green)", self.green + self.bold)
formatter = logging.Formatter(fmt, "%Y-%m-%d %H:%M:%S", style="{")
return formatter.format(record)
def get_full_name(obj: Union[User, Chat]) -> str:
if isinstance(obj, Chat):
if obj.type == ChatType.PRIVATE:
return f"{obj.first_name} {obj.last_name}" if obj.last_name else obj.first_name
return obj.title
elif isinstance(obj, User):
return f"{obj.first_name} {obj.last_name}" if obj.last_name else obj.first_name
else:
raise TypeError("obj must be User or Chat")
def format_exc(e: Exception, suffix="") -> str:
if isinstance(e, errors.RPCError):
return (
f"Telegram API error!\n"
f"[{e.CODE} {e.ID or e.NAME}] — {e.MESSAGE.format(value=e.value)}
\n\n"
f"{suffix}"
)
return f"{e.__class__.__name__}: {e}
\n\n{suffix}"
def with_reply(func):
async def wrapped(client: Client, message: Message):
if not message.reply_to_message:
await message.edit("Reply to message is required")
else:
return await func(client, message)
return wrapped
def with_args(text: str):
def decorator(func):
async def wrapped(client: Client, message: Message):
if message.text and len(message.text.split()) == 1:
await message.edit(text)
else:
return await func(client, message)
return wrapped
return decorator
def with_premium(func):
async def wrapped(client: Client, message: Message):
if not (await client.get_me()).is_premium:
await message.edit("Premium account is required")
else:
return await func(client, message)
return wrapped
async def dpaste(code: str):
async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False)) as session:
data = {"content": code, "lexer": "python", "expires": "never"}
async with session.post("https://dpaste.org/api/", data=data) as resp:
if resp.status != 200:
return "Pasting failed!"
else:
return (await resp.text()).replace('"', "")
async def paste_neko(code: str):
try:
async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False)) as session:
async with session.post(
"https://nekobin.com/api/documents",
json={"content": code},
) as paste:
paste.raise_for_status()
result = await paste.json()
except Exception:
return await dpaste(code=code)
else:
return f"nekobin.com/{result['result']['key']}.py"
def get_args_raw(message: Union[Message, str], use_reply: bool = None) -> str:
"""Returns text after command.
Args:
message (Union[Message, str]): Message or text.
use_reply (bool, optional): Try to get args from reply message if no args in message. Defaults to None.
Returns:
str: Text after command or empty string.
"""
if isinstance(message, Message):
text = message.text or message.caption
args = text.split(maxsplit=1)[1] if len(text.split()) > 1 else ""
if use_reply and not args:
args = message.reply_to_message.text or message.reply_to_message.caption
elif not isinstance(message, str):
return ""
return args or ""
def get_args(
message: Union[Message, str], use_reply: bool = None
) -> Tuple[List[str], Dict[str, str]]:
"""Returns list of common args and a dictionary with named args.
Args:
message (Union[Message, str]): Message or text.
use_reply (bool, optional): Try to get args from reply message if no args in message. Defaults to None.
Returns:
List[str]: List of args.
"""
raw_args = get_args_raw(message, use_reply)
try:
args = list(filter(lambda x: len(x) > 0, shlex.split(raw_args)))
except ValueError:
return [raw_args], {}
common_args = []
named_args = {}
i = 0
while i < len(args):
arg = args[i]
if arg.startswith("-"):
if i + 1 < len(args) and (
not args[i + 1].startswith("-") or len(args[i + 1].split()) > 1
):
named_args[arg] = args[i + 1]
i += 2
else:
i += 1
else:
i += 1
common_args.append(arg)
return common_args, named_args
class ScheduleJob:
def __init__(
self,
func: callable,
trigger: Optional[Union[CronTrigger, IntervalTrigger]] = IntervalTrigger(seconds=3600),
*args,
**kwargs,
):
self.func = func
self.args = args or []
self.kwargs = kwargs or {}
self.id = func.__name__
self.trigger = trigger
def get_ram_usage() -> float:
"""Returns current process tree memory usage in MB"""
try:
import psutil
current_process = psutil.Process(os.getpid())
mem = current_process.memory_info()[0] / 2.0**20
for child in current_process.children(recursive=True):
mem += child.memory_info()[0] / 2.0**20
return round(mem, 1)
except Exception:
return 0
def get_cpu_usage() -> float:
"""Returns current process tree CPU usage in %"""
try:
import psutil
current_process = psutil.Process(os.getpid())
cpu = current_process.cpu_percent()
for child in current_process.children(recursive=True):
cpu += child.cpu_percent()
return round(cpu, 1)
except Exception:
return 0
def humanize_seconds(seconds: Union[int, float]) -> str:
"""Returns humanized time delta from seconds"""
current_time = arrow.get()
target_time = current_time.shift(seconds=-seconds)
return target_time.humanize(current_time, only_distance=True)
class Command:
def __init__(
self,
name: str,
description: Optional[str] = None,
args: Optional[str] = None,
aliases: Optional[List[str]] = None,
):
self.name = name
self.description = description
self.args = args
self.aliases = aliases
self.hidden = False
class Module:
def __init__(self, name: str, path: str):
self.name = name
self.path = path
self.commands = {}
self.hidden = False
def add_command(
self,
command: str,
description: Optional[str] = None,
args: Optional[str] = None,
aliases: Optional[List[str]] = None,
) -> Command:
if command in self.commands:
raise ValueError(f"Command {command} already exists")
self.commands[command] = Command(command, description, args, aliases)
return self.commands[command]
def delete_command(self, command: str) -> None:
if command not in self.commands:
raise ValueError(f"Command {command} not found")
del self.commands[command]
def hide_command(self, command: str) -> None:
if command not in self.commands:
raise ValueError(f"Command {command} not found")
self.commands[command].hidden = True
def show_command(self, command: str) -> None:
if command not in self.commands:
raise ValueError(f"Command {command} not found")
self.commands[command].hidden = False
class ModuleHelp:
def __init__(self) -> None:
self.modules = {}
def add_module(self, name: str, path: str) -> Module:
self.modules[name] = Module(name, path)
return self.modules[name]
def delete_module(self, name: str) -> None:
del self.modules[name]
def hide_module(self, name: str) -> None:
if name not in self.modules:
raise ValueError(f"Module {name} not found")
self.modules[name].hidden = True
def show_module(self, name: str) -> None:
if name not in self.modules:
raise ValueError(f"Module {name} not found")
self.modules[name].hidden = False
def get_module(self, name: str) -> Module:
if name not in self.modules:
raise ValueError(f"Module {name} not found")
return self.modules[name]
def get_module_by_path(self, path: str) -> Module:
for module in self.modules.values():
if module.path == path:
return module
raise ValueError(f"Module with path {path} not found")
def help(self) -> List[str]:
prefix = "."
result = []
help_text = f"For more help on how to use a command, type {prefix}help [module]
\n\nAvailable Modules:\n"
for module_name, module in sorted(self.modules.items(), key=lambda x: x[0]):
help_text += f'• {module_name.title()}: {" ".join([f"{prefix + cmd_name}
" for cmd_name in module.commands.keys()])}\n'
if len(help_text) >= 2048:
result.append(help_text)
help_text = ""
help_text += f"\nThe number of modules in the userbot: {self.modules_count}\n"
help_text += f"The number of commands in the userbot: {self.commands_count}"
result.append(help_text)
return result
def module_help(self, module: str, full: bool = True) -> str:
if module not in self.modules:
raise ValueError(f"Module {module} not found")
prefix = "."
help_text = ""
if full:
help_text += f"Help for |{module}
|\n\n"
help_text += "Usage:\n"
for command in self.modules[module].commands.values():
help_text += f"{prefix}{command.name}"
if command.args:
help_text += f" {command.args}"
if command.description:
help_text += f"
— {command.description}\n"
return help_text
def command_help(self, command: str) -> str:
for module in self.modules.values():
for cmd in module.commands.values():
if cmd.name == command or (cmd.aliases and command in cmd.aliases):
command = cmd
break
else:
continue
break
else:
raise ValueError(f"Command {command} not found")
prefix = "."
help_text = f"Help for command {prefix}{command.name}
\n"
if command.aliases:
help_text += "Aliases: "
help_text += (
f"{' '.join([f'{prefix}{alias}
' for alias in command.aliases])}\n"
)
help_text += (
f"\nModule: {module.name} ({prefix}help {module.name}
)\n\n"
)
help_text += f"{prefix}{command.name}"
if command.args:
help_text += f" {command.args}"
help_text += "
"
if command.description:
help_text += f" — {command.description}"
return help_text
@property
def modules_count(self) -> int:
return len(self.modules)
@property
def commands_count(self) -> int:
return sum(len(module.commands) for module in self.modules.values())
def get_entity_url(
entity: Union[User, Chat],
openmessage: bool = False,
) -> str:
"""
Get link to object, if available
:param entity: Entity to get url of
:param openmessage: Use tg://openmessage link for users
:return: Link to object or empty string
"""
return (
(f"tg://openmessage?user_id={entity.id}" if openmessage else f"tg://user?id={entity.id}")
if isinstance(entity, User)
else (
f"tg://resolve?domain={entity.username}" if getattr(entity, "username", None) else ""
)
)
def get_message_link(
message: Message,
chat: Optional[Chat] = None,
) -> str:
"""
Get link to message
:param message: Message to get link of
:param chat: Chat, where message was sent
:return: Link to message
"""
if message.chat.type == ChatType.PRIVATE:
return f"tg://openmessage?user_id={message.chat.id}&message_id={message.id}"
return (
f"https://t.me/{chat.username}/{message.id}"
if getattr(chat, "username", False)
else f"https://t.me/c/{chat.id}/{message.id}"
)
async def shell_exec(
command: str,
executable: Optional[str] = None,
timeout: Optional[Union[int, float]] = None,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
) -> Tuple[int, str, str]:
"""Executes shell command and returns tuple with return code, decoded stdout and stderr"""
process = await asyncio.create_subprocess_shell(
cmd=command, stdout=stdout, stderr=stderr, shell=True, executable=executable
)
try:
stdout, stderr = await asyncio.wait_for(process.communicate(), timeout)
except asyncio.exceptions.TimeoutError as e:
process.kill()
raise e
return process.returncode, stdout.decode(), stderr.decode()