waidrop-cli 1.0.0__py3-none-any.whl

waidrop_cli/fetch.py

@@ -0,0 +1,214 @@
+"""Range planning, Dropbox downloads, decoding, filtering and dedupe."""
+
+from __future__ import annotations
+
+import json
+import re
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from datetime import date, datetime, timedelta
+
+from cryptography.exceptions import InvalidTag
+
+from .archive import ArchiveCodec
+from .chats import ChatEntry
+from .config import DropboxConfig, Window
+from .errors import GuardError, NoMatchesError, RemoteError
+
+MAX_WORKERS = 4
+PART_RE = re.compile(r"^part-(\d+)(?:-(\d+))?-(\d+)\.(enc|ndjson)$")
+
+
+@dataclass(frozen=True)
+class FetchTask:
+    chat_hash: str
+    key: str
+
+
+def effective_ts(record: dict):
+    value = record.get("ts")
+    if value is None:
+        value = record.get("sortTs")
+    if value is None:
+        value = record.get("capturedAt")
+    return value if isinstance(value, (int, float)) else None
+
+
+def _media_stored(record: dict) -> bool:
+    media = record.get("media")
+    return isinstance(media, dict) and media.get("status") == "stored"
+
+
+def _captured_at(record: dict) -> float:
+    value = record.get("capturedAt")
+    return value if isinstance(value, (int, float)) else 0
+
+
+def _prefer(candidate: dict, current: dict) -> bool:
+    candidate_store = candidate.get("source") == "store"
+    current_store = current.get("source") == "store"
+    if candidate_store != current_store:
+        return candidate_store
+    if _media_stored(candidate) != _media_stored(current):
+        return _media_stored(candidate)
+    return _captured_at(candidate) > _captured_at(current)
+
+
+def _dt_from_folder(name: str) -> date | None:
+    if not name.startswith("dt="):
+        return None
+    try:
+        return date.fromisoformat(name[3:])
+    except ValueError:
+        return None
+
+
+def parse_part_name(name: str, suffix: str | None = None):
+    """Return (first_ts, last_ts|None, seq) for a part object."""
+    name = name or ""
+    if suffix is not None and not name.endswith(suffix):
+        return None
+    match = PART_RE.match(name)
+    if not match:
+        return None
+    first_ts = int(match.group(1))
+    last_ts = int(match.group(2)) if match.group(2) else None
+    seq = int(match.group(3))
+    return first_ts, last_ts, seq
+
+
+def _chat_could_have_window(chat: ChatEntry, window: Window) -> bool:
+    if chat.last_seen is not None and chat.last_seen < window.start_ms:
+        return False
+    if chat.first_seen is not None and chat.first_seen >= window.end_ms:
+        return False
+    return True
+
+
+def _select_parts(parts: list[dict], window: Window) -> list[dict]:
+    """Choose parts that can contain records in [start, end).
+
+    New names carry lastTs (exact interval). Legacy names only carry firstTs:
+    take the last legacy part starting at/before the window, plus every later
+    part. Parts are written in chronological order, so everything before that
+    anchor ends before the window starts (modulo repaired re-uploads).
+    """
+    ordered = sorted(parts, key=lambda part: (part["first_ts"], part["seq"]))
+    selected = [
+        part
+        for part in ordered
+        if part["last_ts"] is not None
+        and part["last_ts"] >= window.start_ms
+        and part["first_ts"] < window.end_ms
+    ]
+
+    anchor = -1
+    for index, part in enumerate(ordered):
+        if part["last_ts"] is None and part["first_ts"] <= window.start_ms:
+            anchor = index
+    for index, part in enumerate(ordered):
+        if part["last_ts"] is not None:
+            continue
+        if part["first_ts"] >= window.end_ms:
+            continue
+        if anchor < 0 or index >= anchor:
+            selected.append(part)
+
+    seen = set()
+    unique = []
+    for part in sorted(selected, key=lambda item: (item["first_ts"], item["seq"])):
+        if part["key"] not in seen:
+            seen.add(part["key"])
+            unique.append(part)
+    return unique
+
+
+def plan_tasks(
+    client, cfg: DropboxConfig, codec: ArchiveCodec, chats: list[ChatEntry], window: Window
+) -> list[FetchTask]:
+    low = datetime.fromtimestamp(window.start_ms / 1000, tz=cfg.tz).date() - timedelta(days=1)
+    high = datetime.fromtimestamp(window.end_ms / 1000, tz=cfg.tz).date() + timedelta(days=1)
+
+    tasks: list[FetchTask] = []
+    for chat in chats:
+        if not _chat_could_have_window(chat, window):
+            continue
+        base = f"{cfg.account_root}/chat={chat.hash}"
+        parts: list[dict] = []
+        for item in client.list_folder(base):
+            if item.get(".tag") != "folder":
+                continue
+            day = _dt_from_folder(item.get("name") or "")
+            if day is None or not (low <= day <= high):
+                continue
+            folder_path = item.get("path_display") or item.get("path_lower") or ""
+            for part in client.list_folder(folder_path):
+                if part.get(".tag") != "file":
+                    continue
+                parsed = parse_part_name(part.get("name") or "", codec.part_suffix)
+                if parsed is None:
+                    continue
+                path = part.get("path_display") or part.get("path_lower") or ""
+                if not path:
+                    continue
+                first_ts, last_ts, seq = parsed
+                parts.append({"key": path, "first_ts": first_ts, "last_ts": last_ts, "seq": seq})
+        for part in _select_parts(parts, window):
+            tasks.append(FetchTask(chat_hash=chat.hash, key=part["key"]))
+
+    if not tasks:
+        raise NoMatchesError("no archived messages found for the matched chats in this window")
+    if len(tasks) > cfg.max_parts:
+        raise GuardError(
+            f"{len(tasks)} part objects exceed --max-parts {cfg.max_parts}; narrow the window or selectors"
+        )
+    return tasks
+
+
+def _parse_part(text: str, key: str) -> list[dict]:
+    records: list[dict] = []
+    for line_number, line in enumerate(text.splitlines(), 1):
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            record = json.loads(line)
+        except json.JSONDecodeError as exc:
+            raise RemoteError(f"invalid NDJSON in {key} line {line_number}: {exc}") from exc
+        if isinstance(record, dict) and record.get("msgId"):
+            records.append(record)
+    return records
+
+
+def collect(
+    client, cfg: DropboxConfig, codec: ArchiveCodec, tasks: list[FetchTask], window: Window
+) -> list[dict]:
+    def fetch(task: FetchTask):
+        payload = client.download(task.key)
+        if payload is None:
+            return task.key, None
+        try:
+            return task.key, codec.decode_text(payload)
+        except (ValueError, InvalidTag) as exc:
+            raise RemoteError(f"cannot decrypt {task.key}: {exc}") from exc
+
+    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
+        payloads = list(pool.map(fetch, tasks))
+
+    merged: dict[str, dict] = {}
+    for key, text in payloads:
+        if text is None:
+            continue
+        for record in _parse_part(text, key):
+            ts = effective_ts(record)
+            if ts is None or not (window.start_ms <= ts < window.end_ms):
+                continue
+            msg_id = str(record["msgId"])
+            current = merged.get(msg_id)
+            if current is None or _prefer(record, current):
+                merged[msg_id] = record
+
+    return sorted(
+        merged.values(),
+        key=lambda record: (effective_ts(record) or 0, str(record.get("msgId") or "")),
+    )

