scripts/_hf_token.py

"""HuggingFace Token 共享 helper - 提供 HF 账户切换和 token 获取

提供以下公共 API:
    - get_hf_token_interactive()  必要时 prompt 后获取 token
    - get_hf_username()           获取当前 HF 用户名

设计:
    - 与 bootstrap-hf.sh 切换流程保持一致
    - 非 TTY 时 (stdin.isatty() == False) 不进入 prompt 循环，避免死锁
    - login 成功后调用 hf auth whoami 验证，切换时输出 switched 信息
    - 不实现 RESTORE_HF_LOGIN_* 恢复机制（单命令工具不需要）

用法:
    import sys, os
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from _hf_token import get_hf_token_interactive

    token = get_hf_token_interactive()
    if not token:
        sys.exit(1)
"""

from __future__ import annotations

import getpass
import os
import re
import subprocess
import sys
from typing import Optional

_HF_CLI_TIMEOUT = 10  # hf CLI 调用超时（秒）


def _parse_whoami_username(output: str) -> str:
    """解析 hf auth whoami 输出，提取 username。

    接受两种格式:
        "user: <name>"
        "Logged in as <name>"
    """
    match = re.search(r"user:\s*(\S+)", output or "")
    if not match:
        match = re.search(r"[Ll]ogged in as\s+(\S+)", output or "")
    return match.group(1) if match else ""


_HF_EXTRA_PATH = "/root/.local/bin"


def _run_hf(args: list[str]) -> tuple[int, str, str]:
    """运行 hf CLI 子命令并返回 (returncode, stdout, stderr)。"""
    import shutil as _shutil

    # Check if hf is available at all before trying subprocess
    if not _shutil.which("hf") and not _shutil.which(os.path.join(_HF_EXTRA_PATH, "hf")):
        return (1, "", "hf not found in PATH")

    try:
        result = subprocess.run(
            ["hf", *args],
            capture_output=True,
            text=True,
            timeout=_HF_CLI_TIMEOUT,
        )
        return result.returncode, result.stdout, result.stderr
    except FileNotFoundError:
        pass
    except (subprocess.TimeoutExpired) as exc:
        return 1, "", str(exc)

    # hf not found — retry with /root/.local/bin in PATH (handles cron/systemd env)
    extra_env = os.environ.copy()
    extra_env["PATH"] = _HF_EXTRA_PATH + os.pathsep + extra_env.get("PATH", "")
    try:
        result = subprocess.run(
            ["hf", *args],
            capture_output=True,
            text=True,
            timeout=_HF_CLI_TIMEOUT,
            env=extra_env,
        )
        return result.returncode, result.stdout, result.stderr
    except (subprocess.TimeoutExpired, FileNotFoundError) as exc:
        return 1, "", str(exc)


def _read_cached_token() -> str:
    """从 env 或 cache 文件读取 token。

    优先级: HUGGINGFACE_HUB_TOKEN > HF_TOKEN > HF_TOKEN_FILE > ~/.cache/huggingface/token
    """
    for var in ("HUGGINGFACE_HUB_TOKEN", "HF_TOKEN"):
        val = os.environ.get(var, "").strip()
        if val:
            return val
    token_file = os.environ.get("HF_TOKEN_FILE")
    if not token_file:
        home = os.environ.get("HOME") or "/root"
        token_file = os.path.join(home, ".cache", "huggingface", "token")
    if os.path.exists(token_file):
        try:
            with open(token_file) as fh:
                return fh.read().strip()
        except OSError:
            return ""
    return ""


def get_hf_username() -> Optional[str]:
    """获取当前 HF 用户名。失败返回 None。"""
    rc, stdout, _ = _run_hf(["auth", "whoami"])
    if rc != 0:
        return None
    username = _parse_whoami_username(stdout)
    if username:
        return username
    # Python fallback: 用 cache token 调 HfApi
    token = _read_cached_token()
    if not token:
        return None
    try:
        from huggingface_hub import HfApi

        api = HfApi(token=token)
        data = api.whoami(token=token)
        name = data.get("name", "") if isinstance(data, dict) else ""
        return (name or "").strip() or None
    except Exception:
        return None


