handle binop case

pixee · Aug 7, 2024 · 463ca4d · 463ca4d
1 parent 979aee8
commit 463ca4d
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 3 deletions.
diff --git a/src/core_codemods/semgrep/semgrep_nan_injection.py b/src/core_codemods/semgrep/semgrep_nan_injection.py
@@ -70,10 +70,8 @@ def _get_target_in_call(self, node: cst.Call) -> cst.CSTNode:
             ):
                 # bool(float(var)), complex(float(var)), bool(float(var)), etc
                 return self._get_target_in_call(wrapped_node)
-            case cst.Call():
+            case cst.Call() | cst.BinaryOperation():
                 return wrapped_node
-            # case cst.BinaryOperation():#
-            #     return node.args[0].value.left.args[0].value.value
 
     def _report_new_lines(self, original_node: cst.SimpleStatementLine):
         self.report_change(original_node)

diff --git a/tests/codemods/semgrep/test_semgrep_nan_injection.py b/tests/codemods/semgrep/test_semgrep_nan_injection.py
@@ -399,6 +399,63 @@ def view(request):
             num_changes=4,
         )
 
+    def test_binop(self, tmpdir):
+        input_code = """\
+        def view(request):
+            tid = request.POST.get("tid")
+            float(tid + str(10))
+        """
+        expected_output = """\
+        def view(request):
+            tid = request.POST.get("tid")
+            if tid + str(10).lower() == "nan":
+                raise ValueError
+            else:
+                float(tid + str(10))
+        """
+        results = {
+            "runs": [
+                {
+                    "results": [
+                        {
+                            "fingerprints": {"matchBasedId/v1": "asd2"},
+                            "locations": [
+                                {
+                                    "physicalLocation": {
+                                        "artifactLocation": {
+                                            "uri": "code.py",
+                                            "uriBaseId": "%SRCROOT%",
+                                        },
+                                        "region": {
+                                            "endColumn": 25,
+                                            "endLine": 3,
+                                            "snippet": {
+                                                "text": "    float(tid + str(10))"
+                                            },
+                                            "startColumn": 5,
+                                            "startLine": 3,
+                                        },
+                                    }
+                                }
+                            ],
+                            "message": {
+                                "text": "Found user input going directly into typecast for bool(), float(), or complex(). This allows an attacker to inject Python's not-a-number (NaN) into the typecast. This results in undefind behavior, particularly when doing comparisons. Either cast to a different type, or add a guard checking for all capitalizations of the string 'nan'."
+                            },
+                            "properties": {},
+                            "ruleId": "python.django.security.nan-injection.nan-injection",
+                        }
+                    ],
+                }
+            ]
+        }
+        self.run_and_assert(
+            tmpdir,
+            input_code,
+            expected_output,
+            results=json.dumps(results),
+            num_changes=4,
+        )
+
 
 def region_data_for_func(func):
     data = {