eink-push/claude_usage.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Claude 账户级用量 —— 复用 Claude Code 自身的 OAuth 登录态，直接调 Anthropic 的 /api/oauth/usage。

这是 Claude Code `/usage` 命令、claude.ai 用量页共用的同一数据源，反映**整个订阅账户**的 5 小时 / 7 天 /
模型限定 7 天额度：不论额度是在 Claude Code CLI、桌面 App 还是网页消耗的，这里都能看到。
与旧方案（statusline 插件把 Claude Code stdin 里的 rate_limits 落盘、本项目读文件）相比：
不依赖 Claude Code 进程在跑、不依赖任何第三方插件，空闲时也持续刷新。

鉴权复用方式与 statusline 一致——**只读** Claude Code 存下来的 access token：
  1. 环境变量 CLAUDE_CODE_OAUTH_TOKEN；
  2. macOS 登录钥匙串 generic password，service「Claude Code-credentials」，JSON 的 claudeAiOauth.accessToken；
  3. 文件 ~/.claude/.credentials.json（Linux / SSH 场景；尊重 CLAUDE_CONFIG_DIR）。
铁律：
  - **绝不刷新 token**。refresh token 每次使用即轮换，第三方刷新而不写回会让 Claude Code 下次 401 并强制重新登录。
    access token 约 8 小时有效，过期就沿用旧快照（标 stale），等 Claude Code 自己下次启动刷新。
  - **快照有效期 CLAUDE_USAGE_TTL 秒内不再调**（.env 可调，默认 300；0=每次实拉），429 后退避（见 usage_common）。
  - 必须与 Claude Code 在同一台机器、同一用户会话下运行（读钥匙串）——所以在形态 A 里它跑在本机，由
    main.py --target push 上报给 server.py，服务端不持任何凭证。
输出快照结构（与旧 statusline 快照兼容，data.usage_from_raw 直接消费）：
  {five_hour:{utilization, resets_at}, seven_day:{...}, seven_day_fable:{...}|None, updated_at, source}
