Address Code Review

Author: yashnap

src/core/src/service_interfaces/CredentialSanitizer.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 #
 # Requires Python 2.7+
-import logging
 import re
 
 
@@ -36,11 +35,10 @@ def sanitize(self, message):
             # (1) scheme: https://, http://, or ftp://
             # (2) username: one or more non-whitespace, non-slash, non-colon, non-@ characters
             # (3) password: zero or more non-whitespace, non-slash, non-@ characters
-            sanitized_message = re.sub(
-                r'(https?://|ftp://)([^:/@\s]+):([^@/\s]*)@',r'\1\2@',message)
+            sanitized_message = re.sub(r'(https?://|ftp://)([^:/@\s]+):([^@/\s]*)@',r'\1\2@',message)
             self.composite_logger.log_verbose("Message was sanitized to remove sensitive information. [InputMessage={0}][SanitizedMessage={1}]".format(str(message), str(sanitized_message)))
             return sanitized_message
         except Exception as error:
-            self.composite_logger.log_error("Error occurred while sanitizing credentials from message: {0}".format(repr(error)))
+            self.composite_logger.log_error("Error occurred while sanitizing credentials from message: [Error={0}]".format(repr(error)))
             return message
 

src/core/tests/Test_TelemetryWriter.py

@@ -71,7 +71,7 @@ def test_write_event(self):
 
         self.assertTrue(telemetry_event_counter_in_first_test_event is not None)
         self.assertTrue(telemetry_event_counter_in_second_test_event is not None)
-        self.assertTrue(int(telemetry_event_counter_in_second_test_event) - int(telemetry_event_counter_in_first_test_event) == 1 if telemetry_event_counter_in_first_test_event and telemetry_event_counter_in_second_test_event else False)
+        self.assertTrue(int(telemetry_event_counter_in_second_test_event) - int(telemetry_event_counter_in_first_test_event) == 1)
 
     def test_write_multiple_events_in_same_file(self):
         time_backup = time.time
@@ -312,103 +312,125 @@ def test_write_event_with_buffer_true_and_empty_string_and_then_flush_with_non_e
             self.assertTrue(text_found.string.startswith("Message 1"))
 
     # ==================== Unit Tests for Credential Sanitization ====================
-    def test_sanitize_credentials_from_uri_https_with_credentials_leak(self):
-        """ Test sanitization of HTTPS URIs with credentials """
-        # Clear any existing event files before test
+    # ==================== Helper functions for Credential Sanitization Tests ====================
+    def _clear_events_folder(self):
+        """
+        Helper method to clear the events folder for sanitization test setup.
+        Removes all existing JSON event files.
+        """
         for f in os.listdir(self.runtime.telemetry_writer.events_folder_path):
             if f.endswith('.json'):
                 os.remove(os.path.join(self.runtime.telemetry_writer.events_folder_path, f))
 
-        # Verify events folder is empty before test
-        self.assertTrue(len(os.listdir(self.runtime.telemetry_writer.events_folder_path)) == 0, "Events folder should be empty before writing")
+    def _read_event_from_file(self, file_index=None, event_index=-1):
+        """
+        Helper method to open and read an event from an event file in the events folder.
+        Args:
+            file_index: Index of the event file to read. If None, uses latest file
+            event_index: Index of the event within the file (default: -1 for last event)
+        Returns: The parsed event dictionary from the JSON file
+        """
+        event_files = [pos_json for pos_json in os.listdir(self.runtime.telemetry_writer.events_folder_path) if re.search('^[0-9]+.json$', pos_json)]
+        if not event_files:
+            raise Exception("No event files found in events folder")
+
+        if file_index is None:
+            event_file_path = os.path.join(self.runtime.telemetry_writer.events_folder_path, event_files[-1])
+        else:
+            event_file_path = os.path.join(self.runtime.telemetry_writer.events_folder_path, event_files[file_index])
+
+        with open(event_file_path, 'r+') as f:
+            events = json.load(f)
+            f.close()
+            if not events:
+                raise Exception("No events found in event file")
+            return events[event_index]
+
+    def _get_message_without_tc(self, event):
+        """
+        Helper method to extract the message without the TC (telemetry counter) portion.
+        Args:
+            event: The event dictionary
+        Returns: The message portion before " [TC=" marker
+        """
+        return event["Message"][:event["Message"].rfind(" [TC=")]
+
+    def _validate_sanitized_event(self, expected_message, task_name=None, event_index=-1, file_index=None):
+        """
+        Helper method to validate an event's message and task name against expected values.
+        Args:
+            expected_message: The expected sanitized message (without TC counter)
+            task_name: The expected task name (optional validation)
+            event_index: Index of the event within the file (default: -1 for last event)
+            file_index: Index of the event file (default: None for latest file)
+        """
+        event = self._read_event_from_file(file_index=file_index, event_index=event_index)
+
+        self.assertIsNotNone(event)
+        message_without_tc = self._get_message_without_tc(event)
+        self.assertEqual(expected_message, message_without_tc)
+        if task_name is not None:
+            self.assertEqual(task_name, event["TaskName"])
+
+    # ==================== Credential Sanitization Test Cases ====================
+    def test_sanitize_credentials_from_uri_https_with_credentials_leak(self):
+        """ Test sanitization of HTTPS URIs with credentials """
+        self._clear_events_folder()
+        self.assertEqual(len([f for f in os.listdir(self.runtime.telemetry_writer.events_folder_path) if re.search('^[0-9]+.json$', f)]), 0)
 
         message = "Error connecting to https://testuser:TESTTOKEN123456@invalid.repo.example/rpm/repodata/repomd.xml"
