First step towards using the last module.

src/RestrictedPython/RCompile.py

@@ -14,7 +14,7 @@
 Python standard library.
 """
 
-from ast import parse
+import ast
 
 from compiler import ast as c_ast
 from compiler import parse as c_parse
@@ -40,7 +40,7 @@ def niceParse(source, filename, mode):
         source = '\xef\xbb\xbf' + source.encode('utf-8')
     try:
         compiler_code = c_parse(source, mode)
-        # ast_code = parse(source, filename, mode)
+        # ast_code = ast.parse(source, filename, mode)
         return compiler_code
     except:
         # Try to make a clean error message using
@@ -117,7 +117,27 @@ def compile_restricted_eval(source, filename='<string>'):
     return compileAndTuplize(gen)
 
 
-def compile_restricted(source, filename, mode):
+class RestrictingNodeTransformer(ast.NodeTransformer):
+
+    whitelist = [
+        ast.Module,
+        ast.FunctionDef,
+        ast.Expr,
+        ast.Num,
+    ]
+
+    def generic_visit(self, node):
+        if node.__class__ not in self.whitelist:
+            raise SyntaxError(
+                'Node {0.__class__.__name__!r} not allowed.'.format(node))
+        else:
+            return super(RestrictingNodeTransformer, self).generic_visit(node)
+
+    def visit_arguments(self, node):
+        return node
+
+
+def compile_restricted(source, filename, mode):  # OLD
     """Replacement for the builtin compile() function."""
     if mode == "single":
         gen = RInteractive(source, filename)
@@ -132,6 +152,13 @@ def compile_restricted(source, filename, mode):
     return gen.getCode()
 
 
+def compile_restricted_ast(source, filename, mode):
+    """Replacement for the builtin compile() function."""
+    c_ast = ast.parse(source, filename, mode)
+    r_ast = RestrictingNodeTransformer().visit(c_ast)
+    return compile(r_ast, filename, mode)
+
+
 class RestrictedCodeGenerator:
     """Mixin for CodeGenerator to replace UNPACK_SEQUENCE bytecodes.
 

src/RestrictedPython/tests/security_yield.py

@@ -0,0 +1,7 @@
+# These are all supposed to raise a SyntaxError when using
+# compile_restricted() but not when using compile().
+# Each function in this module is compiled using compile_restricted().
+
+
+def no_yield():
+    yield 42

src/RestrictedPython/tests/testRestrictions.py

@@ -9,6 +9,7 @@
 from RestrictedPython.Eval import RestrictionCapableEval
 from RestrictedPython.tests import restricted_module, verify
 from RestrictedPython.RCompile import RModule, RFunction
+from RestrictedPython.RCompile import compile_restricted_ast
 
 import os
 import re
@@ -298,6 +299,16 @@ def checkSyntaxSecurity(self):
         if sys.version_info >= (2, 7):
             self._checkSyntaxSecurity('security_in_syntax27.py')
 
+    def checkSyntaxSecurityAst(self):
+        with self.assertRaises(SyntaxError) as err:
+            with open(os.path.join(os.path.join(
+                    _HERE, 'security_yield.py'))) as f:
+                compile_restricted_ast(f.read(), 'security_yield.py', 'exec')
+        self.assertEqual("Node 'Yield' not allowed.", str(err.exception))
+
+    def checkNumberAstOk(self):
+        compile_restricted_ast('42', '<undefined>', 'exec')
+
     def _checkSyntaxSecurity(self, mod_name):
         # Ensures that each of the functions in security_in_syntax.py
         # throws a SyntaxError when using compile_restricted.