Merge pull request #21 from edx/ned/extra-files

jail_code can create extra files for you from content.

codejail/jail_code.py

@@ -116,8 +116,8 @@ class JailResult(object):
         self.stdout = self.stderr = self.status = None
 
 
-def jail_code(command, code=None, files=None, argv=None, stdin=None,
-              slug=None):
+def jail_code(command, code=None, files=None, extra_files=None, argv=None,
+              stdin=None, slug=None):
     """
     Run code in a jailed subprocess.
 
@@ -135,6 +135,11 @@ def jail_code(command, code=None, files=None, argv=None, stdin=None,
     linked-to file is not accessible to the sandbox, the symlink will be
     unreadable as well.
 
+    `extra_files` is a list of pairs, each pair is a filename and a bytestring
+    of contents to write into that file.  These files will be created in the
+    temp directory and cleaned up automatically.  No subdirectories are
+    supported in the filename.
+
     `argv` is the command-line arguments to supply.
 
     `stdin` is a string, the data to provide as the stdin for the process.
@@ -181,11 +186,16 @@ def jail_code(command, code=None, files=None, argv=None, stdin=None,
 
         # Create the main file.
         if code:
-            with open(os.path.join(homedir, "jailed_code"), "w") as jailed:
+            with open(os.path.join(homedir, "jailed_code"), "wb") as jailed:
                 jailed.write(code)
 
             argv = ["jailed_code"] + argv
 
+        # Create extra files requested by the caller:
+        for name, content in extra_files or ():
+            with open(os.path.join(homedir, name), "wb") as extra:
+                extra.write(content)
+
         cmd = []
 
         # Build the command to run.

codejail/tests/test_jail_code.py

@@ -111,6 +111,25 @@ class TestFeatures(JailCodeHelpers, unittest.TestCase):
             "This is doit.py!\nMy args are ['doit.py', '1', '2', '3']\n"
         )
 
+    def test_executing_extra_files(self):
+        res = jailpy(
+            extra_files=[
+                ("run.py", textwrap.dedent("""\
+                            import os
+                            print os.listdir('.')
+                            print open('also.txt').read()
+                            """)),
+                # This file has some non-ASCII, non-UTF8, just binary data.
+                ("also.txt", "also here\xff\x00\xab"),
+            ],
+            argv=["run.py"],
+        )
+        self.assertResultOk(res)
+        self.assertEqual(
+            res.stdout,
+            "['tmp', 'also.txt', 'run.py']\nalso here\xff\x00\xab\n"
+        )
+
 
 class TestLimits(JailCodeHelpers, unittest.TestCase):
     """Tests of the resource limits, and changing them."""