tlgr-cli 2.0.1__py3-none-any.whl

tlgr/core/output.py

@@ -0,0 +1,251 @@
+"""Output formatters for JSON, plain (TSV), and human-readable modes."""
+
+from __future__ import annotations
+
+import base64
+import json
+import sys
+from collections.abc import Sequence
+from typing import Any
+
+
+def _tsv_escape(value: Any) -> str:
+    s = str(value) if value is not None else ""
+    return s.replace("\t", " ").replace("\n", " ")
+
+
+# ---------------------------------------------------------------------------
+# JSON transforms (--results-only, --select)
+# ---------------------------------------------------------------------------
+
+_ENVELOPE_KEYS = frozenset(
+    {
+        "next_page_token",
+        "nextPageToken",
+        "next_cursor",
+        "has_more",
+        "count",
+        "total",
+        "query",
+        "dry_run",
+        "dryRun",
+        "op",
+        "action",
+    }
+)
+
+
+def _unwrap_primary(data: Any) -> Any:
+    """Strip envelope metadata and return only the primary result."""
+    if not isinstance(data, dict):
+        return data
+    if "results" in data:
+        return data["results"]
+
+    candidates = [k for k in data if k not in _ENVELOPE_KEYS]
+    if len(candidates) == 1:
+        return data[candidates[0]]
+
+    for k in candidates:
+        if isinstance(data[k], list):
+            return data[k]
+
+    return data
+
+
+def _get_at_path(obj: Any, path: str) -> tuple[Any, bool]:
+    """Traverse dot-delimited path into a nested dict/list."""
+    segments = [s.strip() for s in path.split(".") if s.strip()]
+    cur = obj
+    for seg in segments:
+        if isinstance(cur, dict):
+            if seg not in cur:
+                return None, False
+            cur = cur[seg]
+        elif isinstance(cur, list):
+            try:
+                cur = cur[int(seg)]
+            except (ValueError, IndexError):
+                return None, False
+        else:
+            return None, False
+    return cur, True
+
+
+def _select_fields(data: Any, fields: list[str]) -> Any:
+    """Project only selected fields from data."""
+    if isinstance(data, list):
+        return [_select_from_item(item, fields) for item in data]
+    return _select_from_item(data, fields)
+
+
+def _select_from_item(item: Any, fields: list[str]) -> Any:
+    if not isinstance(item, dict):
+        return item
+    out: dict[str, Any] = {}
+    for f in fields:
+        val, found = _get_at_path(item, f)
+        if found:
+            out[f] = val
+    return out
+
+
+def apply_json_transforms(
+    data: Any,
+    *,
+    results_only: bool = False,
+    select: str | None = None,
+) -> Any:
+    """Apply --results-only and --select transforms to JSON data."""
+    if results_only:
+        data = _unwrap_primary(data)
+    if select:
+        fields = [f.strip() for f in select.split(",") if f.strip()]
+        if fields:
+            data = _select_fields(data, fields)
+    return data
+
+
+# ---------------------------------------------------------------------------
+# Core output functions
+# ---------------------------------------------------------------------------
+
+
+def output_json(
+    data: Any,
+    *,
+    flood_wait: int | None = None,
+    results_only: bool = False,
+    select: str | None = None,
+) -> None:
+    """Write JSON to stdout."""
+    if flood_wait:
+        if isinstance(data, dict):
+            data["flood_wait"] = flood_wait
+        else:
+            data = {"result": data, "flood_wait": flood_wait}
+
+    data = apply_json_transforms(data, results_only=results_only, select=select)
+    json.dump(data, sys.stdout, default=str, ensure_ascii=False)
+    sys.stdout.write("\n")
+    sys.stdout.flush()
+
+
+def output_plain(rows: Sequence[dict[str, Any]], columns: Sequence[str]) -> None:
+    """Write TSV to stdout (no colors, stable for piping)."""
+    print("\t".join(columns))
+    for row in rows:
+        print("\t".join(_tsv_escape(row.get(c)) for c in columns))
+    sys.stdout.flush()
+
+
+def output_human(
+    rows: Sequence[dict[str, Any]],
+    columns: Sequence[str],
+    *,
+    headers: Sequence[str] | None = None,
+) -> None:
+    """Write space-padded columns to stdout (kubectl / docker style)."""
+    display_headers = headers or columns
+    cells = [[str(row.get(c, "")) for c in columns] for row in rows]
+
+    widths = [len(h) for h in display_headers]
+    for cell_row in cells:
+        for i, v in enumerate(cell_row):
+            widths[i] = max(widths[i], len(v))
+
+    gap = "   "
+    header_line = gap.join(
+        h.upper().ljust(w) for h, w in zip(display_headers, widths, strict=False)
+    )
+    print(header_line.rstrip())
+    for cell_row in cells:
+        line = gap.join(v.ljust(w) for v, w in zip(cell_row, widths, strict=False))
+        print(line.rstrip())
+    sys.stdout.flush()
+
+
+def output_result(
+    data: Any,
+    *,
+    fmt: str = "human",
+    columns: Sequence[str] | None = None,
+    headers: Sequence[str] | None = None,
+    flood_wait: int | None = None,
+    results_only: bool = False,
+    select: str | None = None,
+) -> None:
+    """Dispatch to the correct output formatter.
+
+    *fmt* is one of ``"json"``, ``"plain"``, or ``"human"``.
+    For ``"json"`` *data* is emitted as-is.
+    For ``"plain"`` and ``"human"`` *data* must be a list of dicts
+    and *columns* selects which keys to display.
+    """
+    if fmt == "json":
+        output_json(data, flood_wait=flood_wait, results_only=results_only, select=select)
+        return
+
+    if not isinstance(data, list):
+        data = [data] if isinstance(data, dict) else [{"result": data}]
+
+    cols = columns or (list(data[0].keys()) if data else ["result"])
+
+    if fmt == "plain":
+        output_plain(data, cols)
+    else:
+        output_human(data, cols, headers=headers)
+
+
+def emit(ctx_obj: dict[str, Any], data: Any, **kwargs: Any) -> None:
+    """Convenience wrapper: ``output_result`` with global transforms from *ctx.obj*."""
+    kwargs.setdefault("fmt", ctx_obj.get("fmt", "human"))
+    kwargs.setdefault("results_only", ctx_obj.get("results_only", False))
+    kwargs.setdefault("select", ctx_obj.get("select"))
+    output_result(data, **kwargs)
+
+
+# ---------------------------------------------------------------------------
+# Cursor-based pagination helpers
+# ---------------------------------------------------------------------------
+
+
+def encode_cursor(state: dict[str, Any]) -> str:
+    """Encode pagination state as an opaque base64 cursor token."""
+    raw = json.dumps(state, separators=(",", ":"), sort_keys=True)
+    return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=")
+
+
+def decode_cursor(token: str | None) -> dict[str, Any]:
+    """Decode a cursor token back to pagination state. Returns {} on invalid input."""
+    if not token:
+        return {}
+    try:
+        padded = token + "=" * (-len(token) % 4)
+        raw = base64.urlsafe_b64decode(padded).decode()
+        return json.loads(raw)
+    except Exception:
+        return {}
+
+
+def add_pagination(
+    envelope: dict[str, Any],
+    items: list[Any],
+    limit: int,
+    cursor_state: dict[str, Any],
+    has_more: bool | None = None,
+) -> dict[str, Any]:
+    """Add ``has_more`` and ``next_cursor`` to a JSON envelope.
+
+    ``len(items) >= limit`` is a heuristic for server-paginated endpoints,
+    which cannot know whether more rows exist without asking again. Callers
+    that hold the full result set and slice it themselves know the exact
+    answer — they pass it as ``has_more`` instead of guessing, so the cursor
+    stops being emitted once the list is exhausted.
+    """
+    if has_more is None:
+        has_more = len(items) >= limit
+    envelope["has_more"] = has_more
+    if has_more:
+        envelope["next_cursor"] = encode_cursor(cursor_state)
+    return envelope

