tlgr-cli 2.0.1__py3-none-any.whl

tlgr/daemon/sessions.py

@@ -0,0 +1,274 @@
+"""alias → `AccountSession`, with a lock per alias (COR-12).
+
+v1's `ensure_client()` was `if alias not in self._clients: await connect()`.
+Two concurrent requests for an account that was not yet connected both saw the
+miss, both constructed a `TelegramClient` on the same session file, and the
+second one's connect invalidated the first's auth key — `AUTH_KEY_DUPLICATED`,
+which Telegram treats as a compromised session and revokes.
+
+The fix is boring and complete: one `asyncio.Lock` per alias, a **double
+check** inside it (the whole point — the first waiter has already done the
+work by the time the second gets in), and a `flock` on the session file so
+that even a second *process* cannot do what the second coroutine could not.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from typing import Any
+
+from tlgr.core.accounts import AccountManager
+from tlgr.core.config import AppConfig
+from tlgr.core.errors import AccountNotFoundError, ConfigurationError
+from tlgr.core.identity import load_identity
+from tlgr.core.paths import TlgrPaths, validate_alias
+from tlgr.daemon.ratelimit import RateLimiter
+from tlgr.daemon.session import AccountSession, ClientOptions, SessionState, build_client
+
+log = logging.getLogger("tlgr.daemon.sessions")
+
+__all__ = ["SessionManager"]
+
+
+class SessionManager:
+    def __init__(
+        self,
+        paths: TlgrPaths,
+        config: AppConfig,
+        *,
+        accounts: AccountManager | None = None,
+        client_factory: Any = build_client,
+        on_session_ready: Any = None,
+    ) -> None:
+        self.paths = paths
+        self.config = config
+        self.accounts = accounts or AccountManager(paths.base)
+        self._factory = client_factory
+        self._on_ready = on_session_ready
+        self._sessions: dict[str, AccountSession] = {}
+        self._limiters: dict[str, RateLimiter] = {}
+        self._locks: dict[str, asyncio.Lock] = {}
+
+    # -- accessors ---------------------------------------------------------
+
+    def __contains__(self, alias: str) -> bool:
+        return alias in self._sessions
+
+    def get(self, alias: str) -> AccountSession | None:
+        return self._sessions.get(alias)
+
+    @property
+    def aliases(self) -> list[str]:
+        return list(self._sessions)
+
+    def limiter(self, alias: str) -> RateLimiter:
+        """The account's rate limiter, created on first use.
+
+        Created lazily rather than with the session because a flood deadline
+        outlives a connection: an account that failed to connect still owes the
+        wait it earned before it dropped.
+        """
+        if alias not in self._limiters:
+            validate_alias(alias)
+            buckets = {
+                name: (bucket.rate, bucket.burst) for name, bucket in self.config.rate.items()
+            }
+            self._limiters[alias] = RateLimiter(
+                buckets=buckets,
+                flood_path=self.paths.flood(alias),
+                persist=self.config.flood.persist,
+                sleep_threshold=self.config.flood.sleep_threshold,
+                max_wait=self.config.flood.max_wait,
+            )
+        return self._limiters[alias]
+
+    def _lock(self, alias: str) -> asyncio.Lock:
+        if alias not in self._locks:
+            self._locks[alias] = asyncio.Lock()
+        return self._locks[alias]
+
+    # -- construction ------------------------------------------------------
+
+    def _options(self, alias: str) -> ClientOptions:
+        api_id, api_hash = self.accounts.load_credentials(alias)
+        if not api_id or not api_hash:
+            raise ConfigurationError(
+                f"account {alias!r} has no API credentials. Run: tlgr account add --alias {alias}"
+            )
+        identity = load_identity(
+            self.paths.base,
+            device_model=self.config.identity.device_model,
+            system_version=self.config.identity.system_version,
+            lang_code=self.config.identity.lang_code,
+            system_lang_code=self.config.identity.system_lang_code,
+        )
+        return ClientOptions(
+            api_id=int(api_id),
+            api_hash=str(api_hash),
+            entity_cache_limit=self.config.limits.entity_cache,
+            request_retries=self.config.limits.request_retries,
+            connect_timeout=self.config.network.connect_timeout,
+            flood_sleep_threshold=self.config.flood.sleep_threshold,
+            device_model=identity.device_model,
+            system_version=identity.system_version,
+            app_version=identity.app_version,
+            lang_code=identity.lang_code,
+            system_lang_code=identity.system_lang_code,
+            proxy=_proxy_tuple(self.config.network.proxy),
+            use_ipv6=self.config.network.ipv6,
+            params=identity.params() if self.config.identity.tz_offset else {},
+        )
+
+    def _make(self, alias: str) -> AccountSession:
+        def on_state(state: str, reason: str, user_id: int | None) -> None:
+            with contextlib.suppress(Exception):
+                self.accounts.set_health(alias, state, reason=reason, user_id=user_id)
+
+        return AccountSession(
+            alias,
+            session_path=self.paths.session(alias),
+            lock_path=self.paths.session_lock(alias),
+            options=self._options(alias),
+            client_factory=self._factory,
+            on_state=on_state,
+            state_save_interval=self.config.daemon.state_save_interval,
+            presence=self.config.presence.mode,
+            resync_depth=self.config.daemon.resync_depth,
+            peers_path=self.paths.peers_db(alias),
+            dialog_scan_max=self.config.limits.dialog_scan_max,
+        )
+
+    # -- the one entry point -----------------------------------------------
+
+    async def ensure(self, alias: str) -> AccountSession:
+        """Return the account's session, connecting it on demand — exactly once."""
+        validate_alias(alias)
+        existing = self._sessions.get(alias)
+        if existing is not None:
+            return existing
+
+        async with self._lock(alias):
+            # The double check is the fix: by the time a second caller reaches
+            # here, the first has already built and started the session.
+            existing = self._sessions.get(alias)
+            if existing is not None:
+                return existing
+            if not self.accounts.get_account(alias) and not self.paths.account_dir(alias).exists():
+                raise AccountNotFoundError(f"no account named {alias!r}. Run: tlgr account list")
+            session = self._make(alias)
+            await session.start()
+            self._sessions[alias] = session
+            if self._on_ready is not None:
+                with contextlib.suppress(Exception):
+                    await self._on_ready(session)
+            return session
+
+    async def connect_all(self, aliases: list[str]) -> dict[str, str]:
+        """Connect the start-up list concurrently; failures do not block readiness.
+
+        The list is ordered (§6.1) and the results are reported per alias, so
+        one account with a revoked session cannot stop the others from working
+        — which is what "connect them in a loop and raise" did in v1.
+        """
+        results: dict[str, str] = {}
+
+        async def connect(alias: str) -> None:
+            try:
+                session = await self.ensure(alias)
+                results[alias] = session.state
+            except Exception as exc:
+                results[alias] = f"error: {exc}"
+                log.warning(
+                    "could not connect account %s: %s", alias, exc, extra={"account": alias}
+                )
+
+        await asyncio.gather(*(connect(alias) for alias in aliases))
+        return results
+
+    # -- shutdown ----------------------------------------------------------
+
+    async def stop_all(self, *, timeout: float = 10.0) -> None:
+        await asyncio.gather(
+            *(session.stop(timeout=timeout) for session in list(self._sessions.values())),
+            return_exceptions=True,
+        )
+        self._sessions.clear()
+        for limiter in self._limiters.values():
+            limiter.flood.save()
+
+    async def release(self, alias: str) -> None:
+        session = self._sessions.pop(alias, None)
+        if session is not None:
+            await session.stop()
+
+    # -- reporting ---------------------------------------------------------
+
+    def snapshot(self) -> list[dict[str, Any]]:
+        """One row per account, in the order they were connected."""
+        rows: list[dict[str, Any]] = []
+        for alias, session in self._sessions.items():
+            row = session.snapshot()
+            row.update(self.limiter(alias).snapshot())
+            rows.append(row)
+        for info in self.accounts.list_accounts():
+            if info.alias in self._sessions:
+                continue
+            health = info.health
+            rows.append(
+                {
+                    "alias": info.alias,
+                    "state": health.state if health.state != "unknown" else "not_connected",
+                    "reason": health.reason or None,
+                    "since": health.since or None,
+                    "user_id": info.user_id,
+                }
+            )
+        return rows
+
+    @property
+    def in_flight(self) -> int:
+        return sum(session.in_flight for session in self._sessions.values())
+
+    @property
+    def online(self) -> list[str]:
+        return [
+            alias
+            for alias, session in self._sessions.items()
+            if session.state == SessionState.ONLINE
+        ]
+
+
+def _proxy_tuple(spec: str) -> Any:
+    """`socks5://user:pass@host:1080` → the tuple Telethon wants, or None.
+
+    Returned as a dict for `python-socks`-style proxies because the tuple form
+    silently ignores authentication on some Telethon versions.
+    """
+    if not spec:
+        return None
+    from urllib.parse import urlparse
+
+    parsed = urlparse(spec)
+    scheme = (parsed.scheme or "").lower()
+    if not parsed.hostname or not parsed.port:
+        raise ConfigurationError(f"[network] proxy is not a URL with a host and port: {spec!r}")
+    if scheme in ("socks5", "socks4", "http"):
+        proxy: dict[str, Any] = {
+            "proxy_type": scheme,
+            "addr": parsed.hostname,
+            "port": int(parsed.port),
+            "rdns": True,
+        }
+        if parsed.username:
+            proxy["username"] = parsed.username
+        if parsed.password:
+            proxy["password"] = parsed.password
+        return proxy
+    if scheme == "mtproxy":
+        secret = (parsed.fragment or "").strip()
+        if not secret:
+            raise ConfigurationError("an mtproxy:// proxy needs its secret after a '#'")
+        return (parsed.hostname, int(parsed.port), secret)
+    raise ConfigurationError(f"unsupported proxy scheme {scheme!r}")