def _prompt_for_token() -> Optional[str]:
    """交互式读取一个 token。已确认 stdin 是 TTY。"""
    try:
        token = getpass.getpass("Enter HF_TOKEN: ").strip()
    except (EOFError, KeyboardInterrupt):
        print("\n[WARN] Token input cancelled", file=sys.stderr)
        return None
    if not token:
        print("[WARN] Empty token, please try again", file=sys.stderr)
        return None
    return token


def get_hf_token_interactive() -> Optional[str]:
    """获取 HF token，支持账户切换。

    流程:
        1. 若已登录（hf auth whoami 成功）:
            - TTY: 询问 "Use this user? [Y/n]"，如拒绝则 prompt 新 token 并切换
            - 非 TTY: 静默使用当前用户
        2. 未登录:
            - 尝试 env / cache 文件
            - 仍无: TTY 下 prompt 新 token 登录；非 TTY 返回 None
        3. login 成功后调用 hf auth whoami 验证，输出 switched 信息

    成功返回 token 字符串；非交互且无 token 时返回 None。
    """
    current_user: Optional[str] = None
    is_tty = sys.stdin.isatty()

    # ---- 步骤 1: 检测是否已登录 ----
    rc, stdout, _ = _run_hf(["auth", "whoami"])
    if rc == 0:
        current_user = _parse_whoami_username(stdout) or None
        if current_user:
            print(f"[INFO] HF CLI is logged in as: {current_user}")
            if is_tty:
                response = input("Use this user? (Y/n) [Y]: ").strip().lower()
                if response in ("", "y", "yes"):
                    token = _read_cached_token()
                    if token:
                        return token
                    # cache 缺失却 whoami 成功，尝试从环境拿到
                    env_token = os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get(
                        "HF_TOKEN", ""
                    )
                    if env_token.strip():
                        return env_token.strip()
                # 用户选择切换 → 进入步骤 3
            else:
                # 非 TTY：静默使用当前
                token = _read_cached_token()
                if token:
                    return token
                env_token = os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get(
                    "HF_TOKEN", ""
                )
                if env_token.strip():
                    return env_token.strip()

    # ---- 步骤 2: 未登录时尝试 env / cache ----
    if not current_user:
        token = _read_cached_token()
        if token:
            return token

    # ---- 步骤 3: prompt 新 token 并登录 ----
    if not is_tty:
        print(
            "[ERROR] Cannot prompt for HF_TOKEN in non-interactive mode. "
            "Set HF_TOKEN env var or run: hf auth login",
            file=sys.stderr,
        )
        return None

    print("[INFO] HF CLI is not logged in. Please provide a HF_TOKEN.")
    while True:
        token = _prompt_for_token()
        if not token:
            continue

        login_rc, _, login_stderr = _run_hf(["auth", "login", "--token", token])
        if login_rc != 0:
            print(
                f"[ERROR] HF login failed: {login_stderr.strip() or 'unknown error'}",
                file=sys.stderr,
            )
            continue

        # login 成功 → whoami 验证
        verify_rc, verify_stdout, _ = _run_hf(["auth", "whoami"])
        if verify_rc != 0:
            print(
                "[ERROR] HF login succeeded but whoami verification failed",
                file=sys.stderr,
            )
            continue
        new_user = _parse_whoami_username(verify_stdout) or "unknown"
        if current_user and new_user != current_user:
            print(f"[INFO] HF account switched: {current_user} -> {new_user}")
        elif current_user:
            print(f"[INFO] HF login successful as: {new_user}")
        else:
            print(f"[INFO] HF login successful as: {new_user}")
        return token


if __name__ == "__main__":
    # 简单自测入口
    token = get_hf_token_interactive()
    if token:
        print(f"Got token (len={len(token)})")
        sys.exit(0)
    sys.exit(1)