Implement safe_exec on top of jailpy (old unsafe safe_exec is still here);…

Implement safe_exec on top of jailpy (old unsafe safe_exec is still here); Remove some crazy stuff from the context; always pass globals and locals, locals are the things that can be changed.

codejail/safe_exec.py

 """Safe execution of untrusted Python code."""
 
 import json
+import textwrap
 
-from .lazymod import LazyModule
+import lazymod
+import jailpy
 
-def straw(v):
-    return json.loads(json.dumps(jsonable_dict(v)))
-
-def jsonable_dict(d):
+# If we aren't running safe, then we need to artificially pass the values
+# through a JSON straw to ensure we aren't passing something that won't
+# be executable in the safe context.
+def straw(d):
     jd = {}
     for k,v in d.iteritems():
         try:
@@ -16,9 +18,9 @@ def jsonable_dict(d):
             continue
         else:
             jd[k] = v
-    return jd
+    return json.loads(json.dumps(jd))
 
-def safe_exec(code, globals_dict, locals_dict=None, future_division=False, assumed_imports=None):
+def safe_exec(code, globals_dict, locals_dict, future_division=False, assumed_imports=None):
     """Execute code safely.
 
     Returns None.  The code can modify globals in `global_dict`.
@@ -28,21 +30,57 @@ def safe_exec(code, globals_dict, locals_dict=None, future_division=False, assum
         code = "from __future__ import division\n" + code
 
     g_dict = straw(globals_dict)
-
-    if locals_dict is None:
-        l_dict = g_dict
-    else:
-        l_dict = straw(locals_dict)
+    l_dict = straw(locals_dict)
 
     for modname in assumed_imports or ():
         if isinstance(modname, tuple):
             name, modname = modname
         else:
             name = modname
-        g_dict[name] = LazyModule(modname)
+        g_dict[name] = lazymod.LazyModule(modname)
 
     exec code in g_dict, l_dict
 
     globals_dict.update(straw(g_dict))
-    if locals_dict is not None:
-        locals_dict.update(straw(l_dict))
+    locals_dict.update(straw(l_dict))
+
+# We'll need the code from lazymod.py for use in jailpy, so read it now.
+lazymod_py_file = lazymod.__file__
+if lazymod_py_file.endswith("c"):
+    lazymod_py_file = lazymod_py_file[:-1]
+
+lazymod_py = open(lazymod_py_file).read()
+
+
+def xxxsafe_exec(code, globals_dict, locals_dict, future_division=False, assumed_imports=None):
+    the_code = []
+
+    the_code.append(textwrap.dedent("""\
+        import json
+        import sys
+        code, g_dict, l_dict = json.load(sys.stdin)
+        """))
+
+    if assumed_imports:
+        the_code.append(lazymod_py)
+        for modname in assumed_imports:
+            if isinstance(modname, tuple):
+                name, modname = modname
+            else:
+                name = modname
+            the_code.append("g_dict['{}'] = LazyModule('{}')\n".format(name, modname))
+
+    the_code.append(textwrap.dedent("""\
+        exec code in g_dict, l_dict
+        print >>sys.stderr, l_dict.keys()
+        ok_types = (int, long, float, str, unicode, list, tuple, dict)
+        l_dict = {k:v for k,v in l_dict.iteritems() if isinstance(v, ok_types)}
+        json.dump(l_dict, sys.stdout)
+        """))
+
+    print "".join(the_code)
+    stdin = json.dumps([code, globals_dict, locals_dict])
+    res = jailpy.jailpy("".join(the_code), stdin=stdin)
+    if res.status != 0:
+        raise Exception("Couldn't excecute jailed code: %s" % res.stderr)
+    locals_dict.update(json.loads(res.stdout))

codejail/tests/test_jailpy.py

@@ -33,6 +33,13 @@ class TestFeatures(unittest.TestCase):
             Exception: FAIL
             """))
 
+    def test_stdin_is_provided(self):
+        res = jailpy(
+            "import json,sys; print sum(json.load(sys.stdin))",
+            stdin="[1, 2.5, 33]"
+        )
+        self.assertEqual(res.stdout.strip(), "36.5")
+
 
 class TestLimits(unittest.TestCase):
     def test_cant_use_too_much_memory(self):