from __future__ import annotations import json import mimetypes import secrets 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 LOGGER = get_logger() SESSION_COOKIE = "datahub_session" 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) 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) -> 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") payload = self.hub.api.handle(path, parse_query(query)) self._json(payload, HTTPStatus.OK) def _admin_api(self, method: str, path: str) -> None: if path == "/admin/api/login" and method == "POST": body = self._read_json() result = self.hub.auth.login(str(body.get("username") or "hub_admin"), str(body.get("password") or "")) self._json( {"ok": True, "must_change": result["must_change"], "csrf": result["csrf"]}, HTTPStatus.OK, extra_headers=[self._cookie(result["session"])], ) return user = self.hub.auth.session_user(self._cookie_value(SESSION_COOKIE)) if not user: raise ApiError("UNAUTHORIZED", "请先登录") if method == "POST" and path != "/admin/api/login": csrf = self.headers.get("X-CSRF-Token", "") if not csrf or not secrets.compare_digest(csrf, str(user["csrf_token"])): raise ApiError("UNAUTHORIZED", "CSRF 校验失败") if path == "/admin/api/logout" and method == "POST": self.hub.auth.logout(self._cookie_value(SESSION_COOKIE)) self._json({"ok": True}, HTTPStatus.OK, extra_headers=[self._cookie("", clear=True)]) return if path == "/admin/api/session" and method == "GET": self._json({"username": user["username"], "must_change": user["must_change"], "csrf": user["csrf_token"]}, HTTPStatus.OK) return if path == "/admin/api/change-password" and method == "POST": body = self._read_json() self.hub.auth.change_password(str(body.get("current") or ""), str(body.get("new_password") or "")) self.hub.pipeline.audit(user["username"], "change_password", "hub_admin", "") self._json({"ok": True}, HTTPStatus.OK) return if user["must_change"] and path not in {"/admin/api/change-password", "/admin/api/session"}: raise ApiError("UNAUTHORIZED", "请先修改初始密码") 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.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"], ) 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"], ) self._json(result, HTTPStatus.OK) return raise ApiError("INVALID_ARGUMENT", f"unknown admin endpoint: {path}") 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) -> 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 > 65536: raise ValueError("请求内容为空或过大") return json.loads(self.rfile.read(length).decode("utf-8")) 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 _cookie(self, value: str, clear: bool = False) -> str: max_age = 0 if clear else 12 * 3600 return f"{SESSION_COOKIE}={value}; Path=/; HttpOnly; SameSite=Strict; Max-Age={max_age}" 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()