fix formatting

Author: sage-maker

src/sagemaker/git_utils.py

@@ -37,17 +37,17 @@ def _sanitize_git_url(repo_url):
     Raises:
         ValueError: If the URL contains suspicious patterns that could indicate injection
     """
-    at_count = repo_url.count('@')
+    at_count = repo_url.count("@")
 
-    if repo_url.startswith('git@'):
+    if repo_url.startswith("git@"):
         # git@ format requires exactly one @
         if at_count != 1:
             raise ValueError("Invalid SSH URL format: git@ URLs must have exactly one @ symbol")
-    elif repo_url.startswith('ssh://'):
+    elif repo_url.startswith("ssh://"):
         # ssh:// format can have 0 or 1 @ symbols
         if at_count > 1:
             raise ValueError("Invalid SSH URL format: multiple @ symbols detected")
-    elif repo_url.startswith('https://') or repo_url.startswith('http://'):
+    elif repo_url.startswith("https://") or repo_url.startswith("http://"):
         # HTTPS format allows 0 or 1 @ symbols
         if at_count > 1:
             raise ValueError("Invalid HTTPS URL format: multiple @ symbols detected")
@@ -58,21 +58,23 @@ def _sanitize_git_url(repo_url):
             # Check for suspicious characters in hostname that could indicate injection
             if parsed.hostname:
                 # Check for URL-encoded characters that might be used for obfuscation
-                suspicious_patterns = ['%25', '%40', '%2F', '%3A']  # encoded %, @, /, :
+                suspicious_patterns = ["%25", "%40", "%2F", "%3A"]  # encoded %, @, /, :
                 for pattern in suspicious_patterns:
                     if pattern in parsed.hostname.lower():
                         raise ValueError(f"Suspicious URL encoding detected in hostname: {pattern}")
 
                 # Validate that the hostname looks legitimate
-                if not re.match(r'^[a-zA-Z0-9.-]+$', parsed.hostname):
+                if not re.match(r"^[a-zA-Z0-9.-]+$", parsed.hostname):
                     raise ValueError("Invalid characters in hostname")
 
         except Exception as e:
             if isinstance(e, ValueError):
                 raise
             raise ValueError(f"Failed to parse URL: {str(e)}")
     else:
-        raise ValueError("Unsupported URL scheme: only https://, http://, git@, and ssh:// are allowed")
+        raise ValueError(
+            "Unsupported URL scheme: only https://, http://, git@, and ssh:// are allowed"
+        )
 
     return repo_url
 
@@ -142,10 +144,10 @@ def git_clone_repo(git_config, entry_point, source_dir=None, dependencies=None):
     if entry_point is None:
         raise ValueError("Please provide an entry point.")
     _validate_git_config(git_config)
-    
+
     # SECURITY: Sanitize the repository URL to prevent injection attacks
     git_config["repo"] = _sanitize_git_url(git_config["repo"])
-    
+
     dest_dir = tempfile.mkdtemp()
     _generate_and_run_clone_command(git_config, dest_dir)
 

tests/unit/test_git_utils.py

@@ -501,6 +501,7 @@ def test_git_clone_repo_codecommit_https_creds_not_stored_locally(tempdir, mkdte
 # URL Sanitization Tests - Security vulnerability prevention
 # ============================================================================
 
+
 class TestGitUrlSanitization:
     """Test cases for Git URL sanitization to prevent injection attacks."""
 
@@ -513,7 +514,7 @@ def test_sanitize_git_url_valid_https_urls(self):
             "https://user:[email protected]/user/repo.git",
             "http://internal-git.company.com/repo.git",
         ]
-        
+
         for url in valid_urls:
             # Should not raise any exception
             result = git_utils._sanitize_git_url(url)
@@ -528,7 +529,7 @@ def test_sanitize_git_url_valid_ssh_urls(self):
             "ssh://git-codecommit.us-west-2.amazonaws.com/v1/repos/test-repo/",  # 0 @ symbols - valid for ssh://
             "[email protected]:repo.git",
         ]
-        
+
         for url in valid_urls:
             # Should not raise any exception
             result = git_utils._sanitize_git_url(url)
@@ -542,7 +543,7 @@ def test_sanitize_git_url_blocks_multiple_at_https(self):
             "https://a@b@[email protected]/repo.git",
             "https://user@[email protected]/legit/repo.git",
         ]
-        
+
         for url in malicious_urls:
             with pytest.raises(ValueError) as error:
                 git_utils._sanitize_git_url(url)
@@ -556,34 +557,34 @@ def test_sanitize_git_url_blocks_multiple_at_ssh(self):
             "ssh://git@[email protected]/repo.git",
             "git@a@b@c:repo.git",
         ]
-        
+
         for url in malicious_urls:
             with pytest.raises(ValueError) as error:
                 git_utils._sanitize_git_url(url)
             # git@ URLs should give "exactly one @ symbol" error
             # ssh:// URLs should give "multiple @ symbols detected" error
-            assert any(phrase in str(error.value) for phrase in [
-                "multiple @ symbols detected",
-                "exactly one @ symbol"
-            ])
+            assert any(
+                phrase in str(error.value)
+                for phrase in ["multiple @ symbols detected", "exactly one @ symbol"]
+            )
 
     def test_sanitize_git_url_blocks_invalid_schemes_and_git_at_format(self):
         """Test that invalid schemes and git@ format violations are blocked."""
         # Test unsupported schemes
         unsupported_scheme_urls = [
             "git-github.com:user/repo.git",  # Doesn't start with git@, ssh://, http://, https://
         ]
-        
+
         for url in unsupported_scheme_urls:
             with pytest.raises(ValueError) as error:
                 git_utils._sanitize_git_url(url)
             assert "Unsupported URL scheme" in str(error.value)
-        
+
         # Test git@ URLs with wrong @ count
         invalid_git_at_urls = [
             "[email protected]@evil.com:repo.git",  # 2 @ symbols
         ]
-        
+
         for url in invalid_git_at_urls:
             with pytest.raises(ValueError) as error:
                 git_utils._sanitize_git_url(url)
@@ -597,15 +598,15 @@ def test_sanitize_git_url_blocks_url_encoding_obfuscation(self):
             "https://github.com%2Fevil.com/repo.git",
             "https://github.com%3Aevil.com/repo.git",
         ]
-        
+
         for url in obfuscated_urls:
             with pytest.raises(ValueError) as error:
                 git_utils._sanitize_git_url(url)
             # The error could be either suspicious encoding or invalid characters
-            assert any(phrase in str(error.value) for phrase in [
-                "Suspicious URL encoding detected",
-                "Invalid characters in hostname"
-            ])
+            assert any(
+                phrase in str(error.value)
+                for phrase in ["Suspicious URL encoding detected", "Invalid characters in hostname"]
+            )
 
     def test_sanitize_git_url_blocks_invalid_hostname_chars(self):
         """Test that hostnames with invalid characters are blocked."""
@@ -615,16 +616,19 @@ def test_sanitize_git_url_blocks_invalid_hostname_chars(self):
             "https://github[].com/repo.git",
             "https://github{}.com/repo.git",
         ]
-        
+
         for url in invalid_urls:
             with pytest.raises(ValueError) as error:
                 git_utils._sanitize_git_url(url)
             # The error could be various types due to URL parsing edge cases
-            assert any(phrase in str(error.value) for phrase in [
-                "Invalid characters in hostname",
-                "Failed to parse URL",
-                "does not appear to be an IPv4 or IPv6 address"
-            ])
+            assert any(
+                phrase in str(error.value)
+                for phrase in [
+                    "Invalid characters in hostname",
+                    "Failed to parse URL",
+                    "does not appear to be an IPv4 or IPv6 address",
+                ]
+            )
 
     def test_sanitize_git_url_blocks_unsupported_schemes(self):
         """Test that unsupported URL schemes are blocked."""
@@ -634,7 +638,7 @@ def test_sanitize_git_url_blocks_unsupported_schemes(self):
             "javascript:alert('xss')",
             "data:text/html,<script>alert('xss')</script>",
         ]
-        
+
         for url in unsupported_urls:
             with pytest.raises(ValueError) as error:
                 git_utils._sanitize_git_url(url)
@@ -644,10 +648,10 @@ def test_git_clone_repo_blocks_malicious_https_url(self):
         """Test that git_clone_repo blocks malicious HTTPS URLs."""
         malicious_git_config = {
             "repo": "https://[email protected]@github.com/legit/repo.git",
-            "branch": "main"
+            "branch": "main",
         }
         entry_point = "train.py"
-        
+
         with pytest.raises(ValueError) as error:
             git_utils.git_clone_repo(malicious_git_config, entry_point)
         assert "multiple @ symbols detected" in str(error.value)
@@ -656,10 +660,10 @@ def test_git_clone_repo_blocks_malicious_ssh_url(self):
         """Test that git_clone_repo blocks malicious SSH URLs."""
         malicious_git_config = {
             "repo": "git@[email protected]:sage-maker/temp-sev2.git",
-            "branch": "main"
+            "branch": "main",
         }
         entry_point = "train.py"
-        
+
         with pytest.raises(ValueError) as error:
             git_utils.git_clone_repo(malicious_git_config, entry_point)
         assert "exactly one @ symbol" in str(error.value)
@@ -668,10 +672,10 @@ def test_git_clone_repo_blocks_url_encoded_attack(self):
         """Test that git_clone_repo blocks URL-encoded attacks."""
         malicious_git_config = {
             "repo": "https://github.com%40attacker.com/repo.git",
-            "branch": "main"
+            "branch": "main",
         }
         entry_point = "train.py"
-        
+
         with pytest.raises(ValueError) as error:
             git_utils.git_clone_repo(malicious_git_config, entry_point)
         assert "Suspicious URL encoding detected" in str(error.value)
@@ -690,17 +694,20 @@ def test_sanitize_git_url_comprehensive_attack_scenarios(self):
             "https://github.com%40evil.com/repo.git",
             "https://[email protected]%2Fevil.com/repo.git",
         ]
-        
+
         entry_point = "train.py"
-        
+
         for malicious_url in attack_scenarios:
             git_config = {"repo": malicious_url}
             with pytest.raises(ValueError) as error:
                 git_utils.git_clone_repo(git_config, entry_point)
             # Should be blocked by sanitization
-            assert any(phrase in str(error.value) for phrase in [
-                "multiple @ symbols detected",
-                "exactly one @ symbol",
-                "Suspicious URL encoding detected",
-                "Invalid characters in hostname"
-            ])
+            assert any(
+                phrase in str(error.value)
+                for phrase in [
+                    "multiple @ symbols detected",
+                    "exactly one @ symbol",
+                    "Suspicious URL encoding detected",
+                    "Invalid characters in hostname",
+                ]
+            )