python-security · KevinHock · Jul 31, 2018 · Jul 31, 2018 · KevinHock · Jul 31, 2018
diff --git a/examples/example_inputs/asynchronous.py b/examples/example_inputs/asynchronous.py
@@ -0,0 +1,10 @@
+async def g(x, *args):
+    return await x()
+
+
+async def f(y):
+    z = await g(y, await v)
+    return z
+
+
+f(w)
diff --git a/pyt/core/ast_helper.py b/pyt/core/ast_helper.py
@@ -6,6 +6,8 @@
 import subprocess
 from functools import lru_cache
 
+from .transformer import AsyncTransformer
+
 
 BLACK_LISTED_CALL_NAMES = ['self']
 recursive = False
@@ -32,7 +34,8 @@ def generate_ast(path):
     if os.path.isfile(path):
         with open(path, 'r') as f:
             try:
-                return ast.parse(f.read())
+                tree = ast.parse(f.read())
+                return AsyncTransformer().visit(tree)
             except SyntaxError:  # pragma: no cover
                 global recursive
                 if not recursive:

diff --git a/pyt/core/transformer.py b/pyt/core/transformer.py
@@ -0,0 +1,18 @@
+import ast
+
+
+class AsyncTransformer(ast.NodeTransformer):
+    """Converts all async nodes into their synchronous counterparts."""
+
+    def visit_Await(self, node):
+        """Awaits are treated as if the keyword was absent."""
+        return self.visit(node.value)
+
+    def visit_AsyncFunctionDef(self, node):
+        return self.visit(ast.FunctionDef(**node.__dict__))
+
+    def visit_AsyncFor(self, node):
+        return self.visit(ast.For(**node.__dict__))
+
+    def visit_AsyncWith(self, node):
+        return self.visit(ast.With(**node.__dict__))
diff --git a/pyt/helper_visitors/vars_visitor.py b/pyt/helper_visitors/vars_visitor.py
@@ -68,9 +68,6 @@ def visit_GeneratorComp(self, node):
         for gen in node.generators:
             self.comprehension(gen)
 
-    def visit_Await(self, node):
-        self.visit(node.value)
-
     def visit_Yield(self, node):
         if node.value:
             self.visit(node.value)

diff --git a/tests/cfg/cfg_test.py b/tests/cfg/cfg_test.py
@@ -1446,6 +1446,28 @@ def test_name_constant_if(self):
         self.assertEqual(expected_label, actual_label)
 
 
+class CFGAsync(CFGBaseTestCase):
+    def test_await_keyword_treated_as_if_absent(self):
+        self.cfg_create_from_file('examples/example_inputs/asynchronous.py')
+        enter_g = 8
+        call_x = 9
+        ret_g = 10
+        exit_g = 11
+        call_ret_val = 12
+        set_z_to_g_ret_val = 13
+
+        for i in range(enter_g, set_z_to_g_ret_val + 1):
+            self.assertIn(self.cfg.nodes[i], self.cfg.nodes[i + 1].ingoing)
+            self.assertIn(self.cfg.nodes[i + 1], self.cfg.nodes[i].outgoing)
+
+        self.assertIsInstance(self.cfg.nodes[enter_g], EntryOrExitNode)
+        self.assertEqual(self.cfg.nodes[call_x].label, '~call_3 = ret_x()')
+        self.assertEqual(self.cfg.nodes[ret_g].label, 'ret_g = ~call_3')
+        self.assertIsInstance(self.cfg.nodes[exit_g], EntryOrExitNode)
+        self.assertEqual(self.cfg.nodes[call_ret_val].label, '~call_2 = ret_g')
+        self.assertEqual(self.cfg.nodes[set_z_to_g_ret_val].label, 'z = ~call_2')
+
+
 class CFGName(CFGBaseTestCase):
     """Test is Name nodes are properly handled in different contexts"""
 

diff --git a/tests/core/transformer_test.py b/tests/core/transformer_test.py
@@ -0,0 +1,36 @@
+import ast
+import unittest
+
+from pyt.core.transformer import AsyncTransformer
+
+
+class TransformerTest(unittest.TestCase):
+    """Tests for the AsyncTransformer."""
+
+    def test_async_removed_by_transformer(self):
+        async_tree = ast.parse("\n".join([
+            "async def a():",
+            "   async for b in c():",
+            "       await b()",
+            "   async with d() as e:",
+            "       pass",
+            "   return await y()"
+        ]))
+        self.assertIsInstance(async_tree.body[0], ast.AsyncFunctionDef)
+        self.assertIsInstance(async_tree.body[0].body[-1], ast.Return)
+        self.assertIsInstance(async_tree.body[0].body[-1].value, ast.Await)
+
+        sync_tree = ast.parse("\n".join([
+            "def a():",
+            "   for b in c():",
+            "       b()",
+            "   with d() as e:",
+            "       pass",
+            "   return y()"
+        ]))
+        self.assertIsInstance(sync_tree.body[0], ast.FunctionDef)
+
+        transformed = AsyncTransformer().visit(async_tree)
+        self.assertIsInstance(transformed.body[0], ast.FunctionDef)
+
+        self.assertEqual(ast.dump(transformed), ast.dump(sync_tree))