Backport #66523 to 24.6: Add additional log masking in CI

tests/ci/functional_test_check.py

@@ -24,6 +24,18 @@ from tee_popen import TeePopen
 NO_CHANGES_MSG = "Nothing to run"
 
 
+class SensitiveFormatter(logging.Formatter):
+    @staticmethod
+    def _filter(s):
+        return re.sub(
+            r"(.*)(AZURE_CONNECTION_STRING.*\')(.*)", r"\1AZURE_CONNECTION_STRING\3", s
+        )
+
+    def format(self, record):
+        original = logging.Formatter.format(self, record)
+        return self._filter(original)
+
+
 def get_additional_envs(
     check_name: str, run_by_hash_num: int, run_by_hash_total: int
 ) -> List[str]:
@@ -206,6 +218,9 @@ def parse_args():
 
 def main():
     logging.basicConfig(level=logging.INFO)
+    for handler in logging.root.handlers:
+        # pylint: disable=protected-access
+        handler.setFormatter(SensitiveFormatter(handler.formatter._fmt))  # type: ignore
 
     stopwatch = Stopwatch()
 

tests/ci/stress_check.py

@@ -3,6 +3,7 @@
 import csv
 import logging
 import os
+import re
 import subprocess
 import sys
 from pathlib import Path
@@ -19,6 +20,18 @@ from stopwatch import Stopwatch
 from tee_popen import TeePopen
 
 
+class SensitiveFormatter(logging.Formatter):
+    @staticmethod
+    def _filter(s):
+        return re.sub(
+            r"(.*)(AZURE_CONNECTION_STRING.*\')(.*)", r"\1AZURE_CONNECTION_STRING\3", s
+        )
+
+    def format(self, record):
+        original = logging.Formatter.format(self, record)
+        return self._filter(original)
+
+
 def get_additional_envs(check_name: str) -> List[str]:
     result = []
     azure_connection_string = get_parameter_from_ssm("azure_connection_string")
@@ -114,6 +127,9 @@ def process_results(
 
 def run_stress_test(docker_image_name: str) -> None:
     logging.basicConfig(level=logging.INFO)
+    for handler in logging.root.handlers:
+        # pylint: disable=protected-access
+        handler.setFormatter(SensitiveFormatter(handler.formatter._fmt))  # type: ignore
 
     stopwatch = Stopwatch()
     temp_path = Path(TEMP_PATH)