waidrop_cli/mcp_cli.py

@@ -0,0 +1,293 @@
+"""Run the public waidrop CLI as a subprocess."""
+
+from __future__ import annotations
+
+import os
+import re
+import subprocess
+import sys
+from dataclasses import dataclass
+
+from .errors import CliError, UsageError
+from .mcp_config import ACCOUNT_RE, DEFAULT_ACCOUNT, McpSettings, load_env_file
+
+CREATE_NO_WINDOW = getattr(subprocess, "CREATE_NO_WINDOW", 0)
+CHAT_HEADER = "=== CHAT "
+CHAT_SEPARATOR = " | "
+FAILURE_DETAIL_LIMIT = 2000
+
+
+@dataclass(frozen=True)
+class CliResult:
+    returncode: int
+    stdout: str
+    stderr: str
+    timed_out: bool = False
+
+
+@dataclass(frozen=True)
+class LimitedText:
+    text: str
+    truncated: bool
+    included_chats: int
+    omitted_chats: int
+
+
+def base_command() -> list[str]:
+    return [sys.executable, "-m", "waidrop_cli"]
+
+
+def _base_env(settings: McpSettings) -> dict[str, str]:
+    env = dict(os.environ)
+    for key, value in load_env_file(settings.env_path).items():
+        env.setdefault(key, value)
+    if settings.account:
+        env.setdefault("WAIDROP_ACCOUNT", settings.account)
+    return env
+
+
+def build_env(settings: McpSettings) -> dict[str, str]:
+    env = _base_env(settings)
+    env.setdefault("PYTHONIOENCODING", "utf-8")
+    env.setdefault("PYTHONUTF8", "1")
+    return env
+
+
+def effective_account(settings: McpSettings) -> str:
+    account = (_base_env(settings).get("WAIDROP_ACCOUNT") or "").strip() or DEFAULT_ACCOUNT
+    if not ACCOUNT_RE.match(account):
+        raise UsageError("account alias may only contain letters, digits and . _ -")
+    return account
+
+
+def run_cli(args: list[str], settings: McpSettings) -> CliResult:
+    command = base_command() + args
+    try:
+        completed = subprocess.run(
+            command,
+            capture_output=True,
+            text=True,
+            encoding="utf-8",
+            errors="replace",
+            timeout=settings.cli_timeout_seconds,
+            env=build_env(settings),
+            check=False,
+            creationflags=CREATE_NO_WINDOW,
+        )
+    except subprocess.TimeoutExpired as exc:
+        return CliResult(
+            -1,
+            _as_text(exc.stdout),
+            _as_text(exc.stderr),
+            timed_out=True,
+        )
+    except CliError as exc:
+        return CliResult(exc.exit_code, "", str(exc))
+    except OSError as exc:
+        return CliResult(-1, "", str(exc))
+    return CliResult(completed.returncode, completed.stdout, completed.stderr)
+
+
+def chats_args(
+    *,
+    names: list[str] | None = None,
+    numbers: list[str] | None = None,
+    hashes: list[str] | None = None,
+    jids: list[str] | None = None,
+) -> list[str]:
+    args = ["chats", "--quiet"]
+    for value in names or []:
+        args += ["--name", value]
+    for value in numbers or []:
+        args += ["--number", value]
+    for value in hashes or []:
+        args += ["--chat", value]
+    for value in jids or []:
+        args += ["--jid", value]
+    return args
+
+
+def context_args(
+    *,
+    last: str | None,
+    start: str | None,
+    end: str | None,
+    context_before: str,
+    max_chars_per_chat: int,
+    max_parts: int,
+    tz: str | None,
+    names: list[str] | None = None,
+    numbers: list[str] | None = None,
+    hashes: list[str] | None = None,
+    jids: list[str] | None = None,
+) -> list[str]:
+    args = [
+        "context",
+        "--format",
+        "text",
+        "--quiet",
+        "--context-before",
+        context_before,
+        "--max-chars",
+        str(max_chars_per_chat),
+        "--max-parts",
+        str(max_parts),
+    ]
+    if last:
+        args += ["--last", last]
+    elif start:
+        args += ["--from", start]
+    if end:
+        args += ["--to", end]
+    if tz:
+        args += ["--tz", tz]
+    for value in names or []:
+        args += ["--name", value]
+    for value in numbers or []:
+        args += ["--number", value]
+    for value in hashes or []:
+        args += ["--chat", value]
+    for value in jids or []:
+        args += ["--jid", value]
+    return args
+
+
+def limit_cli_text(
+    text: str,
+    *,
+    max_chats: int,
+    max_total_chars: int,
+) -> LimitedText:
+    identities = _chat_identities(text)
+    if max_chats <= 0 and max_total_chars <= 0:
+        return LimitedText(text, False, len(identities), 0)
+
+    blocks = _split_blocks(text)
+    if not blocks:
+        if max_total_chars > 0 and len(text) > max_total_chars:
+            excerpt = _truncate_at_line(text, max_total_chars)
+            return LimitedText(
+                excerpt + "\n[waidrop-mcp: output truncated]\n",
+                True,
+                0,
+                0,
+            )
+        return LimitedText(text, False, 0, 0)
+
+    allowed = set(identities[:max_chats]) if max_chats > 0 else set(identities)
+    selected: list[str] = []
+    used = 0
+    stopped = False
+
+    for block in blocks:
+        identity = _block_identity(block)
+        if identity not in allowed:
+            continue
+        if max_total_chars > 0 and used + len(block) > max_total_chars:
+            if not selected:
+                excerpt = _truncate_at_line(block, max_total_chars)
+                selected.append(excerpt)
+                used = len(excerpt)
+            stopped = True
+            break
+        selected.append(block)
+        used += len(block)
+
+    included = {_block_identity(block) for block in selected if _block_identity(block)}
+    omitted = len([identity for identity in identities if identity not in included])
+    truncated = stopped or omitted > 0
+    result = "".join(selected)
+
+    if truncated:
+        if omitted:
+            notice = (
+                f"[waidrop-mcp: output truncated; omitted {omitted} chat(s). "
+                "Narrow the window or raise max_total_chars.]\n"
+            )
+        else:
+            notice = "[waidrop-mcp: output truncated; raise max_total_chars.]\n"
+        if not result.endswith("\n"):
+            result += "\n"
+        result += notice
+    return LimitedText(result, truncated, len(included), omitted)
+
+
+def describe_failure(result: CliResult, action: str, timeout_seconds: float) -> str:
+    detail = _clean_detail(result.stderr)
+    suffix = f" {detail}" if detail else ""
+
+    if result.timed_out:
+        return (
+            f"The waidrop CLI timed out after {timeout_seconds:g}s while trying to "
+            f"{action}. Narrow the window, select fewer chats, lower max_parts, "
+            f"or raise cli_timeout_seconds.{suffix}"
+        )
+    if result.returncode == 2:
+        return f"WAiDrop CLI configuration or usage error.{suffix}"
+    if result.returncode == 3:
+        return (
+            "No archived chats or messages matched. Try a longer window or a "
+            f"different chat selector.{suffix}"
+        )
+    if result.returncode == 4:
+        return (
+            "Too many archive parts would be fetched. Narrow the window or chat "
+            f"selection, or lower max_parts.{suffix}"
+        )
+    if result.returncode == 5:
+        return (
+            "Could not read the Dropbox archive. Check `waidrop auth --status` "
+            f"and the archive key.{suffix}"
+        )
+    if result.returncode == -1:
+        return f"Could not start the waidrop CLI.{suffix}"
+    return f"WAiDrop CLI failed with exit code {result.returncode}.{suffix}"
+
+
+def _split_blocks(text: str) -> list[str]:
+    blocks = re.split(r"(?=^=== CHAT )", text, flags=re.MULTILINE)
+    return [block for block in blocks if block.startswith(CHAT_HEADER)]
+
+
+def _block_identity(block: str) -> str:
+    first_line = block.splitlines()[0] if block else ""
+    header = first_line[len(CHAT_HEADER) :] if first_line.startswith(CHAT_HEADER) else first_line
+    identity = header.partition(CHAT_SEPARATOR)[0].strip()
+    return identity or header.strip()
+
+
+def _chat_identities(text: str) -> list[str]:
+    identities: list[str] = []
+    seen: set[str] = set()
+    for block in _split_blocks(text):
+        identity = _block_identity(block)
+        if identity not in seen:
+            seen.add(identity)
+            identities.append(identity)
+    return identities
+
+
+def _truncate_at_line(text: str, limit: int) -> str:
+    if limit <= 0:
+        return ""
+    excerpt = text[:limit]
+    if len(text) > limit:
+        newline = excerpt.rfind("\n")
+        if newline > 0:
+            excerpt = excerpt[: newline + 1]
+    return excerpt
+
+
+def _clean_detail(value: str) -> str:
+    detail = (value or "").strip()
+    if len(detail) > FAILURE_DETAIL_LIMIT:
+        detail = detail[:FAILURE_DETAIL_LIMIT].rstrip() + "..."
+    return detail
+
+
+def _as_text(value) -> str:
+    if value is None:
+        return ""
+    if isinstance(value, bytes):
+        return value.decode("utf-8", errors="replace")
+    return str(value)

