from __future__ import annotations import json import mimetypes from http import HTTPStatus from http.cookies import SimpleCookie from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Any from urllib.parse import unquote, urlparse from datahub.hub import Hub from datahub.logutil import configure_logging, get_logger from datahub.serving import ApiError, parse_query from datahub.settings import DEFAULT_REVIEW_PUBLIC_URL from datahub.siteauth import SITE_SESSION_COOKIE, SiteBridgeError LOGGER = get_logger() class HubRequestHandler(BaseHTTPRequestHandler): hub: Hub def log_message(self, format: str, *args: Any) -> None: LOGGER.info(format % args) def do_GET(self) -> None: # noqa: N802 self._dispatch("GET") def do_POST(self) -> None: # noqa: N802 self._dispatch("POST") def do_OPTIONS(self) -> None: # noqa: N802 self.send_response(HTTPStatus.NO_CONTENT) self.send_header("Allow", "GET, POST, OPTIONS") self.end_headers() def _dispatch(self, method: str) -> None: parsed = urlparse(self.path) path = unquote(parsed.path) try: if path in {"/livez", "/healthz"}: self._json({"status": "ok"}, HTTPStatus.OK) return if path.startswith("/v1/"): self._v1(path, parsed.query, method) return if path.startswith("/admin/api/"): self._admin_api(method, path) return if path.startswith("/admin"): self._admin_static(path) return if path == "/": self.send_response(HTTPStatus.FOUND) self.send_header("Location", "/admin/") self.end_headers() return self._json({"error": {"code": "INVALID_ARGUMENT", "message": "Not found"}}, HTTPStatus.NOT_FOUND) except ApiError as exc: self._json(exc.payload(), exc.status) except PermissionError as exc: self._json({"error": {"code": "UNAUTHORIZED", "message": str(exc)}}, HTTPStatus.UNAUTHORIZED) except ValueError as exc: self._json({"error": {"code": "INVALID_ARGUMENT", "message": str(exc)}}, HTTPStatus.BAD_REQUEST) except Exception: LOGGER.exception("internal error") self._json({"error": {"code": "INTERNAL", "message": "internal error"}}, HTTPStatus.INTERNAL_SERVER_ERROR) def _v1(self, path: str, query: str, method: str = "GET") -> None: token = self.headers.get("X-Datahub-Token", "") if not self.hub.auth.check_api_token(token): self.hub.pipeline.audit("anonymous", "unauthorized", path, "") raise ApiError("UNAUTHORIZED", "missing or invalid X-Datahub-Token") if path == "/v1/query" and method == "POST": body = self._read_json(max_bytes=1_000_000) payload = self.hub.api.query_api(body) self._json(payload, HTTPStatus.OK) return if path == "/v1/credentials/ifind" and method == "POST": body = self._read_json() payload = self.hub.put_ifind_credentials(body) self._json(payload, HTTPStatus.OK) return payload = self.hub.api.handle(path, parse_query(query)) self._json(payload, HTTPStatus.OK) def _admin_api(self, method: str, path: str) -> None: session_token = self._cookie_value(SITE_SESSION_COOKIE) user = self._site_user(session_token) if user is None: if path == "/admin/api/session" and method == "GET": self._json( {"authenticated": False, "login_url": self._login_url()}, HTTPStatus.UNAUTHORIZED, ) return raise ApiError("UNAUTHORIZED", "请先在小白复盘主站登录") if not user["is_admin"]: self._json( { "error": {"code": "PERMISSION_DENIED", "message": "数据中枢仅管理员可进入"}, "authenticated": True, "is_admin": False, "username": user["username"], }, HTTPStatus.FORBIDDEN, ) return if method == "POST" and not self.hub.site_auth.check_csrf(session_token, self.headers.get("X-CSRF-Token", "")): raise ApiError("UNAUTHORIZED", "CSRF 校验失败") if path == "/admin/api/session" and method == "GET": self._json( { "authenticated": True, "is_admin": True, "username": user["username"], "csrf": self.hub.site_auth.csrf_token(session_token), "review_url": self._review_url(), }, HTTPStatus.OK, ) return if path == "/admin/api/logout" and method == "POST": self.hub.site_auth.logout(session_token) self.hub.pipeline.audit(user["username"], "logout", "site_session", "") self._json({"ok": True, "login_url": self._login_url()}, HTTPStatus.OK) return if self._console_api(method, path, user): return if path == "/admin/api/overview" and method == "GET": self._json(self.hub.admin.overview(), HTTPStatus.OK) return if path == "/admin/api/sources" and method == "GET": self._json(self.hub.admin.sources(), HTTPStatus.OK) return if path == "/admin/api/providers/status" and method == "GET": query = parse_query(urlparse(self.path).query) provider = (query.get("provider") or [""])[0] limit = (query.get("limit") or ["50"])[0] self._json(self.hub.admin.providers_status(provider, int(limit or 50)), HTTPStatus.OK) return if path == "/admin/api/source-catalog" and method == "GET": self._json(self.hub.admin.source_catalog(), HTTPStatus.OK) return if path == "/admin/api/lineage" and method == "GET": query = parse_query(urlparse(self.path).query) date = (query.get("date") or [""])[0] self._json(self.hub.admin.lineage(date), HTTPStatus.OK) return if path == "/admin/api/lineage/affected" and method == "GET": query = parse_query(urlparse(self.path).query) provider = (query.get("provider") or [""])[0] interface = (query.get("interface") or [""])[0] self._json(self.hub.admin.lineage_affected(provider, interface), HTTPStatus.OK) return if path.startswith("/admin/api/sources/") and path.endswith("/probe") and method == "POST": provider = path.split("/")[4] self._json(self.hub.admin.probe(provider), HTTPStatus.OK) return if path == "/admin/api/jobs" and method == "GET": self._json(self.hub.admin.jobs(), HTTPStatus.OK) return if path.startswith("/admin/api/jobs/") and path.endswith("/run") and method == "POST": job_id = path.split("/")[4] body = self._read_json(allow_empty=True) self._json(self.hub.admin.run_job(job_id, str(body.get("trade_date") or "")), HTTPStatus.OK) return if path == "/admin/api/batches" and method == "GET": query = parse_query(urlparse(self.path).query) date = (query.get("date") or [""])[0] dataset = (query.get("dataset") or [""])[0] self._json(self.hub.admin.batches(date, dataset), HTTPStatus.OK) return if path == "/admin/api/datasets" and method == "GET": query = parse_query(urlparse(self.path).query) self._json(self.hub.admin.datasets((query.get("date") or [""])[0]), HTTPStatus.OK) return if path == "/admin/api/audit" and method == "GET": self._json(self.hub.admin.audit(), HTTPStatus.OK) return if path == "/admin/api/rollback" and method == "POST": body = self._read_json() result = self.hub.admin.rollback( str(body.get("dataset") or ""), str(body.get("trade_date") or ""), str(body.get("password") or ""), str(body.get("confirm") or ""), user["username"], int(user["id"]), ) self._json(result, HTTPStatus.OK) return if path == "/admin/api/backfill" and method == "POST": body = self._read_json() result = self.hub.admin.backfill( str(body.get("dataset") or ""), str(body.get("trade_date") or ""), str(body.get("password") or ""), str(body.get("confirm") or ""), user["username"], int(user["id"]), ) self._json(result, HTTPStatus.OK) return raise ApiError("INVALID_ARGUMENT", f"unknown admin endpoint: {path}") def _site_user(self, session_token: str) -> dict[str, Any] | None: try: return self.hub.site_auth.verify(session_token) except SiteBridgeError as exc: raise ApiError("UNAVAILABLE", str(exc)) from exc def _review_url(self) -> str: """Browser-reachable review site URL. HEL-564: fixed to REVIEW_PUBLIC_URL (default the intranet address) instead of being derived from the browser's Host header — a domain that only proxies 8765 would send the console back to a dead port. """ configured = self.hub.settings.review_public_url or DEFAULT_REVIEW_PUBLIC_URL return configured.rstrip("/") def _login_url(self) -> str: return f"{self._review_url()}/login/" def _console_api(self, method: str, path: str, user: dict[str, Any]) -> bool: """Endpoints backed by the review site: model pool, members, invites. Returns True when the request was handled so the caller can fall through to the hub-owned endpoints otherwise. """ console = self.hub.site_console if path == "/admin/api/models" and method == "GET": self._json(console.models(), HTTPStatus.OK) elif path == "/admin/api/models/save" and method == "POST": self._json(console.save_models(self._read_json()), HTTPStatus.OK) elif path == "/admin/api/models/test" and method == "POST": self._json(console.test_model(self._read_json()), HTTPStatus.OK) elif path == "/admin/api/models/fetch" and method == "POST": self._json(console.fetch_models(self._read_json()), HTTPStatus.OK) elif path == "/admin/api/members" and method == "GET": self._json(console.members(), HTTPStatus.OK) elif path == "/admin/api/members/save" and method == "POST": self._json(console.save_member(self._read_json()), HTTPStatus.OK) elif path == "/admin/api/members/quota" and method == "POST": self._json(console.save_quota(self._read_json()), HTTPStatus.OK) elif path == "/admin/api/invites" and method == "GET": self._json(console.invites(), HTTPStatus.OK) elif path == "/admin/api/invites/create" and method == "POST": self._json(console.create_invites(self._read_json(allow_empty=True), int(user["id"])), HTTPStatus.OK) elif path == "/admin/api/invites/revoke" and method == "POST": self._json(console.revoke_invite(self._read_json()), HTTPStatus.OK) elif path == "/admin/api/credentials/tushare" and method == "POST": payload = self.hub.put_tushare_credentials(self._read_json()) self.hub.pipeline.audit(user["username"], "store_credential", "tushare_token", "") self._json(payload, HTTPStatus.OK) elif path == "/admin/api/credentials/ifind" and method == "POST": payload = self.hub.put_ifind_credentials(self._read_json()) self.hub.pipeline.audit(user["username"], "store_credential", "ifind_tokens", "") self._json(payload, HTTPStatus.OK) else: return False return True def _admin_static(self, path: str) -> None: relative = path[len("/admin"):].lstrip("/") or "index.html" candidate = (self.hub.static_dir / relative).resolve() try: candidate.relative_to(self.hub.static_dir.resolve()) except ValueError: self.send_error(HTTPStatus.FORBIDDEN) return if candidate.is_dir(): candidate = candidate / "index.html" if not candidate.is_file(): candidate = self.hub.static_dir / "index.html" content = candidate.read_bytes() content_type = mimetypes.guess_type(candidate.name)[0] or "application/octet-stream" if content_type.startswith("text/") or content_type in {"application/javascript", "application/json"}: content_type += "; charset=utf-8" self.send_response(HTTPStatus.OK) self.send_header("Content-Type", content_type) self.send_header("Content-Length", str(len(content))) self.send_header("Cache-Control", "no-cache") self.end_headers() self.wfile.write(content) def _read_json(self, allow_empty: bool = False, max_bytes: int = 65536) -> dict[str, Any]: length = int(self.headers.get("Content-Length", "0") or 0) if length == 0 and allow_empty: return {} if length <= 0 or length > max_bytes: raise ValueError("请求内容为空或过大") raw = self.rfile.read(length) try: payload = json.loads(raw.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError): LOGGER.warning("invalid json request body") raise ValueError("请求不是合法 JSON") from None if not isinstance(payload, dict): raise ValueError("请求不是合法 JSON") return payload def _cookie_value(self, name: str) -> str: cookie = SimpleCookie() try: cookie.load(self.headers.get("Cookie", "")) except Exception: return "" morsel = cookie.get(name) return morsel.value if morsel else "" def _json(self, payload: dict[str, Any], status: HTTPStatus, extra_headers: list[str] | None = None) -> None: raw = json.dumps(payload, ensure_ascii=False).encode("utf-8") self.send_response(status) self.send_header("Content-Type", "application/json; charset=utf-8") self.send_header("Content-Length", str(len(raw))) self.send_header("Cache-Control", "no-store") for header in extra_headers or []: self.send_header("Set-Cookie", header) self.end_headers() self.wfile.write(raw) def make_handler(hub: Hub) -> type[HubRequestHandler]: class BoundHandler(HubRequestHandler): pass BoundHandler.hub = hub BoundHandler.protocol_version = "HTTP/1.1" return BoundHandler def serve(hub: Hub, host: str, port: int) -> None: configure_logging(hub.settings.log_level) handler = make_handler(hub) server = ThreadingHTTPServer((host, port), handler) hub.start() LOGGER.info("xiaobai-datahub listening", extra={"hub": {"host": host, "port": port}}) print(f"xiaobai-datahub is running at http://{host}:{port}/admin/") try: server.serve_forever() except KeyboardInterrupt: pass finally: hub.stop() server.server_close()