Files are properly copied in both implementations of safe_exec, and a new…

Files are properly copied in both implementations of safe_exec, and a new python_path argument adds to the python path.

codejail/jailpy.py

@@ -63,8 +63,11 @@ def jailpy(code, files=None, argv=None, stdin=None):
 
         # All the supporting files are copied into our directory.
         for filename in files or ():
-            dest = os.path.join(tmpdir, os.path.basename(filename))
-            shutil.copyfile(filename, dest)
+            if os.path.isfile(filename):
+                shutil.copy(filename, tmpdir)
+            else:
+                dest = os.path.join(tmpdir, os.path.basename(filename))
+                shutil.copytree(filename, dest)
 
         # Create the main file.
         with open(os.path.join(tmpdir, "jailed_code.py"), "w") as jailed:

codejail/safe_exec.py

 """Safe execution of untrusted Python code."""
 
 import json
+import os.path
+import shutil
+import sys
 import textwrap
 
 import lazymod
 import jailpy
 
+from util import temp_directory, change_directory, TempDirectory
+
 # 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"):
@@ -23,7 +28,7 @@ def names_and_modules(assumed_imports):
             yield modname, modname
 
 
-def safe_exec(code, globals_dict, locals_dict, future_division=False, assumed_imports=None):
+def safe_exec(code, globals_dict, locals_dict, future_division=False, assumed_imports=None, files=None, python_path=None):
     """Execute code as "exec" does, but safely.
 
     `code` is a string of Python code.  `globals_dict` and `locals_dict` are
@@ -41,6 +46,7 @@ def safe_exec(code, globals_dict, locals_dict, future_division=False, assumed_im
 
     """
     the_code = []
+    files = list(files or ())
 
     if future_division:
         the_code.append("from __future__ import division\n")
@@ -51,6 +57,11 @@ def safe_exec(code, globals_dict, locals_dict, future_division=False, assumed_im
         code, g_dict, l_dict = json.load(sys.stdin)
         """))
 
+    for pydir in python_path or ():
+        pybase = os.path.basename(pydir)
+        the_code.append("sys.path.append(%r)\n" % pybase)
+        files.append(pydir)
+
     if assumed_imports:
         the_code.append(lazymod_py)
         for name, modname in names_and_modules(assumed_imports):
@@ -63,25 +74,30 @@ def safe_exec(code, globals_dict, locals_dict, future_division=False, assumed_im
         json.dump(l_dict, sys.stdout)
         """))
 
+    stdin = json.dumps([code, globals_dict, locals_dict])
+    jailed_code = "".join(the_code)
+
     # Turn this on to see what's being executed.
-    if 0:
+    if 1:
         print "--{:-<40}".format(" jailed ")
-        print "".join(the_code)
+        print jailed_code
         print "--{:-<40}".format(" exec ")
         print code
 
-    stdin = json.dumps([code, globals_dict, locals_dict])
-    res = jailpy.jailpy("".join(the_code), stdin=stdin)
+    res = jailpy.jailpy(jailed_code, stdin=stdin, files=files)
     if res.status != 0:
         raise Exception("Couldn't excecute jailed code: %s" % res.stderr)
     locals_dict.update(json.loads(res.stdout))
 
 
-def not_safe_exec(code, globals_dict, locals_dict, future_division=False, assumed_imports=None):
+def not_safe_exec(code, globals_dict, locals_dict, future_division=False, assumed_imports=None, files=None, python_path=None):
     """Another implementation of `safe_exec`, but not safe.
 
     This can be swapped in for debugging problems in sandboxed Python code.
 
+    This is not thread-safe, due to temporarily changing the current directory
+    and modifying sys.path.
+
     """
     def straw(d):
         """Return only the JSON-safe part of d.
@@ -108,7 +124,20 @@ def not_safe_exec(code, globals_dict, locals_dict, future_division=False, assume
     for name, modname in names_and_modules(assumed_imports or ()):
         g_dict[name] = lazymod.LazyModule(modname)
 
-    exec code in g_dict, l_dict
+    with temp_directory(delete_when_done=True) as tmpdir:
+        with change_directory(tmpdir):
+            # Copy the files here.
+            for filename in files or ():
+                dest = os.path.join(tmpdir, os.path.basename(filename))
+                shutil.copyfile(filename, dest)
+
+            original_path = sys.path
+            if python_path:
+                sys.path.extend(python_path)
+            try:
+                exec code in g_dict, l_dict
+            finally:
+                sys.path = original_path
 
     locals_dict.update(straw(l_dict))
 

codejail/tests/hello.txt

+Hello there.

codejail/tests/pylib/module.py

+const = 42

codejail/tests/test_jailpy.py

 """Test jailpy.py"""
 
+import os.path
 import textwrap
 import unittest
 from nose.plugins.skip import SkipTest
@@ -53,6 +54,14 @@ class TestFeatures(JailPyHelpers, unittest.TestCase):
         self.assertResultOk(res)
         self.assertEqual(res.stdout.strip(), "36.5")
 
+    def test_files_are_copied(self):
+        res = jailpy(
+            "print 'Look:', open('hello.txt').read()",
+            files=[os.path.dirname(__file__) + "/hello.txt"]
+        )
+        self.assertResultOk(res)
+        self.assertEqual(res.stdout, 'Look: Hello there.\n\n')
+
 
 class TestLimits(JailPyHelpers, unittest.TestCase):
     def test_cant_use_too_much_memory(self):

codejail/tests/test_safe_exec.py

 """Test safe_exec.py"""
 
-import textwrap
+import os.path
 import unittest
 from nose.plugins.skip import SkipTest
 
 from codejail.safe_exec import safe_exec, not_safe_exec
 
-dedent = textwrap.dedent
-
 class SafeExecTests(object):
     """The tests for `safe_exec`, will be mixed into specific test classes below."""
     def test_set_values(self):
@@ -37,6 +35,22 @@ class SafeExecTests(object):
         self.assertEqual(l['a'][0], 'x')
         self.assertEqual(l['a'][-1], 'y')
 
+    def test_files_are_copied(self):
+        g, l = {}, {}
+        self.safe_exec(
+            "a = 'Look: ' + open('hello.txt').read()", g, l,
+            files=[os.path.dirname(__file__) + "/hello.txt"]
+        )
+        self.assertEqual(l['a'], 'Look: Hello there.\n')
+
+    def test_python_path(self):
+        g, l = {}, {}
+        self.safe_exec(
+            "import module; a = module.const", g, l,
+            python_path=[os.path.dirname(__file__) + "/pylib"]
+        )
+        self.assertEqual(l['a'], 42)
+
 
 class TestSafeExec(SafeExecTests, unittest.TestCase):
     """Run SafeExecTests, with the real safe_exec."""

codejail/util.py

@@ -49,3 +49,23 @@ class ModuleIsolation(object):
         # and delete them all so another import will run code for real again.
         for m in new_mods:
             del sys.modules[m]
+
+
+class ChangeDirectory(object):
+    def __init__(self, new_dir):
+        self.old_dir = os.getcwd()
+        os.chdir(new_dir)
+
+    def clean_up(self):
+        os.chdir(self.old_dir)
+
+@contextlib.contextmanager
+def change_directory(new_dir):
+    """
+    A context manager to change the directory, and then change it back.
+    """
+    cd = ChangeDirectory(new_dir)
+    try:
+        yield new_dir
+    finally:
+        cd.clean_up()