Redact AWS session token from client repr

Author: NguyenCong2k

pymongo/asynchronous/mongo_client.py

@@ -1300,13 +1300,24 @@ def __hash__(self) -> int:
         return hash(self.eq_props())
 
     def _repr_helper(self) -> str:
+        def redact_auth_mechanism_properties(value: Any) -> Any:
+            if isinstance(value, dict):
+                redacted = value.copy()
+                for key in redacted:
+                    if key.upper() == "AWS_SESSION_TOKEN":
+                        redacted[key] = "<redacted>"
+                return redacted
+            return value
+
         def option_repr(option: str, value: Any) -> str:
             """Fix options whose __repr__ isn't usable in a constructor."""
             if option == "document_class":
                 if value is dict:
                     return "document_class=dict"
                 else:
                     return f"document_class={value.__module__}.{value.__name__}"
+            if option == "authmechanismproperties":
+                value = redact_auth_mechanism_properties(value)
             if option in common.TIMEOUT_OPTIONS and value is not None:
                 return f"{option}={int(value * 1000)}"
 

pymongo/synchronous/mongo_client.py

@@ -1300,13 +1300,24 @@ def __hash__(self) -> int:
         return hash(self.eq_props())
 
     def _repr_helper(self) -> str:
+        def redact_auth_mechanism_properties(value: Any) -> Any:
+            if isinstance(value, dict):
+                redacted = value.copy()
+                for key in redacted:
+                    if key.upper() == "AWS_SESSION_TOKEN":
+                        redacted[key] = "<redacted>"
+                return redacted
+            return value
+
         def option_repr(option: str, value: Any) -> str:
             """Fix options whose __repr__ isn't usable in a constructor."""
             if option == "document_class":
                 if value is dict:
                     return "document_class=dict"
                 else:
                     return f"document_class={value.__module__}.{value.__name__}"
+            if option == "authmechanismproperties":
+                value = redact_auth_mechanism_properties(value)
             if option in common.TIMEOUT_OPTIONS and value is not None:
                 return f"{option}={int(value * 1000)}"
 

test/asynchronous/test_client.py

@@ -195,6 +195,19 @@ def test_types(self):
 
         self.assertRaises(ConfigurationError, AsyncMongoClient, [])
 
+    async def test_repr_redacts_aws_session_token(self):
+        token = "SECRET_AWS_SESSION_TOKEN"
+        client = AsyncMongoClient(
+            "mongodb://AKIA:SECRET@localhost:27017/"
+            f"?authMechanism=MONGODB-AWS&authMechanismProperties=AWS_SESSION_TOKEN:{token}",
+            connect=False,
+        )
+
+        the_repr = repr(client)
+
+        self.assertNotIn(token, the_repr)
+        self.assertIn("'AWS_SESSION_TOKEN': '<redacted>'", the_repr)
+
     async def test_max_pool_size_zero(self):
         self.simple_client(maxPoolSize=0)
 

test/test_client.py

@@ -192,6 +192,19 @@ def test_types(self):
 
         self.assertRaises(ConfigurationError, MongoClient, [])
 
+    def test_repr_redacts_aws_session_token(self):
+        token = "SECRET_AWS_SESSION_TOKEN"
+        client = MongoClient(
+            "mongodb://AKIA:SECRET@localhost:27017/"
+            f"?authMechanism=MONGODB-AWS&authMechanismProperties=AWS_SESSION_TOKEN:{token}",
+            connect=False,
+        )
+
+        the_repr = repr(client)
+
+        self.assertNotIn(token, the_repr)
+        self.assertIn("'AWS_SESSION_TOKEN': '<redacted>'", the_repr)
+
     def test_max_pool_size_zero(self):
         self.simple_client(maxPoolSize=0)