AugAssign propagates taint

bcaller · Ben Caller · commit b9f740ce3d3e · 2018-07-27T12:49:55.000+01:00
Before, the variable would be tainted only if the last += was tainted.

Now

url = 'http://'
url += TAINT
url += '?x=y'

Is marked as vulnerable.
diff --git a/pyt/cfg/stmt_visitor.py b/pyt/cfg/stmt_visitor.py
@@ -499,11 +499,12 @@ def visit_AugAssign(self, node):
         rhs_visitor = RHSVisitor()
         rhs_visitor.visit(node.value)
 
+        lhs = extract_left_hand_side(node.target)
         return self.append_node(AssignmentNode(
             label.result,
-            extract_left_hand_side(node.target),
+            lhs,
             node,
-            rhs_visitor.result,
+            rhs_visitor.result + [lhs],
             path=self.filenames[-1]
         ))
 
diff --git a/tests/cfg/cfg_test.py b/tests/cfg/cfg_test.py
@@ -820,6 +820,14 @@ def test_assignment_starred_list(self):
             [('a', ['d']), ('b', ['d']), ('c', ['e'])],
         )
 
+    def test_augmented_assignment(self):
+        self.cfg_create_from_ast(ast.parse('a+=f(b,c)'))
+
+        (node,) = self.cfg.nodes[1:-1]
+        self.assertEqual(node.label, 'a += f(b, c)')
+        self.assertEqual(node.left_hand_side, 'a')
+        self.assertEqual(node.right_hand_side_variables, ['b', 'c', 'a'])
+
 
 class CFGComprehensionTest(CFGBaseTestCase):
     def test_nodes(self):
diff --git a/tests/main_test.py b/tests/main_test.py
@@ -84,11 +84,11 @@ def test_targets_with_recursive(self):
         excluded_files = ""
 
         included_files = discover_files(targets, excluded_files, True)
-        self.assertEqual(len(included_files), 31)
+        self.assertEqual(len(included_files), 32)
 
     def test_targets_with_recursive_and_excluded(self):
         targets = ["examples/vulnerable_code/"]
         excluded_files = "inter_command_injection.py"
 
         included_files = discover_files(targets, excluded_files, True)
-        self.assertEqual(len(included_files), 30)
+        self.assertEqual(len(included_files), 31)
