Add multi-account command ACL system
This commit is contained in:
Codex
347
app/commands/core.py
Normal file
347
app/commands/core.py
Normal file
@@ -0,0 +1,347 @@
|
||||
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()
|
||||
return CheckResult.deny(f"No access to {ctx.command.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")
|
||||
|
||||
|
||||
def combine_checks(*checks: CommandCheck) -> CommandCheck:
|
||||
async def _check(ctx: "CommandContext") -> CheckResult:
|
||||
for check in checks:
|
||||
result = await check.run(ctx)
|
||||
if not result.allowed:
|
||||
return result
|
||||
return CheckResult.allow()
|
||||
|
||||
names = "+".join(check.name for check in checks) or "combined"
|
||||
return CommandCheck(names, _check, "Combined checks")
|
||||
|
||||
|
||||
@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
|
||||
prefix: str = "."
|
||||
aliases: tuple[str, ...] = ()
|
||||
usage: str | None = None
|
||||
category: str = "general"
|
||||
checks: tuple[CommandCheck, ...] = ()
|
||||
|
||||
@property
|
||||
def is_shared(self) -> bool:
|
||||
return self.prefix == ".."
|
||||
|
||||
@property
|
||||
def display_name(self) -> str:
|
||||
return f"{self.prefix}{self.name}"
|
||||
|
||||
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.prefix = "."
|
||||
self.checks: list[CommandCheck] = []
|
||||
|
||||
def private(self) -> "CommandBuilder":
|
||||
self.prefix = "."
|
||||
return self
|
||||
|
||||
def shared(self) -> "CommandBuilder":
|
||||
self.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.prefix == "..":
|
||||
default_checks.append(shared_access_check())
|
||||
else:
|
||||
default_checks.append(owner_only_check())
|
||||
|
||||
return CommandDefinition(
|
||||
name=self.name,
|
||||
description=self.description,
|
||||
handler=handler,
|
||||
prefix=self.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 has_name(self, token: str) -> bool:
|
||||
return self.resolve(token) is not None
|
||||
|
||||
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)
|
||||
Reference in New Issue
Block a user