Use import call modifier pattern for url-sandbox (#750)

Author: drdavella

src/codemodder/codemods/import_modifier_codemod.py

@@ -6,7 +6,7 @@
 
 from codemodder.codemods.api import LibcstResultTransformer
 from codemodder.codemods.imported_call_modifier import ImportedCallModifier
-from codemodder.dependency import Dependency
+from codemodder.dependency import Dependency, Security
 
 
 class MappingImportedCallModifier(ImportedCallModifier[Mapping[str, str]]):
@@ -15,7 +15,7 @@ def update_attribute(self, true_name, original_node, updated_node, new_args):
             return updated_node
 
         import_name = self.matching_functions[true_name]
-        AddImportsVisitor.add_needed_import(self.context, import_name)
+        self.add_import(import_name)
         RemoveImportsVisitor.remove_unused_import_by_node(self.context, original_node)
         return updated_node.with_changes(
             args=new_args,
@@ -30,7 +30,7 @@ def update_simple_name(self, true_name, original_node, updated_node, new_args):
             return updated_node
 
         import_name = self.matching_functions[true_name]
-        AddImportsVisitor.add_needed_import(self.context, import_name)
+        self.add_import(import_name)
         RemoveImportsVisitor.remove_unused_import_by_node(self.context, original_node)
         return updated_node.with_changes(
             args=new_args,
@@ -40,6 +40,9 @@ def update_simple_name(self, true_name, original_node, updated_node, new_args):
             ),
         )
 
+    def add_import(self, import_name):
+        AddImportsVisitor.add_needed_import(self.context, import_name)
+
 
 class ImportModifierCodemod(LibcstResultTransformer, metaclass=ABCMeta):
     call_modifier: type[MappingImportedCallModifier] = MappingImportedCallModifier
@@ -67,3 +70,14 @@ def transform_module_impl(self, tree: cst.Module) -> cst.Module:
             self.add_dependency(dependency)
 
         return result_tree
+
+
+class SecurityCallModifier(MappingImportedCallModifier):
+    def add_import(self, import_name: str) -> None:
+        AddImportsVisitor.add_needed_import(
+            self.context, module=Security.requirement.name, obj=import_name
+        )
+
+
+class SecurityImportModifierCodemod(ImportModifierCodemod, metaclass=ABCMeta):
+    call_modifier: type[SecurityCallModifier] = SecurityCallModifier

src/codemodder/codemods/imported_call_modifier.py

@@ -68,7 +68,9 @@ def update_simple_name(
     def leave_Call(self, original_node: cst.Call, updated_node: cst.Call):
         pos_to_match = self.node_position(original_node)
         line_number = pos_to_match.start.line
-        if self.filter_by_path_includes_or_excludes(pos_to_match):
+        if self.node_is_selected(
+            original_node
+        ) and self.filter_by_path_includes_or_excludes(pos_to_match):
             true_name = self.find_base_name(original_node.func)
             if (
                 self.is_direct_call_from_imported_module(original_node)

src/core_codemods/url_sandbox.py

@@ -1,156 +1,26 @@
-from typing import List, Optional, Union
+from functools import cached_property
 
-import libcst as cst
-from libcst import CSTNode, matchers
-from libcst.codemod import CodemodContext, ContextAwareVisitor
-from libcst.codemod.visitors import AddImportsVisitor, ImportItem
-from libcst.metadata import PositionProvider, ScopeProvider
-
-from codemodder.codemods.base_visitor import UtilsMixin
-from codemodder.codemods.libcst_transformer import (
-    LibcstResultTransformer,
-    LibcstTransformerPipeline,
-)
+from codemodder.codemods.import_modifier_codemod import SecurityImportModifierCodemod
+from codemodder.codemods.libcst_transformer import LibcstTransformerPipeline
 from codemodder.codemods.semgrep import SemgrepRuleDetector
-from codemodder.codemods.transformations.remove_unused_imports import (
-    RemoveUnusedImportsCodemod,
-)
-from codemodder.codemods.utils import ReplaceNodes
-from codemodder.codetf import Change
-from codemodder.dependency import Security
-from codemodder.file_context import FileContext
+from codemodder.dependency import Dependency, Security
 from core_codemods.api import CoreCodemod, Metadata, Reference, ReviewGuidance
 
-replacement_import = "safe_requests"
-
 
-class UrlSandboxTransformer(LibcstResultTransformer):
+class UrlSandboxTransformer(SecurityImportModifierCodemod):
     change_description = "Switch use of requests for security.safe_requests"
-    METADATA_DEPENDENCIES = (PositionProvider, ScopeProvider)
-    adds_dependency = True
-
-    def transform_module_impl(self, tree: cst.Module) -> cst.Module:
-        # we first gather all the nodes we want to change together with their replacements
-        find_requests_visitor = FindRequestCallsAndImports(
-            self.context,
-            self.file_context,
-            self.file_context.results,
-        )
-        tree.visit(find_requests_visitor)
-        if find_requests_visitor.nodes_to_change:
-            self.file_context.codemod_changes.extend(
-                find_requests_visitor.changes_in_file
-            )
-            new_tree = tree.visit(ReplaceNodes(find_requests_visitor.nodes_to_change))
-            self.add_dependency(Security)
-            # if it finds any request.get(...), try to remove the imports
-            if any(
-                (
-                    matchers.matches(n, matchers.Call())
-                    for n in find_requests_visitor.nodes_to_change
-                )
-            ):
-                new_tree = AddImportsVisitor(
-                    self.context,
-                    [ImportItem(Security.name, replacement_import, None, 0)],
-                ).transform_module(new_tree)
-                new_tree = RemoveUnusedImportsCodemod(self.context).transform_module(
-                    new_tree
-                )
-            return new_tree
-        return tree
-
-
-class FindRequestCallsAndImports(ContextAwareVisitor, UtilsMixin):
-    METADATA_DEPENDENCIES = (ScopeProvider,)
-
-    def __init__(
-        self, codemod_context: CodemodContext, file_context: FileContext, results
-    ):
-        self.nodes_to_change: dict[
-            cst.CSTNode, Union[cst.CSTNode, cst.FlattenSentinel, cst.RemovalSentinel]
-        ] = {}
-        self.changes_in_file: List[Change] = []
-        self.file_context = file_context
-        ContextAwareVisitor.__init__(self, codemod_context)
-        UtilsMixin.__init__(
-            self,
-            results=results,
-            line_include=file_context.line_include,
-            line_exclude=file_context.line_exclude,
-        )
 
-    def leave_Call(self, original_node: cst.Call):
-        if not self.node_is_selected(original_node):
-            return
+    @cached_property
+    def mapping(self) -> dict[str, str]:
+        """Build a mapping of functions to their safe_requests imports"""
+        _matching_functions: dict[str, str] = {
+            "requests.get": "safe_requests",
+        }
+        return _matching_functions
 
-        line_number = self.node_position(original_node).start.line
-        match original_node.args[0].value:
-            case cst.SimpleString():
-                return
-
-        match original_node:
-            # case get(...)
-            case cst.Call(func=cst.Name()):
-                # find if get(...) comes from an from requests import get
-                match self.find_single_assignment(original_node):
-                    case cst.ImportFrom() as node:
-                        self.nodes_to_change.update(
-                            {
-                                node: cst.ImportFrom(
-                                    module=cst.Attribute(
-                                        value=cst.Name(Security.name),
-                                        attr=cst.Name(replacement_import),
-                                    ),
-                                    names=node.names,
-                                )
-                            }
-                        )
-                        self.changes_in_file.append(
-                            Change(
-                                lineNumber=line_number,
-                                description=UrlSandboxTransformer.change_description,
-                                findings=self.file_context.get_findings_for_location(
-                                    line_number
-                                ),
-                            )
-                        )
-
-            # case req.get(...)
-            case _:
-                self.nodes_to_change.update(
-                    {
-                        original_node: cst.Call(
-                            func=cst.parse_expression(replacement_import + ".get"),
-                            args=original_node.args,
-                        )
-                    }
-                )
-                self.changes_in_file.append(
-                    Change(
-                        lineNumber=line_number,
-                        description=UrlSandboxTransformer.change_description,
-                        findings=self.file_context.get_findings_for_location(
-                            line_number
-                        ),
-                    )
-                )
-
-    def _find_assignments(self, node: CSTNode):
-        """
-        Given a MetadataWrapper and a CSTNode representing an access, find all the possible assignments that it refers.
-        """
-        scope = self.get_metadata(ScopeProvider, node)
-        return next(iter(scope.accesses[node]))._Access__assignments
-
-    def find_single_assignment(self, node: CSTNode) -> Optional[CSTNode]:
-        """
-        Given a MetadataWrapper and a CSTNode representing an access, find if there is a single assignment that it refers to.
-        """
-        assignments = self._find_assignments(node)
-        if len(assignments) == 1:
-            return next(iter(assignments)).node
-        return None
+    @property
+    def dependency(self) -> Dependency:
+        return Security
 
 
 UrlSandbox = CoreCodemod(
@@ -174,20 +44,20 @@ def find_single_assignment(self, node: CSTNode) -> Optional[CSTNode]:
     ),
     detector=SemgrepRuleDetector(
         """
-    rules:
-      - id: url-sandbox
-        message: Unbounded URL creation
-        severity: WARNING
-        languages:
-          - python
-        pattern-either:
-          - patterns:
-            - pattern: requests.get(...)
-            - pattern-not: requests.get("...")
-            - pattern-inside: |
-                import requests
-                ...
-            """
+         rules:
+           - id: url-sandbox
+             message: Unbounded URL creation
+             severity: WARNING
+             languages:
+               - python
+             pattern-either:
+               - patterns:
+                 - pattern: requests.get(...)
+                 - pattern-not: requests.get("...")
+                 - pattern-inside: |
+                     import requests
+                     ...
+    """
     ),
     transformer=LibcstTransformerPipeline(UrlSandboxTransformer),
 )

tests/codemods/test_url_sandbox.py

@@ -40,10 +40,10 @@ def test_from_requests(self, add_dependency, tmpdir):
         var = "hello"
         """
         expected = """
-        from security.safe_requests import get
+        from security import safe_requests
 
         url = input()
-        get(url)
+        safe_requests.get(url)
         var = "hello"
         """
         self.run_and_assert(tmpdir, input_code, expected)
@@ -160,10 +160,10 @@ def test_from_requests_with_alias(self, add_dependency, tmpdir):
         var = "hello"
         """
         expected = """
-        from security.safe_requests import get as got
+        from security import safe_requests
 
         url = input()
-        got(url)
+        safe_requests.get(url)
         var = "hello"
         """
         self.run_and_assert(tmpdir, input_code, expected)
@@ -198,6 +198,24 @@ def test_ignore_hardcoded(self, _, tmpdir):
 
         self.run_and_assert(tmpdir, input_code, expected)
 
+    def test_ignore_hardcoded_but_not_all(self, _, tmpdir):
+        input_code = """
+        import requests
+
+        requests.get("www.google.com")
+        url = input()
+        requests.get(url)
+        """
+        expected = """
+        import requests
+        from security import safe_requests
+
+        requests.get("www.google.com")
+        url = input()
+        safe_requests.get(url)
+        """
+        self.run_and_assert(tmpdir, input_code, expected)
+
     def test_ignore_hardcoded_from_global_variable(self, _, tmpdir):
         expected = (
             input_code
@@ -262,9 +280,6 @@ def foo():
 
         self.run_and_assert(tmpdir, input_code, expected)
 
-    @pytest.mark.xfail(
-        reason="Does not properly handle 'from' imports with multiple names"
-    )
     def test_multiple_imports(self, add_dependency, tmpdir):
         input_code = """
         from requests import Response, Timeout, get