waidrop_cli/mcp_config.py

@@ -0,0 +1,191 @@
+"""MCP server settings, environment file loading and precedence."""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from types import SimpleNamespace
+
+from .auth import resolve_credentials_path
+from .config import parse_context_before, parse_duration, parse_tz
+from .env_file import load_env_file  # noqa: F401 (re-exported for existing callers)
+from .errors import CliError
+
+DEFAULT_LAST = "6h"
+DEFAULT_CONTEXT_BEFORE = "30m"
+DEFAULT_MAX_CHARS_PER_CHAT = 0
+DEFAULT_MAX_TOTAL_CHARS = 0
+DEFAULT_MAX_CHATS = 0
+DEFAULT_MAX_PARTS = 1000
+DEFAULT_CLI_TIMEOUT_SECONDS = 120.0
+DEFAULT_ACCOUNT = "primary"
+
+ACCOUNT_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
+
+CONFIG_KEYS = {
+    "default_last",
+    "default_context_before",
+    "max_chars_per_chat",
+    "max_total_chars",
+    "max_chats",
+    "max_parts",
+    "cli_timeout_seconds",
+    "account",
+    "credentials_file",
+    "env_file",
+    "tz",
+}
+
+
+@dataclass(frozen=True)
+class McpSettings:
+    default_last: str = DEFAULT_LAST
+    default_context_before: str = DEFAULT_CONTEXT_BEFORE
+    max_chars_per_chat: int = DEFAULT_MAX_CHARS_PER_CHAT
+    max_total_chars: int = DEFAULT_MAX_TOTAL_CHARS
+    max_chats: int = DEFAULT_MAX_CHATS
+    max_parts: int = DEFAULT_MAX_PARTS
+    cli_timeout_seconds: float = DEFAULT_CLI_TIMEOUT_SECONDS
+    account: str = DEFAULT_ACCOUNT
+    credentials_file: str | None = None
+    env_file: str | None = None
+    tz: str | None = None
+
+    @property
+    def credentials_path(self) -> Path:
+        return resolve_credentials_path(SimpleNamespace(credentials_file=self.credentials_file))
+
+    @property
+    def env_path(self) -> Path:
+        if self.env_file:
+            return Path(self.env_file).expanduser()
+        return self.credentials_path.parent / "env"
+
+
+def default_config_path() -> Path:
+    explicit = _env_text("WAIDROP_MCP_CONFIG")
+    if explicit is not None:
+        return Path(explicit).expanduser()
+    credentials = resolve_credentials_path(
+        SimpleNamespace(credentials_file=_env_text("WAIDROP_CREDENTIALS"))
+    )
+    return credentials.parent / "mcp.json"
+
+
+def load_settings() -> McpSettings:
+    config_path = default_config_path()
+    document = _read_config(config_path)
+    values = {key: document[key] for key in CONFIG_KEYS if key in document}
+    values = _apply_env(values)
+    settings = McpSettings(**values)
+    _validate(settings)
+    return settings
+
+
+def _read_config(path: Path) -> dict:
+    try:
+        text = path.read_text(encoding="utf-8")
+    except FileNotFoundError:
+        return {}
+    except OSError as exc:
+        raise ValueError(f"could not read {path}: {exc}") from exc
+
+    try:
+        document = json.loads(text)
+    except json.JSONDecodeError as exc:
+        raise ValueError(f"{path} is not valid JSON: {exc}") from exc
+
+    if not isinstance(document, dict):
+        raise ValueError(f"{path} must contain a JSON object")
+    if document.get("version") != 1:
+        raise ValueError(f'{path} must contain "version": 1')
+    return document
+
+
+def _apply_env(values: dict) -> dict:
+    merged = dict(values)
+    overrides = {
+        "default_last": "WAIDROP_MCP_DEFAULT_LAST",
+        "default_context_before": "WAIDROP_MCP_DEFAULT_CONTEXT_BEFORE",
+        "max_chars_per_chat": "WAIDROP_MCP_MAX_CHARS",
+        "max_total_chars": "WAIDROP_MCP_MAX_TOTAL_CHARS",
+        "max_chats": "WAIDROP_MCP_MAX_CHATS",
+        "max_parts": "WAIDROP_MCP_MAX_PARTS",
+        "cli_timeout_seconds": "WAIDROP_MCP_TIMEOUT",
+        "env_file": "WAIDROP_MCP_ENV_FILE",
+    }
+    for field, env_name in overrides.items():
+        raw = _env_text(env_name)
+        if raw is None:
+            continue
+        if field in {"max_chars_per_chat", "max_total_chars", "max_chats", "max_parts"}:
+            merged[field] = _parse_int(raw, env_name)
+        elif field == "cli_timeout_seconds":
+            merged[field] = _parse_float(raw, env_name)
+        else:
+            merged[field] = raw
+    return merged
+
+
+def _validate(settings: McpSettings) -> None:
+    for field in ("default_last", "default_context_before", "account"):
+        if not isinstance(getattr(settings, field), str):
+            raise ValueError(f"{field} must be a string")
+    for field in ("max_chars_per_chat", "max_total_chars", "max_chats", "max_parts"):
+        value = getattr(settings, field)
+        if not isinstance(value, int) or isinstance(value, bool):
+            raise ValueError(f"{field} must be an integer")
+    if not isinstance(settings.cli_timeout_seconds, (int, float)) or isinstance(
+        settings.cli_timeout_seconds, bool
+    ):
+        raise ValueError("cli_timeout_seconds must be a number")
+    if settings.max_chars_per_chat < 0:
+        raise ValueError("max_chars_per_chat must be zero or positive")
+    if settings.max_total_chars < 0:
+        raise ValueError("max_total_chars must be zero or positive")
+    if settings.max_chats < 0:
+        raise ValueError("max_chats must be zero or positive")
+    if not 1 <= settings.max_parts <= 1000:
+        raise ValueError("max_parts must be between 1 and 1000")
+    if not 1 <= settings.cli_timeout_seconds <= 600:
+        raise ValueError("cli_timeout_seconds must be between 1 and 600")
+    if not ACCOUNT_RE.match(settings.account or ""):
+        raise ValueError("account alias may only contain letters, digits and . _ -")
+    try:
+        parse_duration(settings.default_last)
+    except CliError as exc:
+        raise ValueError(f"invalid default_last: {exc}") from exc
+    try:
+        parse_context_before(settings.default_context_before)
+    except CliError as exc:
+        raise ValueError(f"invalid default_context_before: {exc}") from exc
+    if settings.tz:
+        try:
+            parse_tz(settings.tz)
+        except CliError as exc:
+            raise ValueError(f"invalid tz: {exc}") from exc
+
+
+def _env_text(name: str) -> str | None:
+    value = os.environ.get(name)
+    if value is None:
+        return None
+    value = value.strip()
+    return value or None
+
+
+def _parse_int(raw: str, name: str) -> int:
+    try:
+        return int(raw)
+    except ValueError as exc:
+        raise ValueError(f"{name} must be an integer") from exc
+
+
+def _parse_float(raw: str, name: str) -> float:
+    try:
+        return float(raw)
+    except ValueError as exc:
+        raise ValueError(f"{name} must be a number") from exc