diff --git a/lark_oapi/ws/client.py b/lark_oapi/ws/client.py index 8ee99183..48002687 100644 --- a/lark_oapi/ws/client.py +++ b/lark_oapi/ws/client.py @@ -80,6 +80,42 @@ def _ws_connect_kwargs(): return {} +def _redact_conn_url_for_log(url: str) -> str: + query_start = url.find("?") + fragment_start = url.find("#") + if query_start == -1 or 0 <= fragment_start < query_start: + return url + + query_end = fragment_start + if query_end == -1: + query_end = len(url) + + redacted = [] + cursor = query_start + 1 + while cursor <= query_end: + param_end = url.find("&", cursor, query_end) + if param_end == -1: + param_end = query_end + + equals = url.find("=", cursor, param_end) + if equals != -1: + key = url[cursor:equals] + if key.isascii() and key.lower() in ("access_key", "ticket"): + redacted.append(url[cursor:equals + 1]) + redacted.append("***") + else: + redacted.append(url[cursor:param_end]) + else: + redacted.append(url[cursor:param_end]) + + if param_end == query_end: + break + redacted.append("&") + cursor = param_end + 1 + + return url[:query_start + 1] + "".join(redacted) + url[query_end:] + + def _get_ws_conn_exception_headers(e): headers = getattr(e, "headers", None) if headers is not None: @@ -207,7 +243,7 @@ async def _connect(self) -> None: self._conn_id = conn_id self._service_id = service_id - logger.info(self._fmt_log("connected to {}", conn_url)) + logger.info(self._fmt_log("connected to {}", _redact_conn_url_for_log(conn_url))) loop.create_task(self._receive_message_loop()) except InvalidHandshake as e: _parse_ws_conn_exception(e) @@ -413,7 +449,7 @@ async def _disconnect(self): if self._conn is None: return await self._conn.close() - logger.info(self._fmt_log("disconnected to {}", self._conn_url)) + logger.info(self._fmt_log("disconnected to {}", _redact_conn_url_for_log(self._conn_url))) finally: self._conn = None self._conn_url = "" diff --git a/lark_oapi/ws/tests/test_client_log_redaction.py b/lark_oapi/ws/tests/test_client_log_redaction.py new file mode 100644 index 00000000..ffb0c7e7 --- /dev/null +++ b/lark_oapi/ws/tests/test_client_log_redaction.py @@ -0,0 +1,61 @@ +import logging + +import pytest + +from lark_oapi.ws import client as ws_client + + +class _FakeConn: + async def close(self): + pass + + +def test_redaction_does_not_treat_fragment_text_as_query(): + url = "wss://example.test/callback#section?access_key=visible" + + assert ws_client._redact_conn_url_for_log(url) == url + + +@pytest.mark.asyncio +async def test_connection_lifecycle_logs_redact_only_sensitive_query_values(monkeypatch, caplog): + conn_url = ( + "wss://example.test/callback?device_id=device%20one&service_id=42" + "&access_key=fake%20access%2Fpart&TICKET=Fake%2fticket&ticket&flag&empty=" + "&ticket=second-fake&access_key=&access_key_extra=visible#section%20one" + ) + redacted_url = ( + "wss://example.test/callback?device_id=device%20one&service_id=42" + "&access_key=***&TICKET=***&ticket&flag&empty=" + "&ticket=***&access_key=***&access_key_extra=visible#section%20one" + ) + connected_urls = [] + + async def fake_connect(uri): + connected_urls.append(uri) + return _FakeConn() + + client = ws_client.Client("fake_app_id", "fake_app_secret") + monkeypatch.setattr(client, "_get_conn_url", lambda: conn_url) + monkeypatch.setattr(ws_client.websockets, "connect", fake_connect) + monkeypatch.setattr( + ws_client.loop, + "create_task", + lambda coro: coro.close() if hasattr(coro, "close") else None, + ) + + with caplog.at_level(logging.INFO, logger="Lark"): + await client._connect() + assert client._conn_url == conn_url + await client._disconnect() + + lifecycle_messages = [ + record.getMessage() + for record in caplog.records + if "connected to " in record.getMessage() or "disconnected to " in record.getMessage() + ] + + assert connected_urls == [conn_url] + assert lifecycle_messages == [ + "connected to {} [conn_id=device one]".format(redacted_url), + "disconnected to {} [conn_id=device one]".format(redacted_url), + ]