Added class to hide variables in tracebacks

src/snowflake/cli/_app/secret.py

@@ -0,0 +1,9 @@
+class SecretType:
+    def __init__(self, value):
+        self.value = value
+
+    def __repr__(self):
+        return "SecretType(***)"
+
+    def __str___(self):
+        return "***"

src/snowflake/cli/_app/snow_connector.py

@@ -24,6 +24,7 @@
 from snowflake.cli._app.constants import (
     PARAM_APPLICATION_NAME,
 )
+from snowflake.cli._app.secret import SecretType
 from snowflake.cli._app.telemetry import command_info
 from snowflake.cli.api.config import (
     get_connection_dict,
@@ -205,7 +206,7 @@ def _load_private_key(connection_parameters: Dict, private_key_var_name: str) ->
             connection_parameters[private_key_var_name]
         )
         private_key = _load_pem_to_der(private_key_pem)
-        connection_parameters["private_key"] = private_key
+        connection_parameters["private_key"] = private_key.value
         del connection_parameters[private_key_var_name]
     else:
         raise ClickException(
@@ -217,10 +218,11 @@ def _load_private_key_from_parameters(
     connection_parameters: Dict, private_key_var_name: str
 ) -> None:
     if connection_parameters.get("authenticator") == "SNOWFLAKE_JWT":
-        private_key_pem = connection_parameters[private_key_var_name]
-        private_key_pem = private_key_pem.encode("utf-8")
+        private_key_pem = _load_pem_from_parameters(
+            connection_parameters[private_key_var_name]
+        )
         private_key = _load_pem_to_der(private_key_pem)
-        connection_parameters["private_key"] = private_key
+        connection_parameters["private_key"] = private_key.value
         del connection_parameters[private_key_var_name]
     else:
         raise ClickException(
@@ -236,43 +238,49 @@ def _update_connection_application_name(connection_parameters: Dict):
     connection_parameters.update(connection_application_params)
 
 
-def _load_pem_from_file(private_key_file: str) -> bytes:
+def _load_pem_from_file(private_key_file: str) -> SecretType:
     with SecurePath(private_key_file).open(
         "rb", read_file_limit_mb=DEFAULT_SIZE_LIMIT_MB
     ) as f:
-        private_key_pem = f.read()
+        private_key_pem = SecretType(f.read())
     return private_key_pem
 
 
-def _load_pem_to_der(private_key_pem: bytes) -> bytes:
+def _load_pem_from_parameters(private_key_raw: str) -> SecretType:
+    return SecretType(private_key_raw.encode("utf-8"))
+
+
+def _load_pem_to_der(private_key_pem: SecretType) -> SecretType:
     """
     Given a private key file path (in PEM format), decode key data into DER
     format
     """
-    private_key_passphrase = os.getenv("PRIVATE_KEY_PASSPHRASE", None)
+    private_key_passphrase = SecretType(os.getenv("PRIVATE_KEY_PASSPHRASE", None))
     if (
-        private_key_pem.startswith(ENCRYPTED_PKCS8_PK_HEADER)
-        and private_key_passphrase is None
+        private_key_pem.value.startswith(ENCRYPTED_PKCS8_PK_HEADER)
+        and private_key_passphrase.value is None
     ):
         raise ClickException(
             "Encrypted private key, you must provide the"
             "passphrase in the environment variable PRIVATE_KEY_PASSPHRASE"
         )
 
-    if not private_key_pem.startswith(
+    if not private_key_pem.value.startswith(
         ENCRYPTED_PKCS8_PK_HEADER
-    ) and not private_key_pem.startswith(UNENCRYPTED_PKCS8_PK_HEADER):
+    ) and not private_key_pem.value.startswith(UNENCRYPTED_PKCS8_PK_HEADER):
         raise ClickException(
             "Private key provided is not in PKCS#8 format. Please use correct format."
         )
 
-    if private_key_pem.startswith(UNENCRYPTED_PKCS8_PK_HEADER):
-        private_key_passphrase = None
+    if private_key_pem.value.startswith(UNENCRYPTED_PKCS8_PK_HEADER):
+        private_key_passphrase = SecretType(None)
 
     return prepare_private_key(private_key_pem, private_key_passphrase)
 
 
-def prepare_private_key(private_key_pem, private_key_passphrase=None):
+def prepare_private_key(
+    private_key_pem: SecretType, private_key_passphrase: SecretType = SecretType(None)
+):
     from cryptography.hazmat.backends import default_backend
     from cryptography.hazmat.primitives.serialization import (
         Encoding,
@@ -281,17 +289,21 @@ def prepare_private_key(private_key_pem, private_key_passphrase=None):
         load_pem_private_key,
     )
 
-    private_key = load_pem_private_key(
-        private_key_pem,
-        (
-            str.encode(private_key_passphrase)
-            if private_key_passphrase is not None
-            else private_key_passphrase
-        ),
-        default_backend(),
+    private_key = SecretType(
+        load_pem_private_key(
+            private_key_pem.value,
+            (
+                str.encode(private_key_passphrase.value)
+                if private_key_passphrase.value is not None
+                else private_key_passphrase.value
+            ),
+            default_backend(),
+        )
     )
-    return private_key.private_bytes(
-        encoding=Encoding.DER,
-        format=PrivateFormat.PKCS8,
-        encryption_algorithm=NoEncryption(),
+    return SecretType(
+        private_key.value.private_bytes(
+            encoding=Encoding.DER,
+            format=PrivateFormat.PKCS8,
+            encryption_algorithm=NoEncryption(),
+        )
     )

tests/test_connection.py

@@ -21,6 +21,7 @@
 
 import pytest
 import tomlkit
+from snowflake.cli._app.secret import SecretType
 from snowflake.cli.api.constants import ObjectType
 
 from tests_common import IS_WINDOWS
@@ -705,7 +706,7 @@ def test_key_pair_authentication_from_config(
 ):
     ctx = mock_ctx()
     mock_connector.return_value = ctx
-    mock_convert.return_value = "secret value"
+    mock_convert.return_value = SecretType("secret value")
 
     with NamedTemporaryFile("w+", suffix="toml") as tmp_file:
         tmp_file.write(

tests/test_snow_connector.py

@@ -16,6 +16,7 @@
 from unittest import mock
 
 import pytest
+from snowflake.cli._app.secret import SecretType
 
 
 # Used as a solution to syrupy having some problems with comparing multilines string
@@ -118,9 +119,10 @@ def test_private_key_loading_and_aliases(
         else:
             overrides[user_input] = override_value
 
+    key = SecretType(b"bytes")
     mock_command_info.return_value = "SNOWCLI.SQL"
-    mock_load_pem_from_file.return_value = b"bytes"
-    mock_load_pem_to_der.return_value = b"bytes"
+    mock_load_pem_from_file.return_value = key
+    mock_load_pem_to_der.return_value = key
 
     conn_dict = get_connection_dict(connection_name)
     default_value = conn_dict.get("private_key_file", None) or conn_dict.get(
@@ -135,7 +137,7 @@ def test_private_key_loading_and_aliases(
         expected_private_key_args = (
             {}
             if expected_private_key_file_value is None
-            else dict(private_key=mock_load_pem_to_der.return_value)
+            else dict(private_key=b"bytes")
         )
         mock_connect.assert_called_once_with(
             application=mock_command_info.return_value,
@@ -145,7 +147,7 @@ def test_private_key_loading_and_aliases(
         )
         if expected_private_key_file_value is not None:
             mock_load_pem_from_file.assert_called_with(expected_private_key_file_value)
-            mock_load_pem_to_der.assert_called_with(b"bytes")
+            mock_load_pem_to_der.assert_called_with(key)
 
 
 @mock.patch.dict(os.environ, {}, clear=True)

tests_integration/snowflake_connector.py

@@ -22,6 +22,7 @@
 
 import pytest
 from snowflake import connector
+from snowflake.cli._app.secret import SecretType
 from snowflake.cli.api.exceptions import EnvironmentVariableNotFoundError
 from snowflake.cli._app.snow_connector import update_connection_details_with_private_key