Add test case and update rule for dataflow

pixee · Sep 6, 2023 · 0c31f47 · 0c31f47
1 parent 395054d
commit 0c31f47
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 17 deletions.
diff --git a/src/codemodder/codemods/upgrade_sslcontext_minimum_version.py b/src/codemodder/codemods/upgrade_sslcontext_minimum_version.py
@@ -11,23 +11,25 @@ class UpgradeSSLContextMinimumVersion(SemgrepCodemod):
     def rule(cls):
         return """
         rules:
-          - patterns:
-            - pattern: |
-                $CONTEXT.minimum_version = ssl.TLSVersion.$VERSION
-            - pattern-inside: |
-                import ssl
-                ...
-                $CONTEXT = ssl.SSLContext(...)
-                ...
-            - metavariable-pattern:
-                metavariable: $VERSION
-                patterns:
-                  - pattern-either:
-                    - pattern: SSLv2
-                    - pattern: SSLv3
-                    - pattern: TLSv1
-                    - pattern: TLSv1_1
-                    - pattern: MINIMUM_SUPPORTED
+          - mode: taint
+            pattern-sources:
+              - patterns:
+                - pattern: ssl.SSLContext(...)
+                - pattern-inside: |
+                    import ssl
+                    ...
+            pattern-sinks:
+              - patterns:
+                - pattern: $SINK.minimum_version = ssl.TLSVersion.$VERSION
+                - metavariable-pattern:
+                    metavariable: $VERSION
+                    patterns:
+                      - pattern-either:
+                        - pattern: SSLv2
+                        - pattern: SSLv3
+                        - pattern: TLSv1
+                        - pattern: TLSv1_1
+                        - pattern: MINIMUM_SUPPORTED
         """
 
     def on_result_found(self, original_node, updated_node):

diff --git a/tests/codemods/test_upgrade_sslcontext_minimum_version.py b/tests/codemods/test_upgrade_sslcontext_minimum_version.py
@@ -68,5 +68,20 @@ def test_import_with_alias(self, tmpdir):
 
 context = whatever.SSLContext()
 context.minimum_version = ssl.TLSVersion.TLSv1_2
+"""
+        self.run_and_assert(tmpdir, input_code, expexted_output)
+
+    def test_with_dataflow(self, tmpdir):
+        input_code = """import ssl
+
+context = ssl.SSLContext()
+alias = context
+alias.minimum_version = ssl.TLSVersion.SSLv3
+"""
+        expexted_output = """import ssl
+
+context = ssl.SSLContext()
+alias = context
+alias.minimum_version = ssl.TLSVersion.TLSv1_2
 """
         self.run_and_assert(tmpdir, input_code, expexted_output)