refactor: deduplicate shared helpers and unify QR login state machine

examples/media_download.py

@@ -21,33 +21,21 @@
 from pathlib import Path
 
 from weilink import MessageType, WeiLink
+from weilink._helpers import media_filename
 
 logging.basicConfig(
     level=getattr(logging, os.environ.get("LOGLEVEL", "INFO").upper(), logging.INFO),
     format="%(asctime)s %(levelname)s %(message)s",
 )
 logger = logging.getLogger(__name__)
 
-# File extensions by message type
-_EXT: dict[MessageType, str] = {
-    MessageType.IMAGE: ".jpg",
-    MessageType.VOICE: ".silk",
-    MessageType.VIDEO: ".mp4",
-}
 
-
-def _save_path(save_dir: Path, msg_type: MessageType, msg: object) -> Path:
+def _save_path(save_dir: Path, msg: object) -> Path:
     """Build a unique save path for a media message."""
     from weilink.models import Message
 
     assert isinstance(msg, Message)
-    mid = msg.message_id or 0
-
-    # Use original file name for FILE type
-    if msg_type == MessageType.FILE and msg.file:
-        name = msg.file.file_name or f"{mid}.bin"
-    else:
-        name = f"{mid}{_EXT.get(msg_type, '.bin')}"
+    name = media_filename(msg)
 
     path = save_dir / name
     # Avoid overwriting — append counter if file exists
@@ -100,7 +88,7 @@ def main() -> None:
                     wl.send(user, f"[Download failed: {e}]")
                     continue
 