tlgr/daemon/singleton.py

@@ -0,0 +1,114 @@
+"""`flock`-based single ownership, for the daemon and for each session file.
+
+Two different things are guarded and they fail differently:
+
+* **the daemon** — one process per `~/.tlgr`. A second one exits *0* with a
+  message rather than 1, because under launchd a non-zero exit with
+  `KeepAlive.SuccessfulExit=false` means "respawn me", and v1's exit 1 on
+  "already running" produced an infinite respawn loop (COR-39).
+* **a session file** — one process per `accounts/<alias>/session.session`.
+  Two Telethon clients on one session file is how you earn
+  `AUTH_KEY_DUPLICATED`, which invalidates the session server-side; the
+  daemon holds this lock for the life of the account and the CLI never opens
+  a session file at all.
+
+`flock` is per open file description, so the lock is released by closing the
+fd and, crucially, is *not* inherited across a `Popen`. That is what lets the
+autostart probe hold a lock while spawning a child that takes a different one.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import fcntl
+import os
+from pathlib import Path
+from types import TracebackType
+
+__all__ = ["FileLock", "LockBusy"]
+
+
+class LockBusy(Exception):
+    """Somebody else holds the lock. Carries their pid when we can read it."""
+
+    def __init__(self, path: Path, pid: int | None = None) -> None:
+        self.path = path
+        self.pid = pid
+        holder = f" (held by pid {pid})" if pid else ""
+        super().__init__(f"{path} is locked by another process{holder}")
+
+
+class FileLock:
+    """An exclusive advisory lock on a file, held until `release()`.
+
+    The holder writes its pid into the file so that a refusal can name it —
+    "the lock is busy" is not an actionable message, "pid 8123 holds it" is.
+    """
+
+    def __init__(self, path: Path, *, mode: int = 0o600) -> None:
+        self.path = Path(path)
+        self.mode = mode
+        self._fd: int | None = None
+
+    @property
+    def held(self) -> bool:
+        return self._fd is not None
+
+    def holder_pid(self) -> int | None:
+        try:
+            content = self.path.read_text().strip()
+        except OSError:
+            return None
+        try:
+            return int(content)
+        except ValueError:
+            return None
+
+    def acquire(self, *, blocking: bool = False) -> None:
+        """Take the lock, raising `LockBusy` when someone else has it."""
+        self.path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
+        fd = os.open(str(self.path), os.O_CREAT | os.O_RDWR, self.mode)
+        flags = fcntl.LOCK_EX if blocking else fcntl.LOCK_EX | fcntl.LOCK_NB
+        try:
+            fcntl.flock(fd, flags)
+        except OSError as exc:
+            pid = self.holder_pid()
+            os.close(fd)
+            raise LockBusy(self.path, pid) from exc
+        self._fd = fd
+        self.write_pid()
+
+    def write_pid(self) -> None:
+        """(Re)write our pid into the locked file.
+
+        Called again after `daemonize()`: the pid written before the fork
+        belongs to a process that no longer exists, and a stale pid in a lock
+        file is worse than no pid at all.
+        """
+        if self._fd is None:
+            return
+        os.ftruncate(self._fd, 0)
+        os.lseek(self._fd, 0, os.SEEK_SET)
+        os.write(self._fd, f"{os.getpid()}\n".encode())
+        os.fsync(self._fd)
+
+    def release(self) -> None:
+        if self._fd is None:
+            return
+        with contextlib.suppress(OSError):
+            fcntl.flock(self._fd, fcntl.LOCK_UN)
+        with contextlib.suppress(OSError):
+            os.close(self._fd)
+        self._fd = None
+
+    def __enter__(self) -> FileLock:
+        self.acquire()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        self.release()

tlgr/daemon/stream.py

@@ -0,0 +1,193 @@
+"""NDJSON responses: `--all` walks, progress, and the event stream (§5.3/§5.4).
+
+Two rules make a stream trustworthy and both were missing in v1:
+
+* **exactly one `meta` first and exactly one `end` last.** A stream that ends
+  without an `end` frame is a failure, not a short result — the client raises
+  `RETRYABLE` rather than reporting "there were 400 items" when there were
+  4,120 and the connection dropped.
+* **the walk happens here, inside the daemon.** v1's `--all` was a client
+  loop that re-issued a request per page as fast as the socket allowed, with
+  no account pacing between pages (ROB-01). Walking server-side means the
+  account's own rate limiter sits between pages, and a 10k-dialog enumeration
+  is one request with backpressure instead of a hundred without.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+import time
+from collections.abc import AsyncIterator
+from typing import Any
+
+from aiohttp import web
+
+from tlgr.core.errors import classify, error_body_dict
+from tlgr.models.base import to_builtins
+from tlgr.transport.ndjson import dump_frame
+from tlgr.version import PROTOCOL
+
+log = logging.getLogger("tlgr.daemon.stream")
+
+__all__ = ["NdjsonResponse", "walk_pages"]
+
+#: §13.5 caps: a walk that has produced this much is stopped with a warning
+#: rather than running until the client gives up.
+MAX_WALK_ITEMS = 100_000
+MAX_WALK_SECONDS = 3600
+
+
+class NdjsonResponse:
+    """A chunked `application/x-ndjson` response with the framing rules baked in."""
+
+    def __init__(self, request: web.Request) -> None:
+        self._request = request
+        self._response: web.StreamResponse | None = None
+        self._ended = False
+
+    async def prepare(self, **meta: Any) -> None:
+        response = web.StreamResponse(
+            status=200,
+            headers={
+                "Content-Type": "application/x-ndjson",
+                "Cache-Control": "no-store",
+            },
+        )
+        response.enable_chunked_encoding()
+        await response.prepare(self._request)
+        self._response = response
+        await self.write({"type": "meta", "protocol": PROTOCOL, **meta})
+
+    async def write(self, frame: dict[str, Any]) -> None:
+        if self._response is None:  # pragma: no cover - prepare() comes first
+            raise RuntimeError("the NDJSON response was not prepared")
+        await self._response.write(dump_frame(frame))
+
+    async def end(self, **fields: Any) -> web.StreamResponse:
+        """Write the single `end` frame and close. Idempotent."""
+        if self._response is None:  # pragma: no cover
+            raise RuntimeError("the NDJSON response was not prepared")
+        if not self._ended:
+            self._ended = True
+            await self.write({"type": "end", **fields})
+            with contextlib.suppress(Exception):
+                await self._response.write_eof()
+        return self._response
+
+    async def fail(self, exc: BaseException, *, account: str = "") -> web.StreamResponse:
+        body = error_body_dict(classify(exc, account=account or None))
+        return await self.end(ok=False, error=body)
+
+    @property
+    def started(self) -> bool:
+        return self._response is not None
+
+
+async def walk_pages(
+    pages: AsyncIterator[Any],
+    stream: NdjsonResponse,
+    *,
+    limiter: Any = None,
+    rate_class: str = "read",
+    max_items: int = MAX_WALK_ITEMS,
+    max_seconds: int = MAX_WALK_SECONDS,
+) -> int:
+    """Stream every item from an async page iterator, paced by the limiter.
+
+    Returns the number of items written. The caps are not a policy about what
+    a user may ask for; they are the difference between a walk that ends and
+    one that quietly becomes a permanent background load on the account.
+    """
+    started = time.monotonic()
+    count = 0
+    async for page in pages:
+        # `_attr` checks the dict case first: `getattr({}, "items")` is the
+        # dict *method*, which iterates into a very confusing TypeError.
+        items = _attr(page, "items")
+        for item in items or ():
+            count += 1
+            await stream.write({"type": "item", "seq": count, "data": to_builtins(item)})
+            if count >= max_items:
+                await stream.write(
+                    {
+                        "type": "page",
+                        "has_more": True,
+                        "next_cursor": _cursor_of(page),
+                        "fetched": count,
+                        "truncated": f"stopped at the {max_items} item cap",
+                    }
+                )
+                return count
+        await stream.write(
+            {
+                "type": "page",
+                "has_more": bool(_attr(page, "has_more")),
+                "next_cursor": _cursor_of(page),
+                "fetched": count,
+                "elapsed_ms": int((time.monotonic() - started) * 1000),
+            }
+        )
+        if not _attr(page, "has_more"):
+            break
+        if time.monotonic() - started > max_seconds:
+            await stream.write(
+                {"type": "page", "has_more": True, "truncated": "stopped at the time cap"}
+            )
+            break
+        if limiter is not None:
+            # Backpressure lives here: the next page waits for the account's
+            # own bucket, so a walk paces itself instead of racing the server.
+            await limiter.acquire(rate_class)
+    return count
+
+
+def _attr(page: Any, name: str) -> Any:
+    if isinstance(page, dict):
+        return page.get(name)
+    return getattr(page, name, None)
+
+
+def _cursor_of(page: Any) -> Any:
+    return _attr(page, "next_cursor")
+
+
+async def pump_events(
+    stream: NdjsonResponse,
+    subscriber: Any,
+    *,
+    heartbeat: float = 15.0,
+    timeout: float = 3600.0,
+    shutdown: asyncio.Event | None = None,
+) -> str:
+    """Deliver events until the timeout, the client leaves, or we shut down.
+
+    Returns the reason, which the caller puts in the `end` frame — "the
+    stream ended" without saying why is exactly the ambiguity that makes a
+    consumer guess whether to reconnect.
+    """
+    deadline = time.monotonic() + timeout
+    while True:
+        if shutdown is not None and shutdown.is_set():
+            return "shutdown"
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            return "timeout"
+        try:
+            event = await asyncio.wait_for(
+                subscriber.queue.get(), timeout=min(heartbeat, remaining)
+            )
+        except (TimeoutError, asyncio.TimeoutError):
+            await stream.write({"type": "heartbeat", "ts": _now()})
+            continue
+        lag = subscriber.take_lag()
+        if lag:
+            await stream.write({"type": "lag", "dropped": lag})
+        await stream.write(to_builtins(event))
+
+
+def _now() -> str:
+    from datetime import datetime, timezone
+
+    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")