crypto_okx/lib/hub/hub_bridge_auth.py

"""中控 Bridge 鉴权与 SSO(短期登录链)."""
from __future__ import annotations

import hashlib
import hmac
import os
import secrets
import time
from typing import Optional
from urllib.parse import urlencode


def hub_bridge_token() -> str:
    return (os.getenv("HUB_BRIDGE_TOKEN") or "").strip()


def extract_hub_token_from_request(headers) -> str:
    raw = (headers.get("X-Hub-Token") or "").strip()
    if raw:
        return raw
    auth = (headers.get("Authorization") or "").strip()
    if auth.lower().startswith("bearer "):
        return auth[7:].strip()
    return ""


def hub_bridge_token_ok(headers) -> bool:
    expected = hub_bridge_token()
    if not expected:
        return False
    got = extract_hub_token_from_request(headers)
    if not got:
        return False
    return hmac.compare_digest(got, expected)


def make_sso_token(bridge_token: str, *, ttl_sec: int = 120) -> str:
    """exp.hex_hmac — 中控与客户端共用 bridge_token 签名."""
    exp = int(time.time()) + max(30, int(ttl_sec))
    payload = str(exp)
    sig = hmac.new(
        bridge_token.encode("utf-8"),
        payload.encode("utf-8"),
        hashlib.sha256,
    ).hexdigest()
    return f"{payload}.{sig}"


def verify_sso_token(token: str, bridge_token: Optional[str] = None) -> bool:
    expected = (bridge_token if bridge_token is not None else hub_bridge_token()).strip()
    raw = (token or "").strip()
    if not expected or not raw or "." not in raw:
        return False
    payload, sig = raw.rsplit(".", 1)
    try:
        exp = int(payload)
    except ValueError:
        return False
    if exp < int(time.time()):
        return False
    want = hmac.new(
        expected.encode("utf-8"),
        payload.encode("utf-8"),
        hashlib.sha256,
    ).hexdigest()
    return hmac.compare_digest(sig, want)


def build_embed_open_url(
    base_url: str,
    *,
    bridge_token: str,
    embed_path: str = "/embed",
    ttl_sec: int = 120,
) -> str:
    """中控生成带 SSO 的客户端打开地址."""
    root = (base_url or "").rstrip("/")
    path = embed_path or "/embed"
    if not path.startswith("/"):
        path = "/" + path
    token = make_sso_token(bridge_token, ttl_sec=ttl_sec)
    q = urlencode({"token": token, "next": path})
    return f"{root}/hub_sso?{q}"


def extract_embed_sso(request) -> str:
    """中控 iframe 内用的短期会话令牌(Cookie 被拦时的后备)."""
    raw = (request.headers.get("X-Hub-Embed-SSO") or "").strip()
    if raw:
        return raw
    return (request.args.get("embed_sso") or "").strip()


def hub_embed_sso_ok(request) -> bool:
    tok = extract_embed_sso(request)
    return bool(tok and verify_sso_token(tok))


def make_embed_session_token(bridge_token: str | None = None, *, ttl_sec: int = 604800) -> str:
    """嵌入会话令牌(默认 7 天),算法与 SSO 相同."""
    secret = (bridge_token if bridge_token is not None else hub_bridge_token()).strip()
    return make_sso_token(secret, ttl_sec=ttl_sec)


def frame_ancestors_header_value() -> Optional[str]:
    """
    HUB_EMBED_FRAME_ANCESTORS: 空格分隔 Origin,如 https://hub.local http://192.168.1.10:5100
    未配置则不改响应头(保持现状);配 * 则允许任意父页(仅内网慎用).
    """
    raw = (os.getenv("HUB_EMBED_FRAME_ANCESTORS") or "").strip()
    if not raw:
        return None
    if raw == "*":
        return "*"
    parts = [p.strip() for p in raw.split() if p.strip()]
    return " ".join(parts) if parts else None


def ensure_hub_bridge_token(env_path: str) -> str:
    """若 .env 中 HUB_BRIDGE_TOKEN 为空则生成并写入;返回当前 token."""
    from lib.env.env_file_lib import apply_env_updates, env_get, read_env_lines

    path = (env_path or "").strip()
    if not path or not os.path.isfile(path):
        existing = hub_bridge_token()
        if existing:
            return existing
        token = secrets.token_urlsafe(32)
        os.environ["HUB_BRIDGE_TOKEN"] = token
        return token

    lines = read_env_lines(path)
    cur = (env_get(lines, "HUB_BRIDGE_TOKEN") or "").strip()
    if cur:
        os.environ["HUB_BRIDGE_TOKEN"] = cur
        return cur

    token = secrets.token_urlsafe(32)
    apply_env_updates(path, {"HUB_BRIDGE_TOKEN": token})
    os.environ["HUB_BRIDGE_TOKEN"] = token
    return token