#!/usr/bin/env python3 """Mistral API key management service for Yunwo parsing workers. The service intentionally uses only Python's standard library so it can be deployed next to the production parser without dependency churn. """ from __future__ import annotations import argparse import csv import json import os import re import sqlite3 import sys import threading import traceback import urllib.error import urllib.parse import urllib.request from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta, timezone from http import HTTPStatus from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from typing import Any, Dict, Iterable, List, Optional, Tuple ROOT_DIR = Path(__file__).resolve().parent STATIC_DIR = ROOT_DIR / "static" DEFAULT_DB_PATH = ROOT_DIR / "data" / "keys.sqlite" DEFAULT_MAILBOX_API_BASE = "https://mail-api.sdramago.top" DEFAULT_MISTRAL_API_BASE = "https://api.mistral.ai/v1" DEFAULT_MAILBOX_API_USERNAME = "" DEFAULT_MAILBOX_API_PASSWORD = "" STATUS_ACTIVE = "active" STATUS_UNKNOWN = "unknown" STATUS_QUOTA_EXHAUSTED = "quota_exhausted" STATUS_DISABLED = "disabled" STATUS_ERROR = "error" # Only keys that have passed a check or a successful runtime request may be # handed to parsing workers. Unknown keys are kept in the manager for manual # checking, but must not be advertised as usable credentials. ACTIVE_STATUSES = {STATUS_ACTIVE} class MailboxUnavailable(RuntimeError): pass def utc_now() -> str: return datetime.now(timezone.utc).isoformat(timespec="seconds") def current_month_last_day() -> str: today = datetime.now(timezone.utc).date() if today.month == 12: next_month = today.replace(year=today.year + 1, month=1, day=1) else: next_month = today.replace(month=today.month + 1, day=1) candidate = next_month - timedelta(days=1) if candidate > today: return candidate.isoformat() if next_month.month == 12: month_after = next_month.replace(year=next_month.year + 1, month=1, day=1) else: month_after = next_month.replace(month=next_month.month + 1, day=1) return (month_after - timedelta(days=1)).isoformat() def json_dumps(data: Any) -> bytes: return json.dumps(data, ensure_ascii=False, separators=(",", ":")).encode("utf-8") def read_env_bool(name: str, default: bool = False) -> bool: value = os.getenv(name) if value is None: return default return value.strip().lower() in {"1", "true", "yes", "y", "on"} def read_env_float(name: str, default: float) -> float: value = os.getenv(name) if value is None or value.strip() == "": return default try: return float(value) except ValueError: return default def split_api_keys(raw: str) -> List[str]: candidates = [ item.strip() for item in raw.replace("\n", ",").replace(";", ",").replace(" ", ",").split(",") ] return [item for item in candidates if item] def mask_key(api_key: str) -> str: clean = (api_key or "").strip() if len(clean) <= 10: return "*" * len(clean) return f"{clean[:6]}...{clean[-4:]}" def parse_reset_date(text: str) -> Optional[str]: match = re.search(r"reset\s+on\s+([A-Za-z]+)\s+(\d{1,2}),\s+(\d{4})", text or "", re.I) if not match: return None month, day, year = match.groups() for fmt in ("%B %d %Y", "%b %d %Y"): try: parsed = datetime.strptime(f"{month} {day} {year}", fmt) return parsed.date().isoformat() except ValueError: continue return None def parse_bulk_lines(text: str) -> List[Dict[str, Any]]: rows: List[Dict[str, Any]] = [] for raw_line in (text or "").splitlines(): line = raw_line.strip() if not line or line.startswith("#"): continue if "|" in line and "," not in line: parts = [part.strip() for part in line.split("|")] elif "," in line: parts = next(csv.reader([line])) parts = [part.strip() for part in parts] else: parts = line.split() if not parts: continue api_key = parts[0].strip() email = "" label_parts: List[str] = [] for part in parts[1:]: if "@" in part and not email: email = part.strip() else: label_parts.append(part.strip()) rows.append( { "api_key": api_key, "email": email, "label": " ".join(part for part in label_parts if part), "enabled": True, } ) return rows class KeyStore: def __init__(self, db_path: Path) -> None: self.db_path = db_path self.db_path.parent.mkdir(parents=True, exist_ok=True) self.init_db() def connect(self) -> sqlite3.Connection: conn = sqlite3.connect(str(self.db_path), timeout=30) conn.row_factory = sqlite3.Row conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA foreign_keys=ON") return conn def init_db(self) -> None: with self.connect() as conn: conn.execute( """ CREATE TABLE IF NOT EXISTS api_keys ( id INTEGER PRIMARY KEY AUTOINCREMENT, api_key TEXT NOT NULL UNIQUE, email TEXT NOT NULL DEFAULT '', label TEXT NOT NULL DEFAULT '', enabled INTEGER NOT NULL DEFAULT 1, status TEXT NOT NULL DEFAULT 'active', last_checked_at TEXT, last_success_at TEXT, last_failure_at TEXT, last_status_code INTEGER, last_error TEXT, quota_reset_at TEXT, quota_email_id TEXT, quota_email_subject TEXT, quota_email_seen_at TEXT, notes TEXT NOT NULL DEFAULT '', created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """ ) conn.execute( """ CREATE TABLE IF NOT EXISTS events ( id INTEGER PRIMARY KEY AUTOINCREMENT, key_id INTEGER NOT NULL, event_type TEXT NOT NULL, message TEXT NOT NULL DEFAULT '', metadata_json TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL ) """ ) def row_to_public(self, row: sqlite3.Row, *, include_secret: bool = False) -> Dict[str, Any]: item = dict(row) api_key = item.pop("api_key") item["masked_key"] = mask_key(api_key) item["key_prefix"] = api_key[:6] item["key_suffix"] = api_key[-4:] item["enabled"] = bool(item.get("enabled")) if include_secret: item["api_key"] = api_key return item def add_or_update_key(self, payload: Dict[str, Any]) -> Dict[str, Any]: api_key = str(payload.get("api_key") or "").strip() if not api_key: raise ValueError("api_key is required") now = utc_now() email = str(payload.get("email") or "").strip() label = str(payload.get("label") or "").strip() notes = str(payload.get("notes") or "").strip() enabled = 1 if payload.get("enabled", True) else 0 status = STATUS_ACTIVE if enabled else STATUS_DISABLED with self.connect() as conn: conn.execute( """ INSERT INTO api_keys (api_key, email, label, enabled, status, notes, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(api_key) DO UPDATE SET email=excluded.email, label=CASE WHEN excluded.label != '' THEN excluded.label ELSE api_keys.label END, enabled=excluded.enabled, status=CASE WHEN excluded.enabled = 0 THEN 'disabled' WHEN api_keys.status = 'disabled' AND excluded.enabled = 1 THEN 'active' ELSE api_keys.status END, notes=CASE WHEN excluded.notes != '' THEN excluded.notes ELSE api_keys.notes END, updated_at=excluded.updated_at """, (api_key, email, label, enabled, status, notes, now, now), ) row = conn.execute("SELECT * FROM api_keys WHERE api_key=?", (api_key,)).fetchone() self.add_event(conn, int(row["id"]), "upsert", "key saved", {}) return self.row_to_public(row) def list_keys(self, params: Dict[str, List[str]]) -> List[Dict[str, Any]]: where: List[str] = [] values: List[Any] = [] status = (params.get("status") or [""])[0].strip() enabled = (params.get("enabled") or [""])[0].strip().lower() q = (params.get("q") or [""])[0].strip() if status: where.append("status = ?") values.append(status) if enabled in {"0", "1", "true", "false"}: where.append("enabled = ?") values.append(1 if enabled in {"1", "true"} else 0) if q: where.append("(email LIKE ? OR label LIKE ? OR notes LIKE ? OR api_key LIKE ?)") like = f"%{q}%" values.extend([like, like, like, like]) sql = "SELECT * FROM api_keys" if where: sql += " WHERE " + " AND ".join(where) sql += " ORDER BY enabled DESC, status ASC, id DESC" with self.connect() as conn: return [self.row_to_public(row) for row in conn.execute(sql, values).fetchall()] def get_key(self, key_id: int, *, include_secret: bool = False) -> Optional[Dict[str, Any]]: with self.connect() as conn: row = conn.execute("SELECT * FROM api_keys WHERE id=?", (key_id,)).fetchone() if row is None: return None return self.row_to_public(row, include_secret=include_secret) def get_key_row(self, conn: sqlite3.Connection, key_id: int) -> Optional[sqlite3.Row]: return conn.execute("SELECT * FROM api_keys WHERE id=?", (key_id,)).fetchone() def patch_key(self, key_id: int, payload: Dict[str, Any]) -> Dict[str, Any]: allowed = {"email", "label", "enabled", "status", "notes"} updates: List[str] = [] values: List[Any] = [] current = self.get_key(key_id) if current is None: raise KeyError("key not found") next_status = payload.get("status") if "enabled" in payload and not payload.get("enabled"): next_status = STATUS_DISABLED elif "enabled" in payload and payload.get("enabled") and current.get("status") == STATUS_DISABLED: next_status = STATUS_ACTIVE for field in allowed: if field not in payload: continue if field == "status": continue value = payload[field] if field == "enabled": value = 1 if value else 0 else: value = str(value or "").strip() updates.append(f"{field}=?") values.append(value) if next_status is not None: updates.append("status=?") values.append(str(next_status or "").strip() or STATUS_UNKNOWN) if not updates: return current updates.append("updated_at=?") values.append(utc_now()) values.append(key_id) with self.connect() as conn: result = conn.execute(f"UPDATE api_keys SET {', '.join(updates)} WHERE id=?", values) if result.rowcount == 0: raise KeyError("key not found") row = self.get_key_row(conn, key_id) self.add_event(conn, key_id, "update", "key updated", {"fields": list(payload.keys())}) return self.row_to_public(row) def delete_key(self, key_id: int) -> bool: with self.connect() as conn: row = conn.execute("SELECT id, api_key FROM api_keys WHERE id=?", (key_id,)).fetchone() if row is None: return False self.add_event(conn, key_id, "delete", "key deleted", {"masked_key": mask_key(str(row["api_key"]))}) result = conn.execute("DELETE FROM api_keys WHERE id=?", (key_id,)) return result.rowcount > 0 def active_keys(self) -> Tuple[List[str], List[Dict[str, Any]]]: self.reset_expired_quota_marks() self.quarantine_reported_failures() active_statuses = tuple(sorted(ACTIVE_STATUSES)) placeholders = ",".join("?" for _ in active_statuses) sql = ( "SELECT * FROM api_keys WHERE enabled=1 AND status IN " f"({placeholders}) " "AND COALESCE(last_status_code, 0) NOT IN (401, 403) " "ORDER BY last_success_at DESC, id ASC" ) with self.connect() as conn: rows = conn.execute(sql, active_statuses).fetchall() keys = [str(row["api_key"]) for row in rows] items = [self.row_to_public(row) for row in rows] return keys, items def quarantine_reported_failures(self) -> None: today = datetime.now(timezone.utc).date().isoformat() fallback_reset_at = current_month_last_day() now = utc_now() with self.connect() as conn: rows = conn.execute( """ SELECT id, quota_reset_at FROM api_keys WHERE enabled=1 AND status IN (?, ?) AND last_status_code = 401 AND (quota_reset_at IS NULL OR quota_reset_at > ?) """, (STATUS_ACTIVE, STATUS_UNKNOWN, today), ).fetchall() for row in rows: reset_at = str(row["quota_reset_at"] or fallback_reset_at) conn.execute( "UPDATE api_keys SET status=?, quota_reset_at=?, updated_at=? WHERE id=?", (STATUS_QUOTA_EXHAUSTED, reset_at, now, int(row["id"])), ) self.add_event( conn, int(row["id"]), "auto_quarantine", "key had prior 401; removed from active pool", {"quota_reset_at": reset_at}, ) def reset_expired_quota_marks(self) -> None: today = datetime.now(timezone.utc).date().isoformat() now = utc_now() with self.connect() as conn: rows = conn.execute( """ SELECT id FROM api_keys WHERE enabled=1 AND status=? AND quota_reset_at IS NOT NULL AND quota_reset_at <= ? """, (STATUS_QUOTA_EXHAUSTED, today), ).fetchall() for row in rows: conn.execute( "UPDATE api_keys SET status=?, last_status_code=NULL, last_error='', updated_at=? WHERE id=?", (STATUS_UNKNOWN, now, int(row["id"])), ) self.add_event(conn, int(row["id"]), "reset_due", "quota reset date reached; retry enabled", {}) reset_rows = conn.execute( """ SELECT id FROM api_keys WHERE enabled=1 AND status IN (?, ?) AND last_status_code = 401 AND quota_reset_at IS NOT NULL AND quota_reset_at <= ? """, (STATUS_ACTIVE, STATUS_UNKNOWN, today), ).fetchall() for row in reset_rows: conn.execute( "UPDATE api_keys SET last_status_code=NULL, last_error='', updated_at=? WHERE id=?", (now, int(row["id"])), ) self.add_event(conn, int(row["id"]), "reset_due", "quota reset date reached; retry enabled", {}) def find_quota_email_or_default(self, mailbox: "MailboxClient", address: str) -> Dict[str, Any]: try: return mailbox.find_quota_email(address) except MailboxUnavailable as exc: reset_at = current_month_last_day() return { "reset_at": reset_at, "lookup_error": str(exc), "fallback_reason": "mailbox_unavailable", } def apply_default_quota_reset(self, row: sqlite3.Row, quota_info: Dict[str, Any]) -> Dict[str, Any]: if quota_info.get("reset_at"): return quota_info today = datetime.now(timezone.utc).date().isoformat() existing_reset = str(row["quota_reset_at"] or "").strip() quota_info["reset_at"] = existing_reset if existing_reset > today else current_month_last_day() quota_info.setdefault("fallback_reason", "default_month_end") return quota_info def add_event( self, conn: sqlite3.Connection, key_id: int, event_type: str, message: str, metadata: Dict[str, Any], ) -> None: conn.execute( """ INSERT INTO events (key_id, event_type, message, metadata_json, created_at) VALUES (?, ?, ?, ?, ?) """, (key_id, event_type, message, json.dumps(metadata, ensure_ascii=False), utc_now()), ) def events(self, key_id: int, limit: int = 50) -> List[Dict[str, Any]]: with self.connect() as conn: rows = conn.execute( "SELECT * FROM events WHERE key_id=? ORDER BY id DESC LIMIT ?", (key_id, max(1, min(200, limit))), ).fetchall() return [dict(row) for row in rows] def report_key_status(self, payload: Dict[str, Any], mailbox: "MailboxClient") -> Dict[str, Any]: api_key = str(payload.get("api_key") or "").strip() if not api_key: raise ValueError("api_key is required") status_code = payload.get("status_code") try: status_code_int = int(status_code) if status_code is not None else None except (TypeError, ValueError): status_code_int = None requested_status = str(payload.get("status") or "").strip() error = str(payload.get("error") or payload.get("message") or "").strip() now = utc_now() with self.connect() as conn: row = conn.execute("SELECT * FROM api_keys WHERE api_key=?", (api_key,)).fetchone() if row is None: return {"success": False, "error": "key_not_found"} status = requested_status or STATUS_ERROR quota_info: Dict[str, Any] = {} if status_code_int == 401 or requested_status == STATUS_QUOTA_EXHAUSTED: status = STATUS_QUOTA_EXHAUSTED # Runtime reports must be fast. Mailbox lookups can be slow or # unreachable, so use a safe reset date immediately and let # the HTTP layer enrich it asynchronously. quota_info = self.apply_default_quota_reset(row, quota_info) if quota_info.get("fallback_reason") == "default_month_end": error = ( f"{error}; default quota reset at {quota_info.get('reset_at')}" if error else f"default quota reset at {quota_info.get('reset_at')}" ) conn.execute( """ UPDATE api_keys SET status=?, last_checked_at=?, last_failure_at=?, last_status_code=?, last_error=?, quota_reset_at=COALESCE(?, quota_reset_at), quota_email_id=COALESCE(?, quota_email_id), quota_email_subject=COALESCE(?, quota_email_subject), quota_email_seen_at=COALESCE(?, quota_email_seen_at), updated_at=? WHERE id=? """, ( status, now, now, status_code_int, error[:1000], quota_info.get("reset_at"), quota_info.get("email_id"), quota_info.get("subject"), quota_info.get("seen_at"), now, int(row["id"]), ), ) self.add_event( conn, int(row["id"]), "runtime_report", error[:500] or status, {"status": status, "status_code": status_code_int, "quota": quota_info}, ) updated = conn.execute("SELECT * FROM api_keys WHERE id=?", (int(row["id"]),)).fetchone() return {"success": True, "item": self.row_to_public(updated)} def update_quota_info(self, key_id: int, quota_info: Dict[str, Any]) -> None: """Persist mailbox-derived quota metadata without changing key state.""" reset_at = str(quota_info.get("reset_at") or "").strip() if not reset_at: return now = utc_now() with self.connect() as conn: row = self.get_key_row(conn, key_id) if row is None: return conn.execute( """ UPDATE api_keys SET quota_reset_at=?, quota_email_id=COALESCE(?, quota_email_id), quota_email_subject=COALESCE(?, quota_email_subject), quota_email_seen_at=COALESCE(?, quota_email_seen_at), updated_at=? WHERE id=? """, ( reset_at, quota_info.get("email_id"), quota_info.get("subject"), quota_info.get("seen_at"), now, key_id, ), ) def recheck_candidate_ids(self, limit: int, min_age_seconds: float) -> List[int]: cutoff = ( datetime.now(timezone.utc) - timedelta(seconds=max(0.0, min_age_seconds)) ).isoformat(timespec="seconds") with self.connect() as conn: rows = conn.execute( """ SELECT id FROM api_keys WHERE enabled=1 AND status=? AND (last_checked_at IS NULL OR last_checked_at <= ?) ORDER BY COALESCE(last_checked_at, '') ASC, id ASC LIMIT ? """, (STATUS_UNKNOWN, cutoff, max(1, limit)), ).fetchall() return [int(row["id"]) for row in rows] def check_one( self, key_id: int, mistral: "MistralClient", mailbox: "MailboxClient", *, lookup_mailbox: bool = True, ) -> Dict[str, Any]: now = utc_now() with self.connect() as conn: row = self.get_key_row(conn, key_id) if row is None: raise KeyError("key not found") if not row["enabled"]: conn.execute( "UPDATE api_keys SET status=?, last_checked_at=?, updated_at=? WHERE id=?", (STATUS_DISABLED, now, now, key_id), ) self.add_event(conn, key_id, "check", "key disabled", {}) updated = self.get_key_row(conn, key_id) return self.row_to_public(updated) result = mistral.validate_key(str(row["api_key"])) quota_info: Dict[str, Any] = {} status = STATUS_ERROR last_success_at = row["last_success_at"] last_failure_at = now last_error = result.get("error") or "" if result["ok"]: status = STATUS_ACTIVE last_success_at = now last_failure_at = row["last_failure_at"] last_error = "" elif result.get("status_code") == 401: status = STATUS_QUOTA_EXHAUSTED if lookup_mailbox: quota_info = self.find_quota_email_or_default(mailbox, str(row["email"] or "")) quota_info = self.apply_default_quota_reset(row, quota_info) if quota_info.get("lookup_error"): last_error = f"401 from Mistral; mailbox unavailable, default quota reset at {quota_info.get('reset_at')}" elif quota_info.get("fallback_reason") == "default_month_end": last_error = f"401 from Mistral; quota email reset date was not found; default quota reset at {quota_info.get('reset_at')}" elif result.get("status_code") in {403}: status = STATUS_ERROR elif result.get("status_code") is None: status = STATUS_UNKNOWN conn.execute( """ UPDATE api_keys SET status=?, last_checked_at=?, last_success_at=?, last_failure_at=?, last_status_code=?, last_error=?, quota_reset_at=?, quota_email_id=?, quota_email_subject=?, quota_email_seen_at=?, updated_at=? WHERE id=? """, ( status, now, last_success_at, last_failure_at, result.get("status_code"), last_error[:1000], quota_info.get("reset_at") or row["quota_reset_at"], quota_info.get("email_id") or row["quota_email_id"], quota_info.get("subject") or row["quota_email_subject"], quota_info.get("seen_at") or row["quota_email_seen_at"], now, key_id, ), ) self.add_event( conn, key_id, "check", "mistral validation completed", {"mistral": result, "quota": quota_info}, ) updated = self.get_key_row(conn, key_id) return self.row_to_public(updated) def summary(self) -> Dict[str, Any]: with self.connect() as conn: rows = conn.execute("SELECT status, enabled, COUNT(*) AS count FROM api_keys GROUP BY status, enabled").fetchall() totals: Dict[str, Any] = {"total": 0, "enabled": 0, "disabled": 0, "by_status": {}} for row in rows: count = int(row["count"]) totals["total"] += count if row["enabled"]: totals["enabled"] += count else: totals["disabled"] += count totals["by_status"][row["status"]] = totals["by_status"].get(row["status"], 0) + count return totals class MistralClient: def __init__(self, base_url: str, proxy: str, timeout_sec: float) -> None: self.base_url = base_url.rstrip("/") self.proxy = proxy.strip() self.timeout_sec = max(1.0, timeout_sec) def opener(self) -> urllib.request.OpenerDirector: if not self.proxy: return urllib.request.build_opener() return urllib.request.build_opener( urllib.request.ProxyHandler({"http": self.proxy, "https": self.proxy}) ) def validate_key(self, api_key: str) -> Dict[str, Any]: request = urllib.request.Request( f"{self.base_url}/models", headers={ "Authorization": f"Bearer {api_key}", "Accept": "application/json", "User-Agent": "yunwo-mistral-key-manager/1.0", }, method="GET", ) try: with self.opener().open(request, timeout=self.timeout_sec) as response: body = response.read(4096).decode("utf-8", errors="ignore") return {"ok": 200 <= response.status < 300, "status_code": response.status, "body": body[:500]} except urllib.error.HTTPError as exc: body = exc.read(4096).decode("utf-8", errors="ignore") return {"ok": False, "status_code": exc.code, "error": body[:500] or str(exc)} except Exception as exc: return {"ok": False, "status_code": None, "error": str(exc)} class MailboxClient: def __init__( self, base_url: str, timeout_sec: float, username: str = "", password: str = "", token: str = "", api_key: str = "", ) -> None: self.base_url = base_url.rstrip("/") self.timeout_sec = max(1.0, timeout_sec) self.username = username.strip() self.password = password.strip() self.token = token.strip() self.api_key = api_key.strip() def request_json( self, path: str, *, method: str = "GET", params: Optional[Dict[str, Any]] = None, payload: Optional[Dict[str, Any]] = None, authenticated: bool = True, retry_auth: bool = True, ) -> Dict[str, Any]: if authenticated and not self.api_key: self.ensure_token() url = f"{self.base_url}{path}" if params: url = f"{url}?{urllib.parse.urlencode(params)}" body: Optional[bytes] = None headers = { "Accept": "application/json", "User-Agent": "yunwo-mistral-key-manager/1.0", } if payload is not None: body = json_dumps(payload) headers["Content-Type"] = "application/json; charset=utf-8" if authenticated: bearer = self.api_key or self.token if bearer: headers["Authorization"] = f"Bearer {bearer}" request = urllib.request.Request(url, data=body, headers=headers, method=method) try: with urllib.request.urlopen(request, timeout=self.timeout_sec) as response: return json.loads(response.read().decode("utf-8", errors="ignore")) except urllib.error.HTTPError as exc: if authenticated and not self.api_key and retry_auth and exc.code == 401 and self.username and self.password: self.token = "" self.ensure_token() return self.request_json( path, method=method, params=params, payload=payload, authenticated=authenticated, retry_auth=False, ) raise def ensure_token(self) -> None: if self.token: return if not self.username or not self.password: raise RuntimeError("MAILBOX_API_USERNAME and MAILBOX_API_PASSWORD are required") payload = self.request_json( "/api/v1/auth/token", method="POST", payload={"username": self.username, "password": self.password}, authenticated=False, ) data = payload.get("data") if isinstance(payload, dict) else None token = data.get("token") if isinstance(data, dict) else None if not token: raise RuntimeError("mailbox API did not return an auth token") self.token = str(token) def fetch_emails(self, address: str, size: int = 20) -> List[Dict[str, Any]]: if not address: return [] try: if self.api_key: payload = self.request_json( "/api/emails", params={ "address": address, "page": 1, "limit": min(max(1, size), 100), "folder": "inbox", }, ) messages = payload.get("emails") if isinstance(payload, dict) else None return messages if isinstance(messages, list) else [] payload = self.request_json( "/api/v1/messages", params={"address": address, "page": 1, "size": size}, ) except Exception as exc: message = f"Mailbox API fetch failed for {address}: {exc}" print(message, file=sys.stderr) raise MailboxUnavailable(message) from exc data = payload.get("data") if isinstance(payload, dict) else None messages = data.get("messages") if isinstance(data, dict) else None return messages if isinstance(messages, list) else [] def find_quota_email(self, address: str) -> Dict[str, Any]: emails = self.fetch_emails(address) for listed_email in emails: email = listed_email if self.api_key and listed_email.get("id"): try: detail = self.request_json( f"/api/emails/{urllib.parse.quote(str(listed_email['id']), safe='')}", ) if isinstance(detail, dict): email = {**listed_email, **detail} except Exception as exc: print(f"Mailbox detail fetch failed for {listed_email.get('id')}: {exc}", file=sys.stderr) continue subject = str(email.get("subject") or "") sender = " ".join( str(email.get(field) or "") for field in ("sender", "from_addr", "from_name") ) text = "\n".join( str(email.get(field) or "") for field in ( "subject", "text_body", "html_body", "text_content", "html_content", ) ) lowered = text.lower() if "mistral" not in sender.lower() and "mistral" not in lowered: continue if not ( "monthly quota" in lowered or "reached 100%" in lowered or "quota will reset" in lowered ): continue return { "email_id": email.get("id"), "subject": subject, "reset_at": parse_reset_date(text), "seen_at": utc_now(), } return {} class App: def __init__(self, store: KeyStore, mistral: MistralClient, mailbox: MailboxClient, token: str) -> None: self.store = store self.mistral = mistral self.mailbox = mailbox self.token = token.strip() self._quota_lookup_executor = ThreadPoolExecutor( max_workers=4, thread_name_prefix="mistral-quota-mailbox", ) self._quota_lookup_lock = threading.Lock() self._quota_lookup_inflight: set[int] = set() self._recheck_interval_sec = max( 30.0, read_env_float("MISTRAL_KEY_MANAGER_RECHECK_SEC", 300.0), ) self._recheck_batch_size = max( 1, int(read_env_float("MISTRAL_KEY_MANAGER_RECHECK_BATCH", 16.0)), ) recheck_workers = max( 1, int(read_env_float("MISTRAL_KEY_MANAGER_RECHECK_WORKERS", 4.0)), ) self._recheck_executor = ThreadPoolExecutor( max_workers=recheck_workers, thread_name_prefix="mistral-key-recheck", ) self._recheck_stop = threading.Event() self._recheck_thread = threading.Thread( target=self._recheck_loop, name="mistral-key-recheck-scheduler", daemon=True, ) self._recheck_thread.start() def schedule_quota_lookup(self, key_id: int, email: str) -> None: """Enrich a 401 report without holding the parser request open.""" if not email: return with self._quota_lookup_lock: if key_id in self._quota_lookup_inflight: return self._quota_lookup_inflight.add(key_id) self._quota_lookup_executor.submit(self._lookup_quota_email, key_id, email) def _lookup_quota_email(self, key_id: int, email: str) -> None: try: quota_info = self.store.find_quota_email_or_default(self.mailbox, email) if quota_info.get("reset_at") and not quota_info.get("lookup_error"): self.store.update_quota_info(key_id, quota_info) except Exception: traceback.print_exc() finally: with self._quota_lookup_lock: self._quota_lookup_inflight.discard(key_id) def _recheck_loop(self) -> None: while not self._recheck_stop.is_set(): try: self.store.reset_expired_quota_marks() key_ids = self.store.recheck_candidate_ids( self._recheck_batch_size, self._recheck_interval_sec, ) for key_id in key_ids: self._recheck_executor.submit(self._recheck_key, key_id) except Exception: traceback.print_exc() self._recheck_stop.wait(self._recheck_interval_sec) def _recheck_key(self, key_id: int) -> None: try: item = self.store.check_one( key_id, self.mistral, self.mailbox, lookup_mailbox=False, ) if item.get("status") == STATUS_QUOTA_EXHAUSTED and item.get("email"): self.schedule_quota_lookup(key_id, str(item["email"])) except Exception: traceback.print_exc() def authorized(self, handler: BaseHTTPRequestHandler) -> bool: if not self.token: return True auth = handler.headers.get("Authorization", "") token = handler.headers.get("X-Admin-Token", "") if auth.startswith("Bearer "): token = auth[7:].strip() return token == self.token APP: Optional[App] = None class Handler(BaseHTTPRequestHandler): server_version = "YunwoMistralKeyManager/1.0" def log_message(self, fmt: str, *args: Any) -> None: sys.stderr.write("[%s] %s\n" % (utc_now(), fmt % args)) def send_json(self, data: Any, status: int = 200) -> None: try: body = json_dumps(data) self.send_response(status) self.send_header("Content-Type", "application/json; charset=utf-8") self.send_header("Content-Length", str(len(body))) self.send_header("Access-Control-Allow-Origin", "*") self.send_header("Access-Control-Allow-Headers", "Authorization, X-Admin-Token, Content-Type") self.send_header("Access-Control-Allow-Methods", "GET, POST, PATCH, DELETE, OPTIONS") self.end_headers() self.wfile.write(body) except BrokenPipeError: # The caller may time out while a slow mailbox lookup is running. # A closed client socket is not a server error. return def send_error_json(self, status: int, message: str) -> None: self.send_json({"success": False, "error": message}, status=status) def do_OPTIONS(self) -> None: self.send_response(204) self.send_header("Access-Control-Allow-Origin", "*") self.send_header("Access-Control-Allow-Headers", "Authorization, X-Admin-Token, Content-Type") self.send_header("Access-Control-Allow-Methods", "GET, POST, PATCH, DELETE, OPTIONS") self.end_headers() def read_json_body(self) -> Dict[str, Any]: length = int(self.headers.get("Content-Length") or 0) if length <= 0: return {} if length > 5 * 1024 * 1024: raise ValueError("request body is too large") raw = self.rfile.read(length) return json.loads(raw.decode("utf-8")) def require_app(self) -> App: if APP is None: raise RuntimeError("app not initialized") return APP def ensure_auth(self) -> bool: app = self.require_app() if app.authorized(self): return True self.send_error_json(HTTPStatus.UNAUTHORIZED, "unauthorized") return False def route_api(self, method: str, path: str, query: Dict[str, List[str]]) -> bool: app = self.require_app() if path == "/api/health" and method == "GET": self.send_json( { "success": True, "time": utc_now(), "db": str(app.store.db_path), "mailbox_api_base": app.mailbox.base_url, "mistral_api_base": app.mistral.base_url, "has_token": bool(app.token), "summary": app.store.summary(), } ) return True if path.startswith("/api/") and not self.ensure_auth(): return True if path == "/api/config" and method == "GET": self.send_json( { "success": True, "mailbox_api_base": app.mailbox.base_url, "mistral_api_base": app.mistral.base_url, "has_token": bool(app.token), } ) return True if path == "/api/keys" and method == "GET": items = app.store.list_keys(query) self.send_json({"success": True, "items": items, "summary": app.store.summary()}) return True if path == "/api/keys" and method == "POST": item = app.store.add_or_update_key(self.read_json_body()) self.send_json({"success": True, "item": item}, status=201) return True if path == "/api/keys/bulk" and method == "POST": payload = self.read_json_body() rows = parse_bulk_lines(str(payload.get("text") or "")) items = [app.store.add_or_update_key(row) for row in rows] self.send_json({"success": True, "count": len(items), "items": items}, status=201) return True if path == "/api/keys/active" and method == "GET": keys, items = app.store.active_keys() self.send_json( { "success": True, "keys": keys, "items": items, "generated_at": utc_now(), "count": len(keys), } ) return True if path == "/api/keys/report" and method == "POST": payload = self.read_json_body() result = app.store.report_key_status(payload, app.mailbox) is_quota_report = str(payload.get("status_code") or "").strip() == "401" if ( result.get("success") and (is_quota_report or payload.get("status") == STATUS_QUOTA_EXHAUSTED) ): item = result.get("item") or {} try: key_id = int(item.get("id")) except (TypeError, ValueError): key_id = 0 if key_id and item.get("email"): app.schedule_quota_lookup(key_id, str(item["email"])) self.send_json(result, status=200 if result.get("success") else 404) return True if path == "/api/keys/check-all" and method == "POST": payload = self.read_json_body() enabled_only = bool(payload.get("enabled_only", True)) params = {"enabled": ["1"]} if enabled_only else {} ids = [int(item["id"]) for item in app.store.list_keys(params)] items = [] for key_id in ids: item = app.store.check_one( key_id, app.mistral, app.mailbox, lookup_mailbox=False, ) if item.get("status") == STATUS_QUOTA_EXHAUSTED and item.get("email"): app.schedule_quota_lookup(key_id, str(item["email"])) items.append(item) self.send_json({"success": True, "count": len(items), "items": items}) return True match = re.match(r"^/api/keys/(\d+)$", path) if match: key_id = int(match.group(1)) if method == "GET": item = app.store.get_key(key_id) if item is None: self.send_error_json(404, "key not found") else: self.send_json({"success": True, "item": item}) return True if method == "PATCH": item = app.store.patch_key(key_id, self.read_json_body()) self.send_json({"success": True, "item": item}) return True if method == "DELETE": if app.store.delete_key(key_id): self.send_json({"success": True}) else: self.send_error_json(404, "key not found") return True match = re.match(r"^/api/keys/(\d+)/check$", path) if match and method == "POST": key_id = int(match.group(1)) item = app.store.check_one( key_id, app.mistral, app.mailbox, lookup_mailbox=False, ) if item.get("status") == STATUS_QUOTA_EXHAUSTED and item.get("email"): app.schedule_quota_lookup(key_id, str(item["email"])) self.send_json({"success": True, "item": item}) return True match = re.match(r"^/api/keys/(\d+)/events$", path) if match and method == "GET": events = app.store.events(int(match.group(1))) self.send_json({"success": True, "events": events}) return True return False def do_GET(self) -> None: self.handle_request("GET") def do_POST(self) -> None: self.handle_request("POST") def do_PATCH(self) -> None: self.handle_request("PATCH") def do_DELETE(self) -> None: self.handle_request("DELETE") def handle_request(self, method: str) -> None: try: parsed = urllib.parse.urlparse(self.path) path = parsed.path query = urllib.parse.parse_qs(parsed.query) if path.startswith("/api/"): if not self.route_api(method, path, query): self.send_error_json(404, "not found") return self.serve_static(path) except KeyError as exc: self.send_error_json(404, str(exc)) except ValueError as exc: self.send_error_json(400, str(exc)) except json.JSONDecodeError: self.send_error_json(400, "invalid json") except BrokenPipeError: return except Exception as exc: traceback.print_exc() self.send_error_json(500, str(exc)) def serve_static(self, path: str) -> None: if path in {"", "/"}: file_path = STATIC_DIR / "index.html" else: normalized = urllib.parse.unquote(path).lstrip("/") if normalized.startswith("static/"): normalized = normalized[len("static/") :] file_path = (STATIC_DIR / normalized).resolve() if not str(file_path).startswith(str(STATIC_DIR.resolve())): self.send_error(403) return if not file_path.exists() or not file_path.is_file(): self.send_error(404) return suffix = file_path.suffix.lower() content_type = { ".html": "text/html; charset=utf-8", ".css": "text/css; charset=utf-8", ".js": "application/javascript; charset=utf-8", ".svg": "image/svg+xml", }.get(suffix, "application/octet-stream") body = file_path.read_bytes() self.send_response(200) self.send_header("Content-Type", content_type) self.send_header("Content-Length", str(len(body))) self.end_headers() self.wfile.write(body) def build_app(args: argparse.Namespace) -> App: db_path = Path(args.db or os.getenv("MISTRAL_KEY_MANAGER_DB") or DEFAULT_DB_PATH) mailbox_base = args.mailbox_api_base or os.getenv("MAILBOX_API_BASE") or DEFAULT_MAILBOX_API_BASE mailbox_username = args.mailbox_api_username or os.getenv("MAILBOX_API_USERNAME", DEFAULT_MAILBOX_API_USERNAME) mailbox_password = args.mailbox_api_password or os.getenv("MAILBOX_API_PASSWORD", DEFAULT_MAILBOX_API_PASSWORD) mailbox_token = args.mailbox_api_token or os.getenv("MAILBOX_API_TOKEN", "") mailbox_api_key = args.mailbox_api_key or os.getenv("MAILBOX_API_KEY", "") mistral_base = args.mistral_api_base or os.getenv("MISTRAL_API_BASE") or DEFAULT_MISTRAL_API_BASE proxy = args.mistral_proxy or os.getenv("MISTRAL_API_PROXY") or os.getenv("MISTRAL_OCR_PROXY") or "" timeout = read_env_float("MISTRAL_KEY_MANAGER_TIMEOUT_SEC", 10.0) token = args.token if args.token is not None else os.getenv("MISTRAL_KEY_MANAGER_TOKEN", "") return App( store=KeyStore(db_path), mistral=MistralClient(mistral_base, proxy, timeout), mailbox=MailboxClient( mailbox_base, timeout, mailbox_username, mailbox_password, mailbox_token, mailbox_api_key, ), token=token or "", ) def seed_from_env(store: KeyStore) -> None: raw = os.getenv("MISTRAL_KEY_MANAGER_SEED_KEYS", "") if not raw: raw = os.getenv("MISTRAL_OCR_API_KEYS", "") email = os.getenv("MISTRAL_KEY_MANAGER_SEED_EMAIL", "") label = os.getenv("MISTRAL_KEY_MANAGER_SEED_LABEL", "env-seed") for api_key in split_api_keys(raw): store.add_or_update_key({"api_key": api_key, "email": email, "label": label, "enabled": True}) def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description="Yunwo Mistral API key manager") parser.add_argument("--host", default=os.getenv("MISTRAL_KEY_MANAGER_HOST", "127.0.0.1")) parser.add_argument("--port", type=int, default=int(os.getenv("MISTRAL_KEY_MANAGER_PORT", "8097"))) parser.add_argument("--db", default=os.getenv("MISTRAL_KEY_MANAGER_DB", str(DEFAULT_DB_PATH))) parser.add_argument("--token", default=None) parser.add_argument("--mailbox-api-base", default=os.getenv("MAILBOX_API_BASE", DEFAULT_MAILBOX_API_BASE)) parser.add_argument("--mailbox-api-username", default=os.getenv("MAILBOX_API_USERNAME", DEFAULT_MAILBOX_API_USERNAME)) parser.add_argument("--mailbox-api-password", default=os.getenv("MAILBOX_API_PASSWORD", DEFAULT_MAILBOX_API_PASSWORD)) parser.add_argument("--mailbox-api-token", default=os.getenv("MAILBOX_API_TOKEN", "")) parser.add_argument("--mailbox-api-key", default=os.getenv("MAILBOX_API_KEY", "")) parser.add_argument("--mistral-api-base", default=os.getenv("MISTRAL_API_BASE", DEFAULT_MISTRAL_API_BASE)) parser.add_argument("--mistral-proxy", default=os.getenv("MISTRAL_API_PROXY", os.getenv("MISTRAL_OCR_PROXY", ""))) return parser.parse_args(argv) def main(argv: Optional[List[str]] = None) -> int: global APP args = parse_args(argv) APP = build_app(args) if read_env_bool("MISTRAL_KEY_MANAGER_SEED_ON_START", False): seed_from_env(APP.store) server = ThreadingHTTPServer((args.host, args.port), Handler) print( f"Mistral key manager listening on http://{args.host}:{args.port} " f"db={APP.store.db_path} mailbox={APP.mailbox.base_url}", flush=True, ) try: server.serve_forever() except KeyboardInterrupt: print("shutting down", flush=True) finally: server.server_close() return 0 if __name__ == "__main__": raise SystemExit(main())