"""
import json
import os
import subprocess
import time
import urllib.error
import urllib.request

import config
import usage_common

API_URL = "https://api.anthropic.com/api/oauth/usage"
KEYCHAIN_SERVICE = "Claude Code-credentials"
TIMEOUT = 15


class TokenUnavailable(RuntimeError):
    """找不到 token 或已过期——本机状态而非接口限流：退避只需一个 TTL（Claude Code 随时可能刷新钥匙串），不按 BACKOFF 长等。"""
    backoff = config.CLAUDE_USAGE_TTL


# ---------------------------------------------------------------- 凭证（只读）
def _credentials_blob():
    """按优先级取 claudeAiOauth JSON 串；找不到返回 None。"""
    env_tok = os.environ.get("CLAUDE_CODE_OAUTH_TOKEN")
    if env_tok:
        return json.dumps({"claudeAiOauth": {"accessToken": env_tok}})
    try:
        out = subprocess.run(["security", "find-generic-password", "-s", KEYCHAIN_SERVICE, "-w"],
                             capture_output=True, text=True, timeout=5)
        if out.returncode == 0 and out.stdout.strip():
            return out.stdout.strip()
    except (OSError, subprocess.TimeoutExpired):
        pass
    cfg_dir = os.environ.get("CLAUDE_CONFIG_DIR") or os.path.expanduser("~/.claude")
    path = os.path.join(cfg_dir, ".credentials.json")
    if os.path.exists(path):
        with open(path, encoding="utf-8") as f:
            return f.read()
    return None


def read_token():
    """返回未过期的 access token；拿不到/已过期抛 TokenUnavailable。"""
    blob = _credentials_blob()
    if not blob:
        raise TokenUnavailable("未找到 Claude Code 登录凭证（钥匙串「Claude Code-credentials」/ ~/.claude/.credentials.json）")
    try:
        oauth = json.loads(blob).get("claudeAiOauth") or {}
    except json.JSONDecodeError as e:
        raise TokenUnavailable(f"凭证 JSON 解析失败：{e}") from e
    tok = oauth.get("accessToken")
    if not tok:
        raise TokenUnavailable("凭证里没有 accessToken")
    exp = oauth.get("expiresAt")  # 毫秒 epoch；缺失则不判过期
    if exp and float(exp) / 1000 <= time.time() + 30:
        raise TokenUnavailable("Claude Code 的 access token 已过期（约 8 小时有效），等 Claude Code 下次启动自动刷新；本模块不做刷新")
    return tok


# ---------------------------------------------------------------- 拉取与归一
def fetch_raw():
    """GET /api/oauth/usage，返回原始 JSON。直连、不走代理（与本机其他脚本一致）。"""
    req = urllib.request.Request(API_URL, headers={
        "Authorization": f"Bearer {read_token()}",
        "anthropic-beta": "oauth-2025-04-20",
        "Accept": "application/json",
        "Content-Type": "application/json",
        "User-Agent": config.CLAUDE_USAGE_UA,
    })
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
    try:
        with opener.open(req, timeout=TIMEOUT) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", "replace")[:200]
        if e.code == 429:
            raise RuntimeError(f"HTTP 429 被限流（将退避 {config.CLAUDE_USAGE_BACKOFF}s；勿高频调用，否则约 24 小时内连 Claude Code /usage 也看不到）：{body}") from e
        if e.code == 401:
            raise TokenUnavailable(f"HTTP 401 token 无效：{body}") from e
        raise RuntimeError(f"HTTP {e.code}：{body}") from e


def _seg(percent, resets_at):
    if percent is None:
        return None
    return {"utilization": float(percent), "resets_at": resets_at}


def normalize(raw, now=None):
    """原始响应 → 快照。优先用 limits[]（2026 年中起模型限定额度只在这里：kind=weekly_scoped + scope.model.display_name），
    顶层 five_hour / seven_day 作兜底；两处都没有的窗口为 None。"""
    now = now or time.time()
    five = seven = scoped = None
    for lim in raw.get("limits") or []:
        kind = lim.get("kind")
        seg = _seg(lim.get("percent"), lim.get("resets_at"))
        if kind == "session":
            five = five or seg
        elif kind == "weekly_all":
            seven = seven or seg
        elif kind == "weekly_scoped":
            name = ((lim.get("scope") or {}).get("model") or {}).get("display_name") or ""
            if name.lower() == config.CLAUDE_USAGE_SCOPE_MODEL.lower():
                scoped = scoped or seg
    for key, cur in (("five_hour", five), ("seven_day", seven)):
        if cur is None and isinstance(raw.get(key), dict):
            cur = _seg(raw[key].get("utilization"), raw[key].get("resets_at"))
        if key == "five_hour":
            five = cur
        else:
            seven = cur
    if scoped is None and isinstance(raw.get(config.CLAUDE_USAGE_SCOPE_FALLBACK_KEY), dict):
        s = raw[config.CLAUDE_USAGE_SCOPE_FALLBACK_KEY]
        scoped = _seg(s.get("utilization"), s.get("resets_at"))
    return {"five_hour": five, "seven_day": seven, "seven_day_fable": scoped,
            "updated_at": int(now), "source": "oauth-usage"}


def get_usage(force=False):
    """对外入口：带缓存/最小间隔/退避的账户级用量快照（结构见模块说明）。失败且无旧快照时抛出。"""
    snap, status = usage_common.cached_snapshot(
        config.CLAUDE_USAGE_CACHE_PATH, config.CLAUDE_USAGE_TTL, config.CLAUDE_USAGE_BACKOFF,
        lambda: normalize(fetch_raw()), force=force)
    snap = dict(snap)
    snap["status"] = status
    return snap


if __name__ == "__main__":
    import sys
    d = get_usage(force="--force" in sys.argv)
    print(json.dumps(d, ensure_ascii=False, indent=1))