hmo
c11222e55c
- usage_collector_kimi_local.py: Playwright connect_over_cdp to local Chrome - kimi_login.py: interactive login helper for headless Chrome - Replaces SSH-to-Windows approach, removes single point of failure - Chrome CDP service: chrome-kimi-cdp.service (systemd)
254 lines
9.5 KiB
Python
254 lines
9.5 KiB
Python
#!/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()
|