Add multi-user MAX authentication and tenant isolation

pymax-worker/src/server.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import base64
+import contextvars
 import json
 import mimetypes
 import os
@@ -24,9 +25,8 @@ from pymax.types import ContactInfo
 
 PORT = int(os.environ.get("PORT", "3002"))
 PHONE_NUMBER = os.environ.get("PYMAX_PHONE_NUMBER") or os.environ.get("QMAX_MAX_PHONE_NUMBER") or ""
-SESSION_DIR = pathlib.Path(os.environ.get("PYMAX_SESSION_DIR", "/data/pymax-session"))
+SESSION_ROOT = pathlib.Path(os.environ.get("PYMAX_SESSION_DIR", "/data/pymax-session"))
 SESSION_NAME = os.environ.get("PYMAX_SESSION_NAME", "session.db")
-PHONE_CONTACT_IDS_FILE = SESSION_DIR / "phone-contact-ids.json"
 CHAT_FETCH_LIMIT = int(os.environ.get("PYMAX_CHAT_FETCH_LIMIT", "350"))
 HISTORY_LIMIT = int(os.environ.get("PYMAX_HISTORY_LIMIT", "80"))
 SEND_ROOTS = [pathlib.Path(p) for p in os.environ.get("PYMAX_SEND_ROOTS", "/qmax-data:/tmp").split(":") if p]
@@ -80,16 +80,16 @@ def is_expired_session_error(error: str | None) -> bool:
     )
 
 
-def archive_session_dir() -> str | None:
-    if not SESSION_DIR.exists():
+def archive_session_dir(session_dir: pathlib.Path) -> str | None:
+    if not session_dir.exists():
         return None
     stamp = datetime.now(UTC).strftime("%Y%m%d-%H%M%S")
-    target = SESSION_DIR.with_name(f"{SESSION_DIR.name}.expired-{stamp}")
+    target = session_dir.with_name(f"{session_dir.name}.expired-{stamp}")
     suffix = 1
     while target.exists():
-        target = SESSION_DIR.with_name(f"{SESSION_DIR.name}.expired-{stamp}-{suffix}")
+        target = session_dir.with_name(f"{session_dir.name}.expired-{stamp}-{suffix}")
         suffix += 1
-    shutil.move(str(SESSION_DIR), str(target))
+    shutil.move(str(session_dir), str(target))
     return str(target)
 
 
@@ -372,7 +372,9 @@ class DeferredPasswordProvider:
 
 
 class PyMaxRuntime:
-    def __init__(self) -> None:
+    def __init__(self, session_dir: pathlib.Path, default_phone: str = "") -> None:
+        self.session_dir = session_dir
+        self.phone = default_phone
         self.client: Client | None = None
         self.task: asyncio.Task[None] | None = None
         self.ready = asyncio.Event()
@@ -409,15 +411,22 @@ class PyMaxRuntime:
                 return
             self.ready.clear()
             self.last_error = None
-            use_phone = normalize_phone(phone or PHONE_NUMBER)
+            stored_phone_file = self.session_dir / "phone.txt"
+            stored_phone = ""
+            with suppress(OSError):
+                stored_phone = stored_phone_file.read_text(encoding="utf-8").strip()
+            use_phone = normalize_phone(phone or self.phone or stored_phone or PHONE_NUMBER)
             if not use_phone:
                 self.last_error = "PYMAX_PHONE_NUMBER is not configured."
                 return
