fix: censoring sensitive values from being logged

Author: mohamedzeidan2021

.github/workflows/codeql.yml

@@ -11,9 +11,7 @@ jobs:
     name: Analyze (${{ matrix.language }})
     runs-on: ${{ 'ubuntu-latest' }}
     permissions:
-      # required for all workflows
       security-events: write
-      # required to fetch internal or private CodeQL packs
       packages: read
 
     strategy:
@@ -25,14 +23,13 @@ jobs:
             build-mode: none
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
-      # Initializes the CodeQL tools for scanning.
+        uses: actions/checkout@6ccd57f4c5d15bdc2fef309bd9fb6cc9db2ef1c6
       - name: Initialize CodeQL
-        uses: github/codeql-action/init@v3
+        uses: github/codeql-action/init@4b1d7da102ff94aca014c0245062b1a463356d72
         with:
           languages: ${{ matrix.language }}
           build-mode: ${{ matrix.build-mode }}
       - name: Perform CodeQL Analysis
-        uses: github/codeql-action/analyze@v3
+        uses: github/codeql-action/analyze@4b1d7da102ff94aca014c0245062b1a463356d72
         with:
           category: "/language:${{matrix.language}}"

src/sagemaker/config/config_utils.py

@@ -18,8 +18,10 @@
 from collections import deque
 
 import logging
+import re
 import sys
 from typing import Callable
+from copy import deepcopy
 
 
 def get_sagemaker_config_logger():
