AugAssign propagates taint

Before, the variable would be tainted only if the last += was tainted.

Now

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

url marked as tainted.

Author: Ben Caller

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]
         ))
 

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):