-            SESSION_DIR.mkdir(parents=True, exist_ok=True)
-            SESSION_DIR.chmod(0o700)
+            self.phone = use_phone
+            self.session_dir.mkdir(parents=True, exist_ok=True)
+            self.session_dir.chmod(0o700)
+            stored_phone_file.write_text(use_phone, encoding="utf-8")
+            stored_phone_file.chmod(0o600)
             client = Client(
                 phone=use_phone,
-                work_dir=str(SESSION_DIR),
+                work_dir=str(self.session_dir),
                 session_name=SESSION_NAME,
                 sms_code_provider=self.code_provider,
                 password_provider=self.password_provider,
@@ -474,7 +483,7 @@ class PyMaxRuntime:
         async with self.lock:
             await self._stop_locked()
             self.last_error = None
-            archive_session_dir()
+            archive_session_dir(self.session_dir)
 
     async def _stop_locked(self) -> None:
         client = self.client
@@ -493,7 +502,35 @@ class PyMaxRuntime:
                 await task
 
 
-runtime = PyMaxRuntime()
+class RuntimeRegistry:
+    def __init__(self) -> None:
+        self._runtimes: dict[str, PyMaxRuntime] = {}
+
+    def get(self, account_id: str | None) -> PyMaxRuntime:
+        key = (account_id or "legacy").strip().lower()
+        if key != "legacy" and (len(key) != 32 or any(ch not in "0123456789abcdef" for ch in key)):
+            raise ValueError("X-QMax-Account-Id must be a 32-character hexadecimal UUID.")
+        runtime = self._runtimes.get(key)
+        if runtime is None:
+            session_dir = SESSION_ROOT if key == "legacy" else SESSION_ROOT / "accounts" / key
+            runtime = PyMaxRuntime(session_dir)
+            self._runtimes[key] = runtime
+        return runtime
+
+    async def stop(self) -> None:
+        await asyncio.gather(*(runtime.stop() for runtime in self._runtimes.values()), return_exceptions=True)
+
+
+registry = RuntimeRegistry()
+runtime_context: contextvars.ContextVar[PyMaxRuntime] = contextvars.ContextVar("qmax_runtime")
+
+
+class RuntimeProxy:
+    def __getattr__(self, name: str) -> Any:
+        return getattr(runtime_context.get(), name)
+
+
+runtime = RuntimeProxy()
 
 
 def json_response(data: Any, status: int = 200) -> web.Response:
@@ -890,7 +927,8 @@ async def status(_request: web.Request) -> web.Response:
 async def login_start(request: web.Request) -> web.Response:
     data = await read_json(request)
     phone = str(data.get("phoneNumber") or PHONE_NUMBER).strip()
-    if not runtime.is_authorized and runtime.login_stage() not in {"Code", "Password"}:
+    force = data.get("force") is True
+    if force or (not runtime.is_authorized and runtime.login_stage() not in {"Code", "Password"}):
         await runtime.reset_session_for_login()
     await runtime.ensure_started(phone)
     return json_response(status_payload())
@@ -969,18 +1007,20 @@ async def contacts(_request: web.Request) -> web.Response:
 
 
 def load_phone_contact_ids() -> set[int]:
+    path = runtime.session_dir / "phone-contact-ids.json"
     try:
-        values = json.loads(PHONE_CONTACT_IDS_FILE.read_text(encoding="utf-8"))
+        values = json.loads(path.read_text(encoding="utf-8"))
     except (OSError, ValueError, TypeError):
         return set()
     return {value for item in values if (value := coerce_int(item)) is not None and value > 0}
 
 
 def save_phone_contact_ids(contact_ids: set[int]) -> None:
-    SESSION_DIR.mkdir(parents=True, exist_ok=True)
-    part_path = PHONE_CONTACT_IDS_FILE.with_suffix(".json.part")
+    runtime.session_dir.mkdir(parents=True, exist_ok=True)
+    path = runtime.session_dir / "phone-contact-ids.json"
+    part_path = path.with_suffix(".json.part")
     part_path.write_text(json.dumps(sorted(contact_ids)), encoding="utf-8")
-    part_path.replace(PHONE_CONTACT_IDS_FILE)
+    part_path.replace(path)
 
 
 def append_saved_contacts(client: Client, contacts: list[Any]) -> None:
@@ -1223,8 +1263,21 @@ async def media_fetch(request: web.Request) -> web.StreamResponse:
             )
 
 
+@web.middleware
+async def account_runtime_middleware(request: web.Request, handler: Callable[[web.Request], Awaitable[web.StreamResponse]]) -> web.StreamResponse:
+    try:
+        selected = registry.get(request.headers.get("X-QMax-Account-Id"))
+    except ValueError as exc:
+        return json_response({"success": False, "error": str(exc)}, status=400)
+    token = runtime_context.set(selected)
+    try:
+        return await handler(request)
+    finally:
+        runtime_context.reset(token)
+
+
 def create_app() -> web.Application:
-    app = web.Application(client_max_size=1024 * 1024)
+    app = web.Application(client_max_size=1024 * 1024, middlewares=[account_runtime_middleware])
     app.router.add_get("/health", health)
     app.router.add_get("/status", status)
     app.router.add_post("/login/start", login_start)
@@ -1253,7 +1306,7 @@ def create_app() -> web.Application:
 
 
 async def on_cleanup(_app: web.Application) -> None:
-    await runtime.stop()
+    await registry.stop()
 
 
 if __name__ == "__main__":