From b2c456868bc6cf53d70e1da14763cfbaaf6ca4b8 Mon Sep 17 00:00:00 2001
From: Valerie Becker <vbecker@lsst.org>
Date: Fri, 6 Sep 2024 15:23:22 +0000
Subject: [PATCH] Added clean_url to ConsDbClient's requests.Session response
 hooks

---
 python/lsst/summit/utils/consdbClient.py | 22 ++++++++++++-
 tests/test_consdbClient.py               | 41 ++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/python/lsst/summit/utils/consdbClient.py b/python/lsst/summit/utils/consdbClient.py
index e723ac72..76d9ad9a 100644
--- a/python/lsst/summit/utils/consdbClient.py
+++ b/python/lsst/summit/utils/consdbClient.py
@@ -23,7 +23,7 @@
 import os
 from dataclasses import dataclass
 from typing import Any
-from urllib.parse import quote
+from urllib.parse import quote, urlparse
 
 import requests
 from astropy.table import Table
@@ -77,6 +77,24 @@ def _check_status(r: requests.Response) -> None:
         raise e
 
 
+def clean_url(resp: requests.Response, *args, **kwargs) -> requests.Response:
+    """Parse url from response and remove netloc portion.
+
+    Set new url in response and return response
+
+    Parameters
+    ----------
+    resp : `requests.Response`
+        The response that could contain a URL with tokens
+    """
+    url = urlparse(resp.url)
+    short_user = f"{url.username[:2]}***" if url.username is not None else ""
+    short_pass = f":{url.password[:2]}***" if url.password is not None else ""
+    netloc = f"{short_user}{short_pass}@{url.hostname}"
+    resp.url = url._replace(netloc=netloc).geturl()
+    return resp
+
+
 @dataclass
 class FlexibleMetadataInfo:
     """Description of a flexible metadata value.
@@ -127,6 +145,8 @@ class ConsDbClient:
 
     def __init__(self, url: str | None = None):
         self.session = requests.Session()
+        self.session.hooks["response"].append(clean_url)
+
         if url is None:
             self.url = os.environ["LSST_CONSDB_PQ_URL"]
         else:
diff --git a/tests/test_consdbClient.py b/tests/test_consdbClient.py
index b891d8b3..0b8c0289 100644
--- a/tests/test_consdbClient.py
+++ b/tests/test_consdbClient.py
@@ -28,6 +28,9 @@
 
 @pytest.fixture
 def client():
+    """Initialize client with a fake url
+    Requires mocking connection with @responses.activate decorator
+    """
     return ConsDbClient("http://example.com/consdb")
 
 
@@ -161,6 +164,44 @@ def test_schema(client):
     assert client.schema(instrument, table) == description
 
 
+@responses.activate
+@pytest.mark.parametrize(
+    "secret, redacted",
+    [
+        ("usdf:v987wefVMPz", "us***:v9***"),
+        ("u:v", "u***:v***"),
+        ("ulysses", "ul***"),
+        (":alberta94", "***:al***"),
+    ],
+)
+def test_clean_token_url_response(secret, redacted):
+    """Test tokens URL is cleaned when an error is thrown from requests
+    Use with pytest raises assert an error'
+    assert that url does not contain tokens
+    """
+    domain = "@usdf-fake.slackers.stanford.edu/consdb"
+    complex_client = ConsDbClient(f"https://{secret}{domain}")
+
+    obs_type = "exposure"
+    responses.post(
+        f"https://{secret}{domain}/flex/bad_instrument/exposure/addkey",
+        status=404,
+    )
+    with pytest.raises(HTTPError, match="404") as error:
+        complex_client.add_flexible_metadata_key(
+            "bad_instrument", obs_type, "error", "int", "instrument error"
+        )
+
+    url = error.value.args[0].split()[-1]
+    sanitized = f"https://{redacted}{domain}/flex/bad_instrument/exposure/addkey"
+    assert url == sanitized
+
+
+def test_client(client):
+    """Test ConsDbClient is initialized properly"""
+    assert "clean_url" in str(client.session.hooks["response"])
+
+
 # TODO: more POST tests
 #    client.insert(instrument, table, obs_id, values, allow_update)
 #    client.insert_multiple(instrument, table, obs_dict, allow_update)