344 lines
11 KiB
Python
344 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass, field
|
|
import inspect
|
|
import shlex
|
|
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
|
|
|
if TYPE_CHECKING:
|
|
from app.commands.acl import RuntimeAclStore
|
|
from app.manager import AccountManager, AccountState
|
|
from pyrogram import Client
|
|
from pyrogram.types import Message
|
|
|
|
|
|
CommandHandler = Callable[["CommandContext"], Awaitable[None]]
|
|
CommandCheckFn = Callable[["CommandContext"], Any]
|
|
|
|
|
|
def normalize_command_name(value: str) -> str:
|
|
token = value.strip().lower()
|
|
token = token.lstrip(".")
|
|
if "@" in token:
|
|
token = token.split("@", 1)[0]
|
|
return token
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class CheckResult:
|
|
allowed: bool
|
|
reason: str | None = None
|
|
|
|
@classmethod
|
|
def allow(cls) -> "CheckResult":
|
|
return cls(True, None)
|
|
|
|
@classmethod
|
|
def deny(cls, reason: str) -> "CheckResult":
|
|
return cls(False, reason)
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class CommandCheck:
|
|
name: str
|
|
callback: CommandCheckFn
|
|
description: str = ""
|
|
|
|
async def run(self, ctx: "CommandContext") -> CheckResult:
|
|
result = self.callback(ctx)
|
|
if inspect.isawaitable(result):
|
|
result = await result
|
|
if isinstance(result, CheckResult):
|
|
return result
|
|
if result is None or result is True:
|
|
return CheckResult.allow()
|
|
if result is False:
|
|
return CheckResult.deny(f"Check failed: {self.name}")
|
|
if isinstance(result, str):
|
|
return CheckResult.deny(result)
|
|
raise TypeError(f"Unsupported check result from {self.name}: {type(result)!r}")
|
|
|
|
|
|
def owner_only_check() -> CommandCheck:
|
|
async def _check(ctx: "CommandContext") -> CheckResult:
|
|
if ctx.owner_id is None:
|
|
return CheckResult.deny("Account owner is not available yet")
|
|
if ctx.actor_id != ctx.owner_id:
|
|
return CheckResult.deny("Only the account owner can use this command")
|
|
return CheckResult.allow()
|
|
|
|
return CommandCheck("owner_only", _check, "Allow only the account owner")
|
|
|
|
|
|
def shared_access_check() -> CommandCheck:
|
|
async def _check(ctx: "CommandContext") -> CheckResult:
|
|
if ctx.owner_id is not None and ctx.actor_id == ctx.owner_id:
|
|
return CheckResult.allow()
|
|
if ctx.actor_id is None:
|
|
return CheckResult.deny("This command requires a Telegram user identity")
|
|
store = ctx.acl_store
|
|
if store is None:
|
|
return CheckResult.deny("Permission store is not initialized")
|
|
if store.is_allowed(ctx.actor_id, ctx.command.name):
|
|
return CheckResult.allow()
|
|
required_prefix = ctx.command.expected_prefix(ctx.is_owner)
|
|
return CheckResult.deny(f"No access to {required_prefix}{ctx.command.name}")
|
|
|
|
return CommandCheck("shared_access", _check, "Allow owner or ACL-granted users")
|
|
|
|
|
|
def require_reply_check() -> CommandCheck:
|
|
async def _check(ctx: "CommandContext") -> CheckResult:
|
|
if ctx.replied_user_id is None:
|
|
return CheckResult.deny("Reply to a user's message first")
|
|
return CheckResult.allow()
|
|
|
|
return CommandCheck("require_reply", _check, "Require a replied-to user message")
|
|
|
|
|
|
def require_args_check() -> CommandCheck:
|
|
async def _check(ctx: "CommandContext") -> CheckResult:
|
|
if not ctx.args.strip():
|
|
return CheckResult.deny("This command expects arguments")
|
|
return CheckResult.allow()
|
|
|
|
return CommandCheck("require_args", _check, "Require non-empty arguments")
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class CommandContext:
|
|
client: "Client"
|
|
manager: "AccountManager"
|
|
account: "AccountState"
|
|
message: "Message"
|
|
command: "CommandDefinition"
|
|
prefix: str
|
|
raw_text: str
|
|
args: str
|
|
alias: str
|
|
|
|
@property
|
|
def actor(self) -> Any | None:
|
|
return getattr(self.message, "from_user", None)
|
|
|
|
@property
|
|
def actor_id(self) -> int | None:
|
|
actor = self.actor
|
|
return getattr(actor, "id", None)
|
|
|
|
@property
|
|
def owner_id(self) -> int | None:
|
|
me = self.account.me
|
|
return getattr(me, "id", None)
|
|
|
|
@property
|
|
def is_owner(self) -> bool:
|
|
return self.owner_id is not None and self.actor_id == self.owner_id
|
|
|
|
@property
|
|
def acl_store(self) -> RuntimeAclStore | None:
|
|
return self.account.acl_store
|
|
|
|
@property
|
|
def replied_message(self) -> Any | None:
|
|
return getattr(self.message, "reply_to_message", None)
|
|
|
|
@property
|
|
def replied_user(self) -> Any | None:
|
|
reply = self.replied_message
|
|
if reply is None:
|
|
return None
|
|
return getattr(reply, "from_user", None)
|
|
|
|
@property
|
|
def replied_user_id(self) -> int | None:
|
|
user = self.replied_user
|
|
return getattr(user, "id", None)
|
|
|
|
@property
|
|
def argv(self) -> list[str]:
|
|
if not self.args.strip():
|
|
return []
|
|
try:
|
|
return shlex.split(self.args)
|
|
except ValueError:
|
|
return self.args.split()
|
|
|
|
async def reply(self, text: str, **kwargs: Any) -> Any:
|
|
return await self.message.reply_text(text, **kwargs)
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class CommandDefinition:
|
|
name: str
|
|
description: str
|
|
handler: CommandHandler
|
|
owner_prefix: str = "."
|
|
guest_prefix: str | None = None
|
|
aliases: tuple[str, ...] = ()
|
|
usage: str | None = None
|
|
category: str = "general"
|
|
checks: tuple[CommandCheck, ...] = ()
|
|
|
|
@property
|
|
def is_shared(self) -> bool:
|
|
return self.guest_prefix is not None
|
|
|
|
@property
|
|
def display_name(self) -> str:
|
|
if self.guest_prefix is None:
|
|
return f"{self.owner_prefix}{self.name}"
|
|
return f"{self.owner_prefix}{self.name} / {self.guest_prefix}{self.name}"
|
|
|
|
def expected_prefix(self, is_owner: bool) -> str:
|
|
if self.guest_prefix is None:
|
|
return self.owner_prefix
|
|
return self.owner_prefix if is_owner else self.guest_prefix
|
|
|
|
def help_line(self, ctx: CommandContext | None = None) -> str:
|
|
usage = self.usage or self.display_name
|
|
details: list[str] = []
|
|
if self.aliases:
|
|
details.append(f"aliases: {', '.join(self.aliases)}")
|
|
details.append("shared" if self.is_shared else "owner-only")
|
|
if ctx is not None and self.is_shared and ctx.acl_store is not None:
|
|
allowed = ctx.acl_store.count_allowed_users(self.name)
|
|
denied = ctx.acl_store.count_denied_users(self.name)
|
|
details.append(f"allow={allowed}")
|
|
details.append(f"deny={denied}")
|
|
suffix = f" ({'; '.join(details)})" if details else ""
|
|
return f"{usage} - {self.description}{suffix}"
|
|
|
|
|
|
class CommandBuilder:
|
|
def __init__(
|
|
self,
|
|
registry: "CommandRegistry",
|
|
name: str,
|
|
description: str,
|
|
*,
|
|
aliases: tuple[str, ...] = (),
|
|
usage: str | None = None,
|
|
category: str = "general",
|
|
) -> None:
|
|
self.registry = registry
|
|
self.name = normalize_command_name(name)
|
|
self.description = description.strip()
|
|
self.aliases = [normalize_command_name(alias) for alias in aliases]
|
|
self.usage = usage
|
|
self.category = category
|
|
self.owner_prefix = "."
|
|
self.guest_prefix: str | None = None
|
|
self.checks: list[CommandCheck] = []
|
|
|
|
def private(self) -> "CommandBuilder":
|
|
self.guest_prefix = None
|
|
return self
|
|
|
|
def shared(self) -> "CommandBuilder":
|
|
self.guest_prefix = ".."
|
|
return self
|
|
|
|
def alias(self, *aliases: str) -> "CommandBuilder":
|
|
for alias in aliases:
|
|
normalized = normalize_command_name(alias)
|
|
if normalized and normalized not in self.aliases and normalized != self.name:
|
|
self.aliases.append(normalized)
|
|
return self
|
|
|
|
def check(self, *checks: CommandCheck) -> "CommandBuilder":
|
|
self.checks.extend(checks)
|
|
return self
|
|
|
|
def owner_only(self) -> "CommandBuilder":
|
|
return self.check(owner_only_check())
|
|
|
|
def require_reply(self) -> "CommandBuilder":
|
|
return self.check(require_reply_check())
|
|
|
|
def require_args(self) -> "CommandBuilder":
|
|
return self.check(require_args_check())
|
|
|
|
def build(self, handler: CommandHandler) -> CommandDefinition:
|
|
default_checks: list[CommandCheck] = []
|
|
if self.guest_prefix is not None:
|
|
default_checks.append(shared_access_check())
|
|
else:
|
|
default_checks.append(owner_only_check())
|
|
|
|
return CommandDefinition(
|
|
name=self.name,
|
|
description=self.description,
|
|
handler=handler,
|
|
owner_prefix=self.owner_prefix,
|
|
guest_prefix=self.guest_prefix,
|
|
aliases=tuple(self.aliases),
|
|
usage=self.usage,
|
|
category=self.category,
|
|
checks=tuple(default_checks + self.checks),
|
|
)
|
|
|
|
def __call__(self, handler: CommandHandler) -> CommandHandler:
|
|
self.registry.register(self.build(handler))
|
|
return handler
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class CommandRegistry:
|
|
_commands: dict[str, CommandDefinition] = field(default_factory=dict)
|
|
_aliases: dict[str, str] = field(default_factory=dict)
|
|
_order: list[str] = field(default_factory=list)
|
|
|
|
def command(
|
|
self,
|
|
name: str,
|
|
description: str,
|
|
*,
|
|
aliases: tuple[str, ...] = (),
|
|
usage: str | None = None,
|
|
category: str = "general",
|
|
) -> CommandBuilder:
|
|
return CommandBuilder(
|
|
self,
|
|
name,
|
|
description,
|
|
aliases=aliases,
|
|
usage=usage,
|
|
category=category,
|
|
)
|
|
|
|
def register(self, definition: CommandDefinition) -> None:
|
|
if definition.name in self._commands:
|
|
raise ValueError(f"Command already registered: {definition.name}")
|
|
|
|
for alias in definition.aliases:
|
|
if alias in self._aliases or alias in self._commands:
|
|
raise ValueError(f"Command alias already registered: {alias}")
|
|
|
|
self._commands[definition.name] = definition
|
|
self._order.append(definition.name)
|
|
|
|
for alias in definition.aliases:
|
|
self._aliases[alias] = definition.name
|
|
|
|
def resolve(self, token: str) -> CommandDefinition | None:
|
|
normalized = normalize_command_name(token)
|
|
canonical = self._aliases.get(normalized, normalized)
|
|
return self._commands.get(canonical)
|
|
|
|
def iter_commands(self) -> list[CommandDefinition]:
|
|
return [self._commands[name] for name in self._order]
|
|
|
|
def render_help(self, ctx: CommandContext) -> str:
|
|
grouped: dict[str, list[CommandDefinition]] = defaultdict(list)
|
|
for definition in self.iter_commands():
|
|
grouped[definition.category].append(definition)
|
|
|
|
lines: list[str] = ["Available commands:"]
|
|
for category in sorted(grouped):
|
|
lines.append(f"[{category}]")
|
|
for definition in grouped[category]:
|
|
lines.append(f"- {definition.help_line(ctx)}")
|
|
return "\n".join(lines)
|