+        expected_message = "Error connecting to https://testuser@invalid.repo.example/rpm/repodata/repomd.xml"
+
         self.runtime.telemetry_writer.write_event(message, Constants.TelemetryEventLevel.Error, "Test Task")
 
-        # Verify exactly one event file was created
-        latest_event_file = [pos_json for pos_json in os.listdir(self.runtime.telemetry_writer.events_folder_path) if re.search('^[0-9]+.json$', pos_json)][-1]
+        # Validate exactly one event file was created
         event_files_count = len([f for f in os.listdir(self.runtime.telemetry_writer.events_folder_path) if re.search('^[0-9]+.json$', f)])
-        self.assertEqual(event_files_count, 1, "Events folder should contain exactly one event file")
+        self.assertEqual(event_files_count, 1)
 
-        with open(os.path.join(self.runtime.telemetry_writer.events_folder_path, latest_event_file), 'r+') as f:
-            events = json.load(f)
-            self.assertTrue(events is not None)
-            self.assertEqual(events[-1]["TaskName"], "Test Task")
-            message_without_tc = events[-1]["Message"][:events[-1]["Message"].rfind(" [TC=")]
-            self.assertEqual("Error connecting to https://testuser@invalid.repo.example/rpm/repodata/repomd.xml", message_without_tc)
-            f.close()
+        # Validate using helper
+        self._validate_sanitized_event(expected_message, task_name="Test Task")
 
     def test_sanitize_credentials_from_uri_http_with_credentials_leak(self):
         """ Test sanitization of HTTP URIs with credentials """
         message = "Connection failed to http://user123:password123@example.com/path"
+        expected_message = "Connection failed to http://user123@example.com/path"
+
         self.runtime.telemetry_writer.write_event(message, Constants.TelemetryEventLevel.Error, "Test Task")
 
-        latest_event_file = [pos_json for pos_json in os.listdir(self.runtime.telemetry_writer.events_folder_path) if re.search('^[0-9]+.json$', pos_json)][-1]
-        with open(os.path.join(self.runtime.telemetry_writer.events_folder_path, latest_event_file), 'r+') as f:
-            events = json.load(f)
-            self.assertTrue(events is not None)
-            self.assertEqual(events[-1]["TaskName"], "Test Task")
-            # Verify entire message matches expected output (excluding TC counter)
-            message_without_tc = events[-1]["Message"][:events[-1]["Message"].rfind(" [TC=")]
-            self.assertEqual("Connection failed to http://user123@example.com/path", message_without_tc)
-            f.close()
+        self._validate_sanitized_event(expected_message, task_name="Test Task")
 
     def test_sanitize_credentials_multiple_urls_with_credentials_leak(self):
         """ Test sanitization with multiple URLs containing credentials """
         message = "Failed to fetch from https://user1:pass1@host1.com/api and http://user2:pass2@host2.com/data"
+        expected_message = "Failed to fetch from https://user1@host1.com/api and http://user2@host2.com/data"
+
         self.runtime.telemetry_writer.write_event(message, Constants.TelemetryEventLevel.Error, "Test Task")
 
-        latest_event_file = [pos_json for pos_json in os.listdir(self.runtime.telemetry_writer.events_folder_path) if re.search('^[0-9]+.json$', pos_json)][-1]
-        with open(os.path.join(self.runtime.telemetry_writer.events_folder_path, latest_event_file), 'r+') as f:
-            events = json.load(f)
-            self.assertTrue(events is not None)
-            self.assertEqual(events[-1]["TaskName"], "Test Task")
-            message_without_tc = events[-1]["Message"][:events[-1]["Message"].rfind(" [TC=")]
-            self.assertEqual("Failed to fetch from https://user1@host1.com/api and http://user2@host2.com/data", message_without_tc)
-            f.close()
+        self._validate_sanitized_event(expected_message, task_name="Test Task")
 
     def test_sanitize_credentials_with_error_and_no_credentials(self):
         """  ERROR with 401 status code from jfrog.io """
         message = "ERROR: Failed to download metadata for repo 'packages-microsoft-com-prod': Status code: 401 for https://cec-aa.jfrog.io/artifactory/glib-rpm-hel9-lts-microsoft-com/repodata/repomd.xml"