@@ -67,6 +69,19 @@ def _log_sagemaker_config_single_substitution(source_value, config_value, config
     """
     logger = get_sagemaker_config_logger()
 
+    source_value_log_copy = deepcopy(source_value)
+    config_value_log_copy = deepcopy(config_value)
+
+    if isinstance(source_value_log_copy, dict):
+        for key in source_value_log_copy.keys():
+            if re.search(r'(secret|password|key|token)', key, re.IGNORECASE):
+                source_value_log_copy[key] = '***'
+
+    if isinstance(config_value_log_copy, dict):
+        for key in config_value_log_copy.keys():
+            if re.search(r'(secret|password|key|token)', key, re.IGNORECASE):
+                config_value_log_copy[key] = '***'
+
     if config_value is not None:
 
         if source_value is None:
@@ -79,7 +94,7 @@ def _log_sagemaker_config_single_substitution(source_value, config_value, config
                 logger.debug(
                     "Applied value\n  config key = %s\n  config value that will be used = %s",
                     config_key_path,
-                    config_value,
+                    config_value_log_copy,
                 )
             else:
                 logger.info(
@@ -102,8 +117,8 @@ def _log_sagemaker_config_single_substitution(source_value, config_value, config
                     "  source value that will be used = %s"
                 ),
                 config_key_path,
-                config_value,
-                source_value,
+                config_value_log_copy,
+                source_value_log_copy,
             )
         elif source_value is not None and config_value != source_value:
             # Sagemaker Config had a value defined that is NOT going to be used
@@ -117,8 +132,8 @@ def _log_sagemaker_config_single_substitution(source_value, config_value, config
                     "  source value that will be used = %s",
                 ),
                 config_key_path,
-                config_value,
-                source_value,
+                config_value_log_copy,
+                source_value_log_copy,
             )
     else:
         # nothing was specified in the config and nothing is being automatically applied

src/sagemaker/fw_utils.py

@@ -152,7 +152,6 @@
     "2.1.0",
     "2.1.2",
     "2.2.0",
-    "2.3.0",
     "2.3.1",
 ]
 

src/sagemaker/image_uri_config/pytorch-smp.json

@@ -8,7 +8,7 @@
             "2.1": "2.1.2",
             "2.2": "2.3.1",
             "2.2.0": "2.3.1",
-            "2.3": "2.4.0"
+            "2.3.1": "2.4.0"
         },
         "versions": {
             "2.0.1": {

src/sagemaker/image_uri_config/pytorch.json

@@ -1008,6 +1008,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
@@ -1052,6 +1053,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
@@ -2331,6 +2333,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
@@ -2419,6 +2422,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"

src/sagemaker/image_uri_config/tensorflow.json

@@ -2140,6 +2140,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
@@ -2181,6 +2182,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
@@ -4354,6 +4356,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"

src/sagemaker/session.py

@@ -8348,7 +8348,9 @@ def _logs_for_job(  # noqa: C901 - suppress complexity warning for this method
     """
     sagemaker_client = sagemaker_session.sagemaker_client
     request_end_time = time.time() + timeout if timeout else None
-    description = sagemaker_client.describe_training_job(TrainingJobName=job_name)
+    description = _wait_until(
+        lambda: sagemaker_client.describe_training_job(TrainingJobName=job_name)
+    )
     print(secondary_training_status_message(description, None), end="")
 
     instance_count, stream_names, positions, client, log_group, dot, color_wrap = _logs_init(

tests/data/marketplace/iris/scoring_logic.py

@@ -3,7 +3,7 @@
 import logging
 import re
 from flask import Flask
-from flask import request
+from flask import request, escape
 from joblib import dump, load
 import numpy as np
 import os
@@ -106,4 +106,4 @@ def endpoint_invocations():
 
         return response
     except Exception as e:
-        return f"Error during model invocation: {str(e)} for input: {request.get_data()}"
+        return f"Error during model invocation: {str(e)} for input: {escape(request.get_data())}"

tests/unit/sagemaker/image_uris/test_smp_v2.py

@@ -27,6 +27,7 @@ def test_smp_v2(load_config):
         "torch_distributed": {"enabled": True},
         "smdistributed": {"modelparallel": {"enabled": True}},
     }
+
     for processor in PROCESSORS:
         for version in VERSIONS:
             ACCOUNTS = load_config["training"]["versions"][version]["registries"]
@@ -38,6 +39,11 @@ def test_smp_v2(load_config):
                         if "2.1" in version or "2.2" in version or "2.3" in version:
                             cuda_vers = "cu121"
 
+                        if "2.3.1" == version:
+                            py_version = "py311"
+
+                        print(version, py_version)
+
                         uri = image_uris.get_training_image_uri(
                             region,
                             framework="pytorch",

tests/unit/test_utils.py

@@ -15,6 +15,7 @@
 from __future__ import absolute_import
 
 import copy
+import logging
 import shutil
 import tarfile
 from datetime import datetime
@@ -59,6 +60,8 @@
     _validate_new_tags,
     remove_tag_with_key,
 )
+
+from src.sagemaker.config.config_utils import _log_sagemaker_config_single_substitution
 from tests.unit.sagemaker.workflow.helpers import CustomStep
 from sagemaker.workflow.parameters import ParameterString, ParameterInteger
 
@@ -1279,6 +1282,91 @@ def test_resolve_value_from_config():
     mock_info_logger.reset_mock()
 
 
+class TestLogSagemakerConfig(TestCase):
+
+    def test_sensitive_info_masking(self):
+        logger = logging.getLogger('sagemaker.config')
+        logger.setLevel(logging.DEBUG)
+
+        stream_handler = logging.StreamHandler()
+        logger.addHandler(stream_handler)
+
+        # source value is None
+        with self.assertLogs(logger, level='DEBUG') as log:
+            _log_sagemaker_config_single_substitution(
+                None,
+                {"apiKey": "topsecretkey"},
+                "config/path"
+            )
+
+        self.assertIn("config value that will be used = {'apiKey': '***'}", log.output[0])
+
+        # source value is None and config_value == source_value
+        with self.assertLogs(logger, level='DEBUG') as log:
+            _log_sagemaker_config_single_substitution(
+                {"secretword": "topsecretword"},
+                {"secretword": "topsecretword"},
+                "config/path"
+            )
+
+        self.assertIn("Skipped value", log.output[0])
+        self.assertIn("source value that will be used = {'secretword': '***'}", log.output[0])
+        self.assertIn("config value = {'secretword': '***'}", log.output[0])
+
+        # source value is not None and config_value != source_value
+        with self.assertLogs(logger, level='DEBUG') as log:
+            _log_sagemaker_config_single_substitution(
+                {"password": "supersecretpassword"},
+                {"apiKey": "topsecretkey"},
+                "config/path"
+            )
+
+        self.assertIn("Skipped value", log.output[0])
+        self.assertIn("source value that will be used = {'password': '***'}", log.output[0])
+        self.assertIn("config value = {'apiKey': '***'}", log.output[0])
+
+    def test_non_sensitive_info_masking(self):
+        logger = logging.getLogger('sagemaker.config')
+        logger.setLevel(logging.DEBUG)
+
+        stream_handler = logging.StreamHandler()
+        logger.addHandler(stream_handler)
+
+        # source value is None
+        with self.assertLogs(logger, level='DEBUG') as log:
+            _log_sagemaker_config_single_substitution(
+                None,
+                {"username": "randomvalue"},
+                "config/path"
+            )
+
+        self.assertIn("config value that will be used = {'username': 'randomvalue'}", log.output[0])
+
+        # source value is not None and config_value == source_value
+        with self.assertLogs(logger, level='DEBUG') as log:
+            _log_sagemaker_config_single_substitution(
+                {"nonsensitivevalue": "randomvalue"},
+                {"nonsensitivevalue": "randomvalue"},
+                "config/path"
+            )
+
+        self.assertIn("Skipped value", log.output[0])
+        self.assertIn("source value that will be used = {'nonsensitivevalue': 'randomvalue'}", log.output[0])
+        self.assertIn("config value = {'nonsensitivevalue': 'randomvalue'}", log.output[0])
+
+        # source value is not None and config_value != source_value
+        with self.assertLogs(logger, level='DEBUG') as log:
+            _log_sagemaker_config_single_substitution(
+                {"username": "nonsensitiveinfo"},
+                {"configvalue": "nonsensitivevalue"},
+                "config/path/non_sensitive"
+            )
+
+            self.assertIn("Skipped value", log.output[0])
+            self.assertIn("source value that will be used = {'username': 'nonsensitiveinfo'}", log.output[0])
+            self.assertIn("config value = {'configvalue': 'nonsensitivevalue'}", log.output[0])
+
+
 def test_get_sagemaker_config_value():
     mock_config_logger = Mock()