-                path = _save_path(save_dir, msg.msg_type, msg)
+                path = _save_path(save_dir, msg)
                 path.write_bytes(data)
                 logger.info(
                     "Saved %s (%d bytes) -> %s", msg.msg_type.name, len(data), path

src/weilink/_helpers.py

@@ -0,0 +1,140 @@
+"""Shared constants and helpers used across CLI, server, and admin."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from weilink.models import BotInfo, Message, MessageDirection, MessageType
+
+
+# -- Media constants ---------------------------------------------------------
+
+MEDIA_EXT_MAP: dict[MessageType, str] = {
+    MessageType.IMAGE: ".jpg",
+    MessageType.VOICE: ".amr",
+    MessageType.VIDEO: ".mp4",
+}
+
+MEDIA_MIME_MAP: dict[MessageType, str] = {
+    MessageType.IMAGE: "image/jpeg",
+    MessageType.VOICE: "audio/amr",
+    MessageType.VIDEO: "video/mp4",
+}
+
+
+# -- QR login ----------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class QRResult:
+    """Result of interpreting a ``poll_qr_status`` response.
+
+    Attributes:
+        status: Normalized status string — one of
+            ``"confirmed"``, ``"scanned"``, ``"expired"``, or ``"waiting"``.
+        bot_info: Populated only when *status* is ``"confirmed"``.
+    """
+
+    status: str
+    bot_info: BotInfo | None = None
+
+
+def process_qr_status(status_resp: dict[str, Any]) -> QRResult:
+    """Interpret a raw ``poll_qr_status`` response.
+
+    Args:
+        status_resp: Dict returned by ``_protocol.poll_qr_status()``.
+
+    Returns:
+        A :class:`QRResult` with normalized status and, on confirmation,
+        the extracted :class:`BotInfo`.
+    """
+    from weilink._protocol import BASE_URL
+
+    status = status_resp.get("status", "")
+
+    if status == "confirmed":
+        bot_info = BotInfo(
+            bot_id=status_resp.get("ilink_bot_id", ""),
+            base_url=status_resp.get("baseurl", BASE_URL),
+            token=status_resp.get("bot_token", ""),
+            user_id=status_resp.get("ilink_user_id", ""),
+        )
+        return QRResult(status="confirmed", bot_info=bot_info)
+
+    if status == "scaned":
+        return QRResult(status="scanned")
+
+    if status == "expired":
+        return QRResult(status="expired")
+
+    return QRResult(status="waiting")
+
+
+# -- Media helpers -----------------------------------------------------------
+
+
+def media_filename(msg: Message) -> str:
+    """Derive a filename for a media message.
+
+    Uses the original ``file_name`` when available, otherwise falls back
+    to ``{message_id}{ext}`` using :data:`MEDIA_EXT_MAP`.
+    """
+    if msg.file and msg.file.file_name:
+        return msg.file.file_name
+    ext = MEDIA_EXT_MAP.get(msg.msg_type, ".bin")
+    return f"{msg.message_id}{ext}"
+
+
+# -- Parsing helpers ---------------------------------------------------------
+
+
+def parse_direction(s: str) -> int | None:
+    """Parse a direction string to its integer value.
+
+    Returns:
+        ``MessageDirection.USER`` (1) for ``"received"``,
+        ``MessageDirection.BOT`` (2) for ``"sent"``,
+        or ``None`` for unrecognized input.
+    """
+    d = s.lower()
+    if d == "received":
+        return MessageDirection.USER
+    if d == "sent":
+        return MessageDirection.BOT
+    return None
+
+
+def parse_message_type(s: str) -> int | None:
+    """Parse a message type name to its integer value.
+
+    Returns:
+        The ``MessageType`` integer (e.g. 2 for ``"IMAGE"``),
+        or ``None`` for unrecognized input.
+    """
+    try:
+        return MessageType[s.upper()].value
+    except KeyError:
+        return None
+
+
+def parse_time(s: str) -> int | None:
+    """Parse an ISO 8601 string or unix milliseconds to *int* milliseconds.
+
+    Returns:
+        Millisecond timestamp, or ``None`` on failure.
+    """
+    if not s:
+        return None
+    try:
+        return int(s)
+    except ValueError:
+        pass
+    try:
+        from datetime import datetime
+
+        dt = datetime.fromisoformat(s.replace("Z", "+00:00"))
+        return int(dt.timestamp() * 1000)
+    except (ValueError, TypeError):
+        return None

src/weilink/admin/handlers.py

@@ -12,7 +12,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar
 
 from weilink import __version__
-from weilink.models import BotInfo
+from weilink._helpers import MEDIA_MIME_MAP, media_filename, process_qr_status
 
 from .static import ADMIN_HTML, load_locale
 
@@ -248,13 +248,10 @@ def _handle_poll_login(self, query: dict[str, list[str]]) -> None:
             self._send_json({"status": "waiting"})
             return
 
-        status = status_resp.get("status", "")
+        qr = process_qr_status(status_resp)
 
-        if status == "confirmed":
-            bot_token = status_resp.get("bot_token", "")
-            base_url = status_resp.get("baseurl", proto.BASE_URL)
-            bot_id = status_resp.get("ilink_bot_id", "")
-            user_id = status_resp.get("ilink_user_id", "")
+        if qr.status == "confirmed":
+            assert qr.bot_info is not None
             name = pending["name"]
 
             with self._lock:
@@ -265,28 +262,23 @@ def _handle_poll_login(self, query: dict[str, list[str]]) -> None:
                     token_path = wl._base_path / name / "token.json"
                     session = wl._create_session(name, token_path)
 
-                session.bot_info = BotInfo(
-                    bot_id=bot_id,
-                    base_url=base_url,
-                    token=bot_token,
-                    user_id=user_id,
-                )
+                session.bot_info = qr.bot_info
                 session.cursor = ""
                 wl._save_session_state(session)
 
             del self._pending_logins[qrcode]
             self._send_json(
                 {
                     "status": "confirmed",
-                    "bot_id": bot_id,
+                    "bot_id": qr.bot_info.bot_id,
                     "session_name": name,
                 }
             )
 
-        elif status == "scaned":
+        elif qr.status == "scanned":
             self._send_json({"status": "scaned"})
 
-        elif status == "expired":
+        elif qr.status == "expired":
             del self._pending_logins[qrcode]
             self._send_json({"status": "expired"})
 
@@ -427,17 +419,6 @@ def _handle_get_messages(self, query: dict[str, list[str]]) -> None:
                 m["message_id"] = str(m["message_id"])
         self._send_json({"messages": messages, "total": total})
 
-    _MIME_MAP: ClassVar[dict[str, str]] = {
-        "IMAGE": "image/jpeg",
-        "VOICE": "audio/amr",
-        "VIDEO": "video/mp4",
-    }
-    _EXT_MAP: ClassVar[dict[str, str]] = {
-        "IMAGE": ".jpg",
-        "VOICE": ".amr",
-        "VIDEO": ".mp4",
-    }
-
     def _handle_download_media(self, message_id_str: str) -> None:
         """Download media from a stored message and serve the bytes."""
         store = self.weilink._message_store
@@ -463,13 +444,8 @@ def _handle_download_media(self, message_id_str: str) -> None:
             return
 
         # Derive filename and MIME type
-        mt = msg.msg_type.name
-        if msg.file and msg.file.file_name:
-            filename = msg.file.file_name
-        else:
-            ext = self._EXT_MAP.get(mt, ".bin")
-            filename = f"{msg.message_id}{ext}"
-        content_type = self._MIME_MAP.get(mt, "application/octet-stream")
+        filename = media_filename(msg)
+        content_type = MEDIA_MIME_MAP.get(msg.msg_type, "application/octet-stream")
 
         self.send_response(200)
         self.send_header("Content-Type", content_type)

src/weilink/cli.py

@@ -264,17 +264,9 @@ def _run_download(args: argparse.Namespace) -> None:
     out_dir.mkdir(parents=True, exist_ok=True)
 
     # Derive filename
-    from weilink.models import MessageType
-
-    ext_map = {
-        MessageType.IMAGE: ".jpg",
-        MessageType.VOICE: ".amr",
-        MessageType.VIDEO: ".mp4",
-    }
-    if msg.file and msg.file.file_name:
-        name = msg.file.file_name
-    else:
-        name = f"{msg.message_id}{ext_map.get(msg.msg_type, '.bin')}"
+    from weilink._helpers import media_filename
+
+    name = media_filename(msg)
 
     out_path = out_dir / name
     out_path.write_bytes(data)
@@ -299,40 +291,34 @@ def _run_history(args: argparse.Namespace) -> None:
         wl.close()
         sys.exit(1)
 
-    from weilink.models import MessageType
+    from weilink._helpers import parse_direction, parse_message_type, parse_time
 
     kwargs: dict[str, Any] = {}
     if args.user:
         kwargs["user_id"] = args.user
     if args.bot:
         kwargs["bot_id"] = args.bot
     if args.type:
-        try:
-            kwargs["msg_type"] = MessageType[args.type.upper()].value
-        except KeyError:
+        mt = parse_message_type(args.type)
+        if mt is None:
             err = f"Unknown message type: {args.type}"
             if _json_flag(args):
                 print(json.dumps({"error": err}))
             else:
                 print(f"Error: {err}", file=sys.stderr)
             wl.close()
             sys.exit(1)
+        kwargs["msg_type"] = mt
     if args.direction:
-        d = args.direction.lower()
-        if d == "received":
-            kwargs["direction"] = 1
-        elif d == "sent":
-            kwargs["direction"] = 2
+        d = parse_direction(args.direction)
+        if d is not None:
+            kwargs["direction"] = d
     if args.since:
-        from weilink.server.app import _parse_time
-
-        ts = _parse_time(args.since)
+        ts = parse_time(args.since)
         if ts is not None:
             kwargs["since_ms"] = ts
     if args.until:
-        from weilink.server.app import _parse_time
-
-        ts = _parse_time(args.until)
+        ts = parse_time(args.until)
         if ts is not None:
             kwargs["until_ms"] = ts
     if args.text: