excel-toolkit/scripts/auth_runtime.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Python 版本的 OpenClaw Auth Runtime
基于 ~/clawd/skills/_shared/auth-runtime 的 TypeScript 实现

功能：
- 使用 CLIENT_KEY 获取访问令牌
- 支持令牌缓存（可配置 TTL）
- 自动刷新过期令牌
- 401/403 自动重试
"""

import hashlib
import json
import os
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional

import requests  # type: ignore


@dataclass
class EnvConfig:
    """环境配置"""
    auth_base: str = "https://api-gw-test.yuanwei-lnc.com"
    client_key: str = ""
    auth_cache_dir: str = "/tmp/skill-auth-cache"
    auth_min_ttl_sec: int = 60


@dataclass
class SessionResponse:
    """会话响应"""
    access_token: str
    hook_url: Optional[str] = None
    hook_token: Optional[str] = None
    expires_in: int = 900


@dataclass
class CachedTokenData:
    """缓存的令牌数据"""
    access_token: str
    expires_at: float


@dataclass
class ApiResponse:
    """HTTP 响应"""
    status: int
    body: str
    headers: dict[str, str]


# 可重试的状态码
RETRYABLE_STATUS = {401, 403}

# 可重试的响应体标记
RETRYABLE_BODY_MARKERS = [
    'session not found or expired',
    'invalid or expired token',
    'unauthorized',
    'client key expired',
    'client key revoked',
]


def create_env_config() -> EnvConfig:
    """
    从环境变量创建配置
    
    环境变量：
    - AUTH_BASE: 认证基础 URL（默认：https://api-gw-test.yuanwei-lnc.com）
    - CLIENT_KEY: 客户端密钥（必需）
    - AUTH_CACHE_DIR: 缓存目录（默认：/tmp/skill-auth-cache）
    - AUTH_MIN_TTL_SEC: 最小令牌 TTL 秒数（默认：60）
    """
    auth_base = os.getenv("AUTH_BASE", "https://api-gw-test.yuanwei-lnc.com").rstrip("/")
    client_key = os.getenv("CLIENT_KEY", "")
    auth_cache_dir = os.getenv("AUTH_CACHE_DIR", "/tmp/skill-auth-cache")
    auth_min_ttl_sec = int(os.getenv("AUTH_MIN_TTL_SEC", "60"))
    
    return EnvConfig(
        auth_base=auth_base,
        client_key=client_key,
        auth_cache_dir=auth_cache_dir,
        auth_min_ttl_sec=auth_min_ttl_sec,
    )


def get_cache_file(auth_base: str, client_key: str, cache_dir: str) -> Path:
    """
    获取缓存文件路径
    
    使用 auth_base 和 client_key 的哈希值生成唯一的缓存文件名
    """
    cache_path = Path(cache_dir)
    cache_path.mkdir(parents=True, exist_ok=True)
    
    # 生成缓存文件名的哈希
    key_str = f"{auth_base}:{client_key}"
    hash_value = hashlib.sha256(key_str.encode()).hexdigest()[:16]
    
    return cache_path / f"token_{hash_value}.json"


def read_cached_token(cache_file: Path, min_ttl_sec: int) -> Optional[str]:
    """
    读取缓存的令牌
    
    如果令牌存在且剩余 TTL 大于最小 TTL，则返回令牌
    """
    if not cache_file.exists():
        return None
    
    try:
        with open(cache_file, "r", encoding="utf-8") as f:
            data = json.load(f)
        
        cached = CachedTokenData(
            access_token=data["access_token"],
            expires_at=data["expires_at"],
        )
        
        # 检查是否过期（考虑最小 TTL）
        if time.time() + min_ttl_sec < cached.expires_at:
            return cached.access_token
        else:
            # 令牌已过期或即将过期，删除缓存
            cache_file.unlink(missing_ok=True)
            return None
    
    except (json.JSONDecodeError, KeyError, Exception):
        # 缓存损坏，删除并重新获取
        cache_file.unlink(missing_ok=True)
        return None


def write_cache(cache_file: Path, session: SessionResponse) -> None:
    """
    写入缓存
    
    缓存令牌和过期时间
    """
    cache_file.parent.mkdir(parents=True, exist_ok=True)
    
    data = {
        "access_token": session.access_token,
        "expires_at": time.time() + session.expires_in,
        "expires_in": session.expires_in,
    }
    
    with open(cache_file, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)


def delete_cache(cache_file: Path) -> None:
    """删除缓存"""
    cache_file.unlink(missing_ok=True)


def fetch_session_json(dry_run: bool, config: EnvConfig) -> SessionResponse:
    """
    从认证端点获取会话 JSON
    
    使用 CLIENT_KEY 请求访问令牌
    """
    if dry_run:
        return SessionResponse(
            access_token="<dry-run-token>",
            hook_url="<dry-run-hook-url>",
            hook_token="<dry-run-hook-token>",
            expires_in=900,
        )
    
    if not config.client_key:
        raise ValueError("CLIENT_KEY is required")
    
    # 请求认证端点
    payload = {"clientKey": config.client_key}
    url = f"{config.auth_base}/auth/skill-credit/session"
    
    response = requests.post(url, json=payload, timeout=30)
    
    if response.status_code < 200 or response.status_code >= 300:
        raise RuntimeError(
            f"Auth session request failed: HTTP {response.status_code} - {response.text}"
        )
    
    session_data = response.json()
    
    if not session_data.get("accessToken"):
        raise RuntimeError(f"Missing accessToken in session response: {response.text}")
    
    return SessionResponse(
        access_token=session_data["accessToken"],
        hook_url=session_data.get("hookUrl"),
        hook_token=session_data.get("hookToken"),
        expires_in=session_data.get("expiresIn", 900),
    )


def get_access_token(dry_run: bool, config: EnvConfig) -> str:
    """
    获取访问令牌（带缓存）
    
    1. 检查缓存
    2. 如果缓存有效，返回缓存的令牌
    3. 否则请求新令牌并缓存
    """
    if dry_run:
        return "<dry-run-token>"
    
    if not config.client_key:
        raise ValueError("CLIENT_KEY is required")
    
    cache_file = get_cache_file(config.auth_base, config.client_key, config.auth_cache_dir)
    cached_token = read_cached_token(cache_file, config.auth_min_ttl_sec)
    
    if cached_token:
        print(f"✓ 使用缓存的访问令牌")
        return cached_token
    
    print(f"🔑 请求新的访问令牌...")
    session = fetch_session_json(dry_run, config)
    write_cache(cache_file, session)
    print(f"✓ 令牌已缓存到：{cache_file}")
    
    return session.access_token


def refresh_access_token(dry_run: bool, config: EnvConfig) -> str:
    """
    刷新访问令牌（绕过缓存）
    
    删除缓存并重新请求新令牌
    """
    if dry_run:
        return "<dry-run-token>"
    
    if not config.client_key:
        raise ValueError("CLIENT_KEY is required")
    
    cache_file = get_cache_file(config.auth_base, config.client_key, config.auth_cache_dir)
    delete_cache(cache_file)
    print(f"🔄 刷新访问令牌...")
    
    return get_access_token(dry_run, config)


def is_retryable_session_error(response: ApiResponse) -> bool:
    """
    检查响应是否表示会话过期/无效
    
    如果是 401/403 或响应体包含特定标记，则可以重试
    """
    if response.status not in RETRYABLE_STATUS:
        return False
    
    body = (response.body or "").lower()
    if not body:
        return True
    
    return any(marker in body for marker in RETRYABLE_BODY_MARKERS)


def request_api(
    method: str,
    url: str,
    auth_token: Optional[str] = None,
    body: Optional[dict[str, Any]] = None,
    timeout: int = 30,
) -> ApiResponse:
    """
    发送 HTTP 请求
    
    Args:
        method: HTTP 方法（GET, POST, PUT, DELETE 等）
        url: 请求 URL
        auth_token: 访问令牌（可选）
        body: 请求体（可选，自动序列化为 JSON）
        timeout: 超时时间（秒）
    
    Returns:
        ApiResponse: 响应对象
    """
    headers = {"Content-Type": "application/json"}
    
    if auth_token:
        headers["Authorization"] = f"Bearer {auth_token}"
    
    try:
        if method.upper() == "GET":
            response = requests.get(url, headers=headers, timeout=timeout)
        elif method.upper() == "POST":
            response = requests.post(
                url, headers=headers, json=body, timeout=timeout
            )
        elif method.upper() == "PUT":
            response = requests.put(
                url, headers=headers, json=body, timeout=timeout
            )
        elif method.upper() == "DELETE":
            response = requests.delete(url, headers=headers, timeout=timeout)
        else:
            raise ValueError(f"Unsupported HTTP method: {method}")
        
        return ApiResponse(
            status=response.status_code,
            body=response.text,
            headers=dict(response.headers),
        )
    
    except requests.exceptions.RequestException as e:
        return ApiResponse(
            status=0,
            body=str(e),
            headers={},
        )


def request_api_with_auto_refresh(
    method: str,
    url: str,
    dry_run: bool,
    config: EnvConfig,
    body: Optional[dict[str, Any]] = None,
    access_token: Optional[str] = None,
) -> ApiResponse:
    """
    发送 API 请求并自动刷新令牌
    
    1. 使用当前令牌发送请求
    2. 如果是 401/403 错误，刷新令牌并重试一次
    """
    token = access_token or get_access_token(dry_run, config)
    
    # 第一次请求
    first = request_api(method, url, token, body)
    
    # 检查是否需要重试
    if not is_retryable_session_error(first):
        return first
    
    # 刷新令牌并重试
    print(f"⚠️ 检测到会话过期，刷新令牌后重试...")
    fresh_token = refresh_access_token(dry_run, config)
    return request_api(method, url, fresh_token, body)


# 便捷函数：从 .env 文件加载 CLIENT_KEY
def load_client_key_from_env(env_path: Optional[Path] = None) -> str:
    """
    从 .env 文件加载 CLIENT_KEY
    
    优先级：
    1. 环境变量 CLIENT_KEY
    2. .env 文件中的 CLIENT_KEY
    3. 抛出异常
    """
    # 先检查环境变量
    client_key = os.getenv("CLIENT_KEY")
    if client_key:
        return client_key
    
    # 查找 .env 文件
    if env_path is None:
        possible_paths = [
            Path(".env"),
            Path.cwd() / ".env",
            Path(__file__).parent.parent / ".env",
        ]
        for p in possible_paths:
            if p.exists():
                env_path = p
                break
    
    if env_path and env_path.exists():
        with open(env_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                if "=" in line:
                    key, value = line.split("=", 1)
                    if key.strip() == "CLIENT_KEY":
                        return value.strip()
    
    raise ValueError(
        "CLIENT_KEY not found. Please set CLIENT_KEY environment variable "
        "or add it to .env file"
    )