Add multi-account command ACL system

app/manager.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+from pyrogram import Client, idle
+
+from app.commands import DEFAULT_REGISTRY, CommandRegistry, RuntimeAclStore
+from app.config import AccountConfig, AppConfig
+from app.handlers import register_common_handlers
+
+log = logging.getLogger(__name__)
+
+
+@dataclass(slots=True)
+class AccountState:
+    config: AccountConfig
+    client: Client | None = None
+    acl_store: RuntimeAclStore | None = None
+    status: str = "pending"
+    error: str | None = None
+    started_at: datetime | None = None
+    me: Any | None = None
+
+
+@dataclass(slots=True)
+class AccountManager:
+    config: AppConfig
+    command_registry: CommandRegistry = field(init=False)
+    states: list[AccountState] = field(init=False)
+
+    def __post_init__(self) -> None:
+        self.command_registry = DEFAULT_REGISTRY
+        self.states = [AccountState(config=account) for account in self.config.accounts]
+
+    def _make_client(self, account: AccountConfig) -> Client:
+        api_id = account.api_id if account.api_id is not None else self.config.api_id
+        api_hash = account.api_hash if account.api_hash is not None else self.config.api_hash
+        if api_id is None or api_hash is None:
+            raise ValueError(f"Missing API credentials for account {account.name}")
+
+        workdir = account.workdir or self.config.session_dir
+        workdir.mkdir(parents=True, exist_ok=True)
+
+        client = Client(
+            account.session,
+            api_id=api_id,
+            api_hash=api_hash,
+            workdir=str(workdir),
+            session_string=account.session_string,
+            phone_number=account.phone_number,
+            password=account.password,
+            phone_code=account.phone_code,
+            bot_token=account.bot_token,
+            proxy=account.proxy,
+            in_memory=account.in_memory,
+            no_updates=account.no_updates,
+            takeout=account.takeout,
+            hide_password=account.hide_password,
+        )
+        return client
+
+    async def start_all(self) -> None:
+        for state in self.states:
+            await self.start_account(state)
+
+    async def start_account(self, state: AccountState) -> None:
+        account = state.config
+        state.status = "starting"
+        state.error = None
+
+        acl_path = self.config.session_dir / f"{account.session}.acl.json"
+        state.acl_store = RuntimeAclStore(acl_path)
+        client = self._make_client(account)
+        state.client = client
+        register_common_handlers(client, self, state)
+
+        try:
+            await client.start()
+            state.me = await client.get_me()
+            state.started_at = datetime.now(timezone.utc)
+            state.status = "running"
+            log.info("Started account %s (%s)", account.name, account.session)
+        except Exception as exc:
+            state.status = "failed"
+            state.error = str(exc)
+            log.exception("Failed to start account %s (%s)", account.name, account.session)
+            await self._safe_stop(client)
+
+    async def _safe_stop(self, client: Client) -> None:
+        try:
+            await client.stop()
+        except Exception:
+            pass
+
+    async def stop_all(self) -> None:
+        for state in reversed(self.states):
+            if state.client is None:
+                continue
+            if state.status not in {"running", "starting"}:
+                continue
+            try:
+                await state.client.stop()
+                state.status = "stopped"
+                log.info("Stopped account %s (%s)", state.config.name, state.config.session)
+            except Exception:
+                log.exception("Failed to stop account %s (%s)", state.config.name, state.config.session)
+
+    def build_status_report(self) -> str:
+        lines = ["Configured accounts:"]
+        for state in self.states:
+            account = state.config
+            parts = [
+                f"{account.name} [{account.session}]",
+                f"status={state.status}",
+            ]
+            if state.acl_store is not None:
+                stats = state.acl_store.stats()
+                parts.append(
+                    f"acl={stats['allow_rules']} allow / {stats['deny_rules']} deny"
+                )
+            if state.me is not None:
+                display_name = (
+                    getattr(state.me, "full_name", None)
+                    or getattr(state.me, "first_name", None)
+                    or getattr(state.me, "username", None)
+                )
+                if display_name:
+                    parts.append(f"user={display_name}")
+                parts.append(f"user_id={getattr(state.me, 'id', 'unknown')}")
+            if state.error:
+                parts.append(f"error={state.error}")
+            lines.append(" - " + ", ".join(parts))
+        return "\n".join(lines)
+
+    async def run(self) -> None:
+        await self.start_all()
+
+        if not any(state.status == "running" for state in self.states):
+            await self.stop_all()
+            raise RuntimeError("No accounts were started successfully")
+
+        try:
+            await idle()
+        finally:
+            await self.stop_all()