crypto_okx/lib/hub/hub_bridge_register.py

"""中控只读 Bridge 路由 + SSO 登录入口."""
from __future__ import annotations

from typing import Any

from flask import Flask, jsonify, redirect, request, session
from urllib.parse import quote

from lib.hub.hub_bridge_auth import (
    frame_ancestors_header_value,
    hub_bridge_token,
    hub_bridge_token_ok,
    make_embed_session_token,
    verify_sso_token,
)


def _safe_next_path(raw: str, default: str = "/embed") -> str:
    p = (raw or "").strip() or default
    if not p.startswith("/"):
        p = "/" + p
    if "://" in p or p.startswith("//"):
        return default
    return p


def _enrich_hub_protocol(snap: dict[str, Any], *, app_module: Any) -> dict[str, Any]:
    """对齐中控设计文档 version=1 字段,同时保留原 snapshot 键."""
    out = dict(snap or {})
    out["version"] = 1
    mode = "live"
    try:
        from lib.sim.mode_lib import is_sim_mode

        get_db = getattr(app_module, "get_db", None)
        if callable(get_db) and is_sim_mode(get_db):
            mode = "sim"
    except Exception:
        mode = "live"
    label = ""
    try:
        from lib.sim.mode_lib import exchange_mode_label

        get_db = getattr(app_module, "get_db", None)
        if callable(get_db):
            label = exchange_mode_label(get_db=get_db, exchange_key="okx") or ""
    except Exception:
        label = ""
    if not label:
        label = "okx_sim" if mode == "sim" else "okx_live"

    positions = out.get("positions") if isinstance(out.get("positions"), list) else []
    slim_pos = []
    for p in positions:
        if not isinstance(p, dict):
            continue
        slim_pos.append(
            {
                "inst_id": p.get("inst_id"),
                "opt_type": p.get("opt_type"),
                "sheets": p.get("sheets") or p.get("pos"),
                "pnl_usdt": p.get("display_pnl_usdt")
                if p.get("display_pnl_usdt") is not None
                else p.get("net_pnl_usdt"),
                "mark": p.get("mark") or p.get("mark_px"),
                "underlying": p.get("underlying"),
                "source_label": p.get("source_label"),
            }
        )

    out["instance"] = {
        "label": label,
        "mode": mode,
        "options_enabled": bool(out.get("enabled")),
        "account_label": out.get("account_label") or "",
        "margin_mode": out.get("options_margin_mode") or "coin",
        "underly": out.get("options_underly") or "ETH",
    }
    out["funds"] = {
        "funding_usdt": out.get("funding_usdt"),
        "trading_usdt": out.get("trading_usdt"),
        "funding_usdc": out.get("funding_usdc"),
        "trading_usdc": out.get("trading_usdc"),
    }
    out["pnl"] = {
        "unrealized_usdt": out.get("upl_total_usdc"),
    }
    out["positions_brief"] = slim_pos
    try:
        from lib.options.options_time_lib import options_app_now_str

        out["ts"] = options_app_now_str()
    except Exception:
        from datetime import datetime

        out["ts"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    return out


def register_hub_bridge_routes(app: Flask, app_module: Any) -> None:
    @app.route("/api/hub/options/snapshot")
    def api_hub_options_snapshot():
        if not hub_bridge_token_ok(request.headers):
            return jsonify({"ok": False, "msg": "未授权", "version": 1}), 401
        hub_ctx = app.config.get("HUB_CTX") or {}
        fn = hub_ctx.get("options_snapshot_fn")
        if not callable(fn):
            return jsonify(
                {"ok": False, "msg": "期权快照未注册", "version": 1, "enabled": False}
            ), 503
        try:
            snap = fn() or {}
        except Exception as e:
            return jsonify({"ok": False, "msg": str(e), "version": 1, "enabled": True}), 500
        return jsonify(_enrich_hub_protocol(snap, app_module=app_module))

    @app.route("/hub_sso")
    def hub_sso_login():
        token = (request.args.get("token") or "").strip()
        next_path = _safe_next_path(request.args.get("next") or "/embed")
        if not verify_sso_token(token):
            return (
                "<!doctype html><meta charset=utf-8><title>SSO 失败</title>"
                "<p>中控登录链无效或已过期,请从中控重新打开.</p>",
                403,
            )
        session["logged_in"] = True
        session.permanent = True
        session.modified = True
        # iframe 跨站时常丢 Cookie:附带可放入 sessionStorage 的 embed_sso
        embed_tok = make_embed_session_token(hub_bridge_token(), ttl_sec=604800)
        joiner = "&" if "?" in next_path else "?"
        return redirect(f"{next_path}{joiner}embed_sso={quote(embed_tok, safe='')}")

    @app.after_request
    def _hub_frame_ancestors(resp):
        ancestors = frame_ancestors_header_value()
        if ancestors:
            resp.headers["Content-Security-Policy"] = f"frame-ancestors {ancestors}"
            resp.headers.pop("X-Frame-Options", None)
        return resp