tlgr/core/pagination.py

@@ -0,0 +1,227 @@
+"""Opaque, versioned, op-bound cursors.
+
+A cursor is state the caller is not supposed to read, edit or reuse elsewhere.
+v1's cursor was plain base64 JSON that decoded to `{}` on any problem — a
+truncated token silently restarted the walk from message 0, which looks like
+"the chat only has these 40 messages" rather than like an error.
+
+Here a cursor carries its version, the op it belongs to, an account
+fingerprint and an expiry, and is signed. The HMAC is integrity, not secrecy:
+it exists so that a hand-edited cursor produces `USAGE: invalid cursor`
+instead of a plausible wrong answer.
+"""
+
+from __future__ import annotations
+
+import base64
+import binascii
+import hashlib
+import hmac
+import os
+import secrets
+import time
+from enum import Enum
+from pathlib import Path
+from typing import Any, TypeVar
+
+import msgspec
+
+from tlgr.core.errors import UsageError
+from tlgr.models.page import Page
+
+__all__ = [
+    "CURSOR_VERSION",
+    "PageKind",
+    "build_page",
+    "cursor_key",
+    "decode_cursor",
+    "encode_cursor",
+]
+
+CURSOR_VERSION = 1
+
+#: Default lifetimes. A LOCAL cursor indexes into a materialised snapshot that
+#: goes stale fast; a server-side offset stays meaningful much longer.
+_TTL_LOCAL = 3600
+_TTL_DEFAULT = 86400
+
+T = TypeVar("T")
+
+
+class PageKind(str, Enum):
+    """Which offset state a paginated op carries (ARCHITECTURE §3.5)."""
+
+    HISTORY = "HISTORY"
+    SEARCH = "SEARCH"
+    RATE = "RATE"
+    DIALOGS = "DIALOGS"
+    PARTICIPANTS = "PARTICIPANTS"
+    LOCAL = "LOCAL"
+
+
+#: PageKinds whose offsets are dates as well as ids, so `--since/--until` mean
+#: something and the generator injects them.
+DATE_OFFSET_KINDS = frozenset({PageKind.HISTORY, PageKind.SEARCH, PageKind.DIALOGS})
+
+
+def _config_dir() -> Path:
+    """Where the cursor key lives. `TLGR_HOME` exists so tests never touch ~."""
+    override = os.environ.get("TLGR_HOME", "").strip()
+    if override:
+        return Path(override)
+    from tlgr.core.config import CONFIG_DIR
+
+    return Path(CONFIG_DIR)
+
+
+def cursor_key(base: Path | None = None) -> bytes:
+    """The signing key, generated once at 0600.
+
+    Rotating it invalidates outstanding cursors, which is the correct
+    behaviour: they describe a walk of data this installation no longer
+    guarantees anything about.
+    """
+    directory = base or _config_dir()
+    path = directory / "cursor.key"
+    try:
+        return path.read_bytes()
+    except FileNotFoundError:
+        pass
+    directory.mkdir(parents=True, exist_ok=True)
+    key = secrets.token_bytes(32)
+    # Write through a private temp file so the key is never briefly world-readable.
+    tmp = path.with_suffix(".key.tmp")
+    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+    try:
+        os.write(fd, key)
+    finally:
+        os.close(fd)
+    os.replace(tmp, path)
+    return key
+
+
+def account_fingerprint(account: str) -> str:
+    """First 8 hex of sha256(alias) — enough to catch a cursor crossing accounts.
+
+    The alias itself is not embedded because a cursor is pasted into shell
+    history, logs and bug reports.
+    """
+    return hashlib.sha256(account.encode()).hexdigest()[:8]
+
+
+def _b64(data: bytes) -> str:
+    return base64.urlsafe_b64encode(data).decode().rstrip("=")
+
+
+def _unb64(text: str) -> bytes:
+    return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))
+
+
+def _sign(payload: bytes, key: bytes) -> str:
+    return _b64(hmac.new(key, payload, hashlib.sha256).digest()[:16])
+
+
+def encode_cursor(
+    *,
+    op: str,
+    kind: PageKind | str,
+    state: dict[str, Any],
+    account: str = "",
+    ttl: int | None = None,
+    key: bytes | None = None,
+    now: int | None = None,
+) -> str:
+    """Sign pagination *state* into a token bound to this op and account."""
+    kind_name = kind.value if isinstance(kind, PageKind) else str(kind)
+    if ttl is None:
+        ttl = _TTL_LOCAL if kind_name == PageKind.LOCAL.value else _TTL_DEFAULT
+    payload = msgspec.json.encode(
+        {
+            "v": CURSOR_VERSION,
+            "op": op,
+            "kind": kind_name,
+            "acct": account_fingerprint(account),
+            "st": state,
+            "exp": int(now or time.time()) + ttl,
+        }
+    )
+    return f"{_b64(payload)}.{_sign(payload, key or cursor_key())}"
+
+
+def decode_cursor(
+    token: str,
+    *,
+    op: str,
+    kind: PageKind | str | None = None,
+    account: str = "",
+    key: bytes | None = None,
+    now: int | None = None,
+) -> dict[str, Any]:
+    """Validate a token and return its state, or raise USAGE.
+
+    Every rejection is a USAGE error rather than a silent restart: a cursor
+    that cannot be trusted must stop the walk, not quietly begin a new one.
+    """
+
+    def reject(why: str) -> UsageError:
+        return UsageError(f"invalid cursor: {why}", field="cursor")
+
+    head, dot, signature = token.partition(".")
+    if not dot:
+        raise reject("not a tlgr cursor (missing signature)")
+    try:
+        payload = _unb64(head)
+    except (binascii.Error, ValueError) as exc:
+        raise reject("corrupt encoding") from exc
+    if not hmac.compare_digest(signature, _sign(payload, key or cursor_key())):
+        raise reject("signature does not match (truncated or hand-edited?)")
+
+    try:
+        data = msgspec.json.decode(payload, type=dict)
+    except msgspec.DecodeError as exc:
+        raise reject("corrupt payload") from exc
+
+    if data.get("v") != CURSOR_VERSION:
+        raise reject(f"version {data.get('v')!r} is not supported by this tlgr")
+    if data.get("op") != op:
+        raise reject(f"it belongs to {data.get('op')!r}, not {op!r}")
+    if kind is not None:
+        want = kind.value if isinstance(kind, PageKind) else str(kind)
+        if data.get("kind") != want:
+            raise reject(f"it paginates {data.get('kind')!r}, not {want!r}")
+    if data.get("acct") != account_fingerprint(account):
+        raise reject("it was issued for a different account")
+    expiry = data.get("exp")
+    if not isinstance(expiry, int) or expiry < int(now or time.time()):
+        raise reject("it has expired; start the listing again")
+
+    state = data.get("st")
+    if not isinstance(state, dict):
+        raise reject("corrupt state")
+    return state
+
+
+def build_page(
+    items: list[T],
+    *,
+    op: str,
+    kind: PageKind | str,
+    state: dict[str, Any] | None = None,
+    account: str = "",
+    has_more: bool | None = None,
+    limit: int | None = None,
+    total: int | None = None,
+    key: bytes | None = None,
+) -> Page[T]:
+    """Assemble a `Page[T]`, emitting a cursor only when there is more to fetch.
+
+    `len(items) >= limit` is the fallback guess for server-paginated endpoints,
+    which cannot know whether more rows exist without asking again. A caller
+    holding the full set knows the exact answer and passes `has_more`.
+    """
+    if has_more is None:
+        has_more = limit is not None and len(items) >= limit
+    page: Page[T] = Page(items=items, has_more=has_more, total=total)
+    if has_more and state is not None:
+        page.next_cursor = encode_cursor(op=op, kind=kind, state=state, account=account, key=key)
+    return page