+        expected_message = "ERROR: Failed to download metadata for repo 'packages-microsoft-com-prod': Status code: 401 for https://cec-aa.jfrog.io/artifactory/glib-rpm-hel9-lts-microsoft-com/repodata/repomd.xml"
+
         self.runtime.telemetry_writer.write_event(message, Constants.TelemetryEventLevel.Error, "Test Task")
 
-        latest_event_file = [pos_json for pos_json in os.listdir(self.runtime.telemetry_writer.events_folder_path) if re.search('^[0-9]+.json$', pos_json)][-1]
-        with open(os.path.join(self.runtime.telemetry_writer.events_folder_path, latest_event_file), 'r+') as f:
-            events = json.load(f)
-            self.assertTrue(events is not None)
-            self.assertEqual(events[-1]["TaskName"], "Test Task")
-            # Verify entire message matches expected output (excluding TC counter)
-            message_without_tc = events[-1]["Message"][:events[-1]["Message"].rfind(" [TC=")]
-            self.assertEqual("ERROR: Failed to download metadata for repo 'packages-microsoft-com-prod': Status code: 401 for https://cec-aa.jfrog.io/artifactory/glib-rpm-hel9-lts-microsoft-com/repodata/repomd.xml", message_without_tc)
-            f.close()
+        self._validate_sanitized_event(expected_message, task_name="Test Task")
 
     def test_sanitize_credentials_with_error_and_credentials_leak(self):
         """  Curl error with buildbot:BuildBotToken credentials """
         message = ("Curl error (6): Couldn't resolve host 'packages.microsoft.com' Could not "
                    "retrieve mirrorlist https://buildbot:BuildBotToken@mirror.example.com/repodata/repomd.xml")
-        self.runtime.telemetry_writer.write_event(message, Constants.TelemetryEventLevel.Error, "Test Task")
+        expected_message = ("Curl error (6): Couldn't resolve host 'packages.microsoft.com' Could not "
+                           "retrieve mirrorlist https://buildbot@mirror.example.com/repodata/repomd.xml")
 
-        latest_event_file = [pos_json for pos_json in os.listdir(self.runtime.telemetry_writer.events_folder_path) if re.search('^[0-9]+.json$', pos_json)][-1]
-        with open(os.path.join(self.runtime.telemetry_writer.events_folder_path, latest_event_file), 'r+') as f:
-            events = json.load(f)
-            self.assertTrue(events is not None)
-            self.assertEqual(events[-1]["TaskName"], "Test Task")
-            # Verify entire message matches expected output (excluding TC counter)
-            message_without_tc = events[-1]["Message"][:events[-1]["Message"].rfind(" [TC=")]
-            expected_message = ("Curl error (6): Couldn't resolve host 'packages.microsoft.com' Could not "
-                               "retrieve mirrorlist https://buildbot@mirror.example.com/repodata/repomd.xml")
-            self.assertEqual(expected_message, message_without_tc)
-            f.close()
+        self.runtime.telemetry_writer.write_event(message, Constants.TelemetryEventLevel.Error, "Test Task")
+        self._validate_sanitized_event(expected_message, task_name="Test Task")
 
     def test_sanitize_credentials_with_credentials_leak(self):
         """ ERROR with expired SSL certs and TESTTOKEN123456 """
-        # Clear any existing event files before test
-        for f in os.listdir(self.runtime.telemetry_writer.events_folder_path):
-            if f.endswith('.json'):
-                os.remove(os.path.join(self.runtime.telemetry_writer.events_folder_path, f))
-
-        # Verify events folder is empty before test
-        self.assertTrue(len([f for f in os.listdir(self.runtime.telemetry_writer.events_folder_path) if re.search('^[0-9]+.json$', f)]) == 0, "Events folder should be empty before writing")
+        self._clear_events_folder()
+        self.assertEqual(len([f for f in os.listdir(self.runtime.telemetry_writer.events_folder_path) if re.search('^[0-9]+.json$', f)]), 0)
 
         message = ("ERROR: Customer environment error (expired SSL certs): "
                    "Command=sudo yum update -y --disablerepo='*' "
@@ -419,30 +441,22 @@ def test_sanitize_credentials_with_credentials_leak(self):
                    "for https://testuser:TESTTOKEN123456@packages-microsoft-com-prod/CENTRAL.rpm "
                    "Error: Failed to download metadata for repo 'packages-microsoft-com-prod': "
                    "Cannot download repomd.xml: All mirrors were tried")
