From c4302f4e6e25408a6bc15c0fb17797e9cc74a465 Mon Sep 17 00:00:00 2001
From: Afroz Alam <afroz.alam@snowflake.com>
Date: Fri, 3 Mar 2023 09:03:41 -0800
Subject: [PATCH] Mask JWT tokens in logs in case of post-auth errors (#1457)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Mask JWT tokens in logs in case of post-auth errors

The JWT tokens are valid for 1 minute, but still can be enough to steal the tokens from logs (if a hacker has access to the log stream or the log service) and use the token to access the data, potentially to change the user's credentials — if done in an automated way.

* use secret detector and add tests

---------

Co-authored-by: Sergey Vasilyev <sv@datafold.com>
---
 src/snowflake/connector/network.py | 12 ++++------
 test/unit/test_retry_network.py    | 38 ++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/src/snowflake/connector/network.py b/src/snowflake/connector/network.py
index e68145792..6720d2373 100644
--- a/src/snowflake/connector/network.py
+++ b/src/snowflake/connector/network.py
@@ -20,6 +20,7 @@
 
 import OpenSSL.SSL
 
+from snowflake.connector.secret_detector import SecretDetector
 from snowflake.connector.vendored.requests.models import PreparedRequest
 from snowflake.connector.vendored.urllib3.connectionpool import (
     HTTPConnectionPool,
@@ -980,14 +981,9 @@ def handle_invalid_certificate_error(self, conn, full_url, cause) -> None:
     def _handle_unknown_error(self, method, full_url, headers, data, conn) -> None:
         """Handles unknown errors."""
         if data:
-            try:
-                decoded_data = json.loads(data)
-                if decoded_data.get("data") and decoded_data["data"].get("PASSWORD"):
-                    # masking the password
-                    decoded_data["data"]["PASSWORD"] = "********"
-                    data = json.dumps(decoded_data)
-            except Exception:
-                logger.info("data is not JSON")
+            _, masked_data, err_str = SecretDetector.mask_secrets(data)
+            if err_str is None:
+                data = masked_data
         logger.error(
             f"Failed to get the response. Hanging? "
             f"method: {method}, url: {full_url}, headers:{headers}, "
diff --git a/test/unit/test_retry_network.py b/test/unit/test_retry_network.py
index ba983bdf6..7f7e69b52 100644
--- a/test/unit/test_retry_network.py
+++ b/test/unit/test_retry_network.py
@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import errno
+import logging
 import os
 import time
 from unittest.mock import MagicMock, Mock, PropertyMock
@@ -217,3 +218,40 @@ def fake_request_exec(**kwargs):
     assert ret == {}
     assert cnt.c == 1  # failed on first call - did not retry
     assert rest._connection.errorhandler.called  # error
+
+
+def test_secret_masking(caplog):
+    connection = MagicMock()
+    connection.errorhandler = Mock(return_value=None)
+
+    rest = SnowflakeRestful(
+        host="testaccount.snowflakecomputing.com", port=443, connection=connection
+    )
+
+    data = (
+        '{"code": 12345,'
+        ' "data": {"TOKEN": "_Y1ZNETTn5/qfUWj3Jedb", "PASSWORD": "dummy_pass"}'
+        "}"
+    )
+    default_parameters = {
+        "method": "POST",
+        "full_url": "https://testaccount.snowflakecomputing.com/",
+        "headers": {},
+        "data": data,
+    }
+
+    class NotRetryableException(Exception):
+        pass
+
+    def fake_request_exec(**kwargs):
+        return None
+
+    # inject a fake method
+    rest._request_exec = fake_request_exec
+
+    # first two attempts will fail but third will success
+    with caplog.at_level(logging.ERROR):
+        ret = rest.fetch(timeout=10, **default_parameters)
+    assert '"TOKEN": "****' in caplog.text
+    assert '"PASSWORD": "****' in caplog.text
+    assert ret == {}