#!/usr/bin/env python3 """ usage_collector_kimi_local.py — Kimi 用量采集（Playwright + 本地 CDP） ===================================================================== 连接本地 Chrome CDP (127.0.0.1:9222, chrome-kimi-cdp.service)， 打开 Kimi 控制台，提取用量数据，保存到 usage_stats_kimi.json。 替代旧的 SSH-to-Windows 方案，直接在 246 本地运行。 调用方式: python3 usage_collector_kimi_local.py """ import json import os import re import sys import time from pathlib import Path from datetime import datetime, timezone from playwright.sync_api import sync_playwright CDP_URL = "http://127.0.0.1:9222" KIMI_CONSOLE = "https://www.kimi.com/membership/subscription?tab=quota" SCRIPT_DIR = Path(__file__).resolve().parent GATEWAY_DIR = SCRIPT_DIR.parent TEMP_DIR = GATEWAY_DIR / "temp" os.makedirs(str(TEMP_DIR), exist_ok=True) OUTPUT_FILE = TEMP_DIR / "usage_stats_kimi.json" KEY_ID = "key7" PROVIDER = "kimi" def extract_metric(text, keyword, after_window=200): """Find percentage and date reset after a keyword in page text.""" idx = text.find(keyword) if idx < 0: return None, None, "unknown", f"keyword '{keyword}' not found" after = text[idx:idx + after_window] m_pct = re.search(r'(\d+(?:\.\d+)?)%', after) m_date = re.search(r'(\d{4}-\d{2}-\d{2})|(\d{2}-\d{2}\s+\d{2}:\d{2})', after) pct = float(m_pct.group(1)) if m_pct else None date_str = (m_date.group(1) or m_date.group(2)) if m_date else None if pct is None: return None, None, "parse_error", f"no % found near '{keyword}'" status = "ok" if pct >= 95: status = "rate-limited" elif pct >= 80: status = "warn" return pct, date_str, status, None def parse_date(dstr): """Parse date string to seconds remaining from now.""" if not dstr: return 0 now = datetime.now(timezone.utc) try: if '-' in dstr and ':' not in dstr: d = datetime.strptime(dstr, "%Y-%m-%d") return max(0, int((d.replace(tzinfo=timezone.utc) - now).total_seconds())) elif '-' in dstr and ':' in dstr: year = now.year d = datetime.strptime(dstr, "%m-%d %H:%M").replace(year=year, tzinfo=timezone.utc) if d < now: d = d.replace(year=year + 1) return max(0, int((d - now).total_seconds())) except Exception: pass return 0 def collect_metrics(page_text): """Parse usage metrics from page text.""" monthly_keywords = ["总使用量", "总用量", "月额度", "月"] rolling_keywords = ["5 小时用量", "5小时用量", "5小时", "5h", "5 Hour"] weekly_keywords = ["7 天用量", "7天用量", "7天", "7d", "7 Day"] monthly_pct = monthly_date = monthly_status = monthly_err = None rolling_pct = rolling_date = rolling_status = rolling_err = None weekly_pct = weekly_date = weekly_status = weekly_err = None for kw in monthly_keywords: monthly_pct, monthly_date, monthly_status, monthly_err = extract_metric(page_text, kw) if monthly_pct is not None: break for kw in rolling_keywords: rolling_pct, rolling_date, rolling_status, rolling_err = extract_metric(page_text, kw) if rolling_pct is not None: break for kw in weekly_keywords: weekly_pct, weekly_date, weekly_status, weekly_err = extract_metric(page_text, kw) if weekly_pct is not None: break # Fallback: positional assignment if keywords miss none_count = sum(1 for x in [monthly_pct, rolling_pct, weekly_pct] if x is None) if none_count > 0: print(f"[kimi] Keyword matched {3 - none_count}/3 metrics, filling rest positionally") pct_matches = list(re.finditer(r'(\d+(?:\.\d+)?)%', page_text)) date_matches = list(re.finditer(r'(\d{4}-\d{2}-\d{2})|(\d{2}-\d{2}\s+\d{2}:\d{2})', page_text)) entries = [] for m in pct_matches: pct_val = float(m.group(1)) pct_pos = m.start() best_date = None for d in date_matches: if d.start() > pct_pos: best_date = d.group(1) or d.group(2) break entries.append((pct_val, best_date, pct_pos)) metric_keys = ["monthly", "rolling", "weekly"] current = {"monthly": monthly_pct, "rolling": rolling_pct, "weekly": weekly_pct} for i, key in enumerate(metric_keys): if current[key] is None and i < len(entries): if key == "monthly": monthly_pct = entries[i][0] monthly_date = entries[i][1] monthly_status = "ok" elif key == "rolling": rolling_pct = entries[i][0] rolling_date = entries[i][1] rolling_status = "ok" elif key == "weekly": weekly_pct = entries[i][0] weekly_date = entries[i][1] weekly_status = "ok" return { "monthly_pct": monthly_pct if monthly_pct is not None else 0, "monthly_reset": parse_date(monthly_date), "monthly_status": monthly_status or "ok", "rolling_pct": rolling_pct if rolling_pct is not None else 0, "rolling_reset": parse_date(rolling_date), "rolling_status": rolling_status or "ok", "weekly_pct": weekly_pct if weekly_pct is not None else 0, "weekly_reset": parse_date(weekly_date), "weekly_status": weekly_status or "ok", "errors": [e for e in [monthly_err, rolling_err, weekly_err] if e], } def main(): print("[kimi] Starting local Playwright collection...") with sync_playwright() as p: # Connect to the persistent Chrome CDP service try: browser = p.chromium.connect_over_cdp(CDP_URL) except Exception as e: print(f"[kimi] ERROR: Cannot connect to CDP at {CDP_URL}: {e}") print("[kimi] Is chrome-kimi-cdp.service running? (systemctl status chrome-kimi-cdp)") sys.exit(1) # Use the default browser context (has persistent profile with login session) context = browser.contexts[0] if browser.contexts else browser.new_context() # Check if already logged in by looking for existing kimi pages kimi_page = None for page in context.pages: if "kimi.com" in page.url: kimi_page = page print(f"[kimi] Reusing existing Kimi tab: {page.url}") break if not kimi_page: kimi_page = context.new_page() # Navigate to the quota page print(f"[kimi] Navigating to {KIMI_CONSOLE}") kimi_page.goto(KIMI_CONSOLE, wait_until="networkidle", timeout=30000) # Wait for usage data to render (poll for % signs) page_text = "" for attempt in range(20): page_text = kimi_page.inner_text("body") or "" pcts = re.findall(r'(\d+(?:\.\d+)?)%', page_text) non_zero = sum(1 for pct in pcts if float(pct) > 0) if non_zero >= 2: print(f"[kimi] Page ready — {non_zero} non-zero %s (attempt {attempt + 1})") break time.sleep(1) else: print(f"[kimi] Warning: page may not have fully loaded") # Check for login redirect current_url = kimi_page.url if "login" in current_url or "passport" in current_url: print(f"[kimi] ERROR: Redirected to login page: {current_url}") print("[kimi] Need to login first — run: python3 kimi_login.py") browser.close() sys.exit(1) # Parse metrics metrics = collect_metrics(page_text) print(f"[kimi] Monthly: {metrics['monthly_pct']}% | " f"5h: {metrics['rolling_pct']}% | 7d: {metrics['weekly_pct']}%") if metrics["errors"]: print(f"[kimi] Warnings: {'; '.join(metrics['errors'])}") # Screenshot for debugging try: kimi_page.screenshot(path=str(TEMP_DIR / "kimi_console.png")) print("[kimi] Screenshot saved") except Exception: pass # Don't close the page — leave it open for next collection # browser.close() would kill the entire CDP service # Build output output = { "ok": True, "last_refresh_iso": datetime.now(timezone.utc).isoformat(timespec="seconds"), "provider": PROVIDER, "accounts": [{ "key_id": KEY_ID, "label": "key7 (Kimi)", "provider": PROVIDER, "workspace_id": "kimi", "rolling": { "usage_percent": metrics["rolling_pct"], "reset_in_sec": metrics["rolling_reset"], "status": metrics["rolling_status"], }, "weekly": { "usage_percent": metrics["weekly_pct"], "reset_in_sec": metrics["weekly_reset"], "status": metrics["weekly_status"], }, "monthly": { "usage_percent": metrics["monthly_pct"], "reset_in_sec": metrics["monthly_reset"], "status": metrics["monthly_status"], }, "session_expired": False, "last_update_iso": datetime.now(timezone.utc).isoformat(timespec="seconds"), "error": None, }], } with open(str(OUTPUT_FILE), "w", encoding="utf-8") as f: json.dump(output, f, ensure_ascii=False, indent=2) print(f"[kimi] Output saved to {OUTPUT_FILE}") if __name__ == "__main__": main()