+        expected_message = ("ERROR: Customer environment error (expired SSL certs): "
+                           "Command=sudo yum update -y --disablerepo='*' "
+                           "--enablerepo='microsoft' !!Code=11 Out- Updating "
+                           "Subscription Management repositories. "
+                           "Unable to read consumer identity This system is not registered "
+                           "with an entitlement server. Status code: 401 "
+                           "for https://testuser@packages-microsoft-com-prod/CENTRAL.rpm "
+                           "Error: Failed to download metadata for repo 'packages-microsoft-com-prod': "
+                           "Cannot download repomd.xml: All mirrors were tried")
+
         self.runtime.telemetry_writer.write_event(message, Constants.TelemetryEventLevel.Error, "Test Task")
 
-        # Verify exactly one event file was created
+        # Validate exactly one event file was created
         event_files_count = len([f for f in os.listdir(self.runtime.telemetry_writer.events_folder_path) if re.search('^[0-9]+.json$', f)])
-        self.assertEqual(event_files_count, 1, "Events folder should contain exactly one event file")
-
-        latest_event_file = [pos_json for pos_json in os.listdir(self.runtime.telemetry_writer.events_folder_path) if re.search('^[0-9]+.json$', pos_json)][-1]
-        with open(os.path.join(self.runtime.telemetry_writer.events_folder_path, latest_event_file), 'r+') as f:
-            events = json.load(f)
-            self.assertTrue(events is not None)
-            self.assertEqual(events[-1]["TaskName"], "Test Task")
-            # Verify entire message matches expected output (excluding TC counter)
-            message_without_tc = events[-1]["Message"][:events[-1]["Message"].rfind(" [TC=")]
-            expected_message = ("ERROR: Customer environment error (expired SSL certs): "
-                               "Command=sudo yum update -y --disablerepo='*' "
-                               "--enablerepo='microsoft' !!Code=11 Out- Updating "
-                               "Subscription Management repositories. "
-                               "Unable to read consumer identity This system is not registered "
-                               "with an entitlement server. Status code: 401 "
-                               "for https://testuser@packages-microsoft-com-prod/CENTRAL.rpm "
-                               "Error: Failed to download metadata for repo 'packages-microsoft-com-prod': "
-                               "Cannot download repomd.xml: All mirrors were tried")
-            self.assertEqual(expected_message, message_without_tc)
-            f.close()
+        self.assertEqual(event_files_count, 1)
+        self._validate_sanitized_event(expected_message, task_name="Test Task")
 
     def test_sanitize_credentials_exception_handling(self):
         """ Test exception handling: passing None should return the input unchanged """

src/extension/src/CredentialSanitizer.py

@@ -36,12 +36,10 @@ def sanitize(self, message):
             # (1) scheme: https://, http://, or ftp://
             # (2) username: one or more non-whitespace, non-slash, non-colon, non-@ characters
             # (3) password: zero or more non-whitespace, non-slash, non-@ characters
-            sanitized_message = re.sub(
-                r'(https?://|ftp://)([^:/@\s]+):([^@/\s]*)@',r'\1\2@',message)
-            self.logger.log_verbose(
-                "Message was sanitized to remove sensitive information. [InputMessage={0}][SanitizedMessage={1}]".format(str(message), str(sanitized_message)))
+            sanitized_message = re.sub(r'(https?://|ftp://)([^:/@\s]+):([^@/\s]*)@',r'\1\2@',message)
+            self.logger.log_verbose("Message was sanitized to remove sensitive information. [InputMessage={0}][SanitizedMessage={1}]".format(str(message), str(sanitized_message)))
             return sanitized_message
         except Exception as error:
-            self.logger.log_error("Error occurred while sanitizing credentials from message: {0}".format(repr(error)))
+            self.logger.log_error("Error occurred while sanitizing credentials from message: [Error={0}]".format(repr(error)))
             return message
 

src/extension/src/TelemetryWriter.py

@@ -74,7 +74,6 @@ def __ensure_message_restriction_compliance(self, full_message):
             self.logger.log_telemetry_module_error("Error occurred while formatting message for a telemetry event. [Error={0}]".format(repr(e)))
             raise
 
-
     def __get_agent_supports_telemetry_from_env_var(self):
         """ Returns True if the env var AZURE_GUEST_AGENT_EXTENSION_SUPPORTED_FEATURES has a key of
             ExtensionTelemetryPipeline in the list. Value of the env var looks like this: