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()