xiaobai-review/xiaobai-datahub/datahub/httpapp.py

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, 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:
        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 == "/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"],
            )
            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, 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 _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()