xiaobai-review/xiaobai-datahub/tests/test_admin.py

from __future__ import annotations

import io
import json
import logging
import tempfile
import threading
import unittest
from http.server import ThreadingHTTPServer
from pathlib import Path
from urllib.error import HTTPError
from urllib.request import Request, urlopen

from datahub.adapters.tushare import TushareAdapter
from datahub.crypto import SecretVault
from datahub.httpapp import make_handler
from datahub.hub import Hub
from datahub.logutil import JsonFormatter
from datahub.settings import Settings
from tests.fixtures import fake_transport


class AdminTests(unittest.TestCase):
    def setUp(self) -> None:
        self.tmp = tempfile.TemporaryDirectory()
        settings = Settings(
            encryption_key=SecretVault.generate_key(),
            api_token="z" * 32,
            admin_password="StartPass1",
            tushare_token="real-tushare-token-abcdef",
            db_path=Path(self.tmp.name) / "hub.db",
            scheduler_enabled=False,
        )
        self.hub = Hub(settings, adapter=TushareAdapter("real-tushare-token-abcdef", transport=fake_transport))
        handler = make_handler(self.hub)
        self.server = ThreadingHTTPServer(("127.0.0.1", 0), handler)
        threading.Thread(target=self.server.serve_forever, daemon=True).start()
        self.base = f"http://127.0.0.1:{self.server.server_address[1]}"

    def tearDown(self) -> None:
        self.server.shutdown()
        self.server.server_close()
        self.tmp.cleanup()

    def _json(self, path, method="GET", body=None, cookie="", csrf=""):
        data = None if body is None else json.dumps(body).encode()
        headers = {"Content-Type": "application/json"}
        if cookie:
            headers["Cookie"] = cookie
        if csrf:
            headers["X-CSRF-Token"] = csrf
        req = Request(self.base + path, data=data, headers=headers, method=method)
        with urlopen(req, timeout=5) as resp:
            set_cookie = resp.headers.get("Set-Cookie", "")
            return resp.status, json.loads(resp.read().decode()), set_cookie

    def test_login_change_password_and_secret_masking(self) -> None:
        status, body, cookie_header = self._json(
            "/admin/api/login", "POST", {"username": "hub_admin", "password": "StartPass1"}
        )
        self.assertEqual(status, 200)
        self.assertTrue(body["must_change"])
        cookie = cookie_header.split(";")[0]
        csrf = body["csrf"]
        status, _, _ = self._json(
            "/admin/api/change-password",
            "POST",
            {"current": "StartPass1", "new_password": "NewPass123"},
            cookie=cookie,
            csrf=csrf,
        )
        self.assertEqual(status, 200)
        _, sources, _ = self._json("/admin/api/sources", cookie=cookie, csrf=csrf)
        blob = json.dumps(sources)
        self.assertNotIn("real-tushare-token-abcdef", blob)
        self.assertTrue(sources["items"][0]["credential"]["configured"])
        self.assertTrue(str(sources["items"][0]["credential"]["last4"]).endswith("cdef") or "****" in str(sources["items"][0]["credential"]["last4"]))

    def test_rollback_requires_password_and_confirm(self) -> None:
        _, body, cookie_header = self._json(
            "/admin/api/login", "POST", {"username": "hub_admin", "password": "StartPass1"}
        )
        cookie = cookie_header.split(";")[0]
        csrf = body["csrf"]
        self._json("/admin/api/change-password", "POST", {"current": "StartPass1", "new_password": "NewPass123"}, cookie, csrf)
        from urllib.error import HTTPError

        with self.assertRaises(HTTPError) as ctx:
            self._json(
                "/admin/api/rollback",
                "POST",
                {"dataset": "daily", "trade_date": "20240902", "password": "wrong", "confirm": "daily:20240902"},
                cookie,
                csrf,
            )
        self.assertEqual(ctx.exception.code, 401)

    def test_invalid_json_does_not_log_request_body_secrets(self) -> None:
        secret = "SuperSecretPass1!"
        token = "hub-token-should-not-leak"
        raw = json.dumps({"password": secret, "token": token, "username": "hub_admin"}) + "{not-json"
        stream = io.StringIO()
        logger = logging.getLogger("datahub")
        handler = logging.StreamHandler(stream)
        handler.setFormatter(JsonFormatter())
        logger.addHandler(handler)
        previous_level = logger.level
        logger.setLevel(logging.DEBUG)
        try:
            req = Request(
                self.base + "/admin/api/login",
                data=raw.encode("utf-8"),
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with self.assertRaises(HTTPError) as ctx:
                urlopen(req, timeout=5)
            body = ctx.exception.read().decode("utf-8")
            self.assertEqual(ctx.exception.code, 400)
            self.assertNotIn(secret, body)
            self.assertNotIn(token, body)
            blob = stream.getvalue() + body
            self.assertNotIn(secret, blob)
            self.assertNotIn(token, blob)
            self.assertNotIn(raw, blob)
        finally:
            logger.removeHandler(handler)
            logger.setLevel(previous_level)

    def test_json_formatter_drops_decode_error_document(self) -> None:
        secret = "ParseSecretTokenXYZ"
        formatter = JsonFormatter()
        logger = logging.getLogger("datahub.test")
        record = logger.makeRecord(
            "datahub.test", logging.ERROR, __file__, 1, "parse failed", (), None
        )
        try:
            json.loads('{"password": "%s"}{' % secret)
        except json.JSONDecodeError as exc:
            record.exc_info = (type(exc), exc, exc.__traceback__)
        blob = formatter.format(record)
        self.assertNotIn(secret, blob)
        self.assertIn("invalid json", blob)


if __name__ == "__main__":
    unittest.main()