Refactor to move assumed_imports into capa, so that code_jail is more pure.

common/lib/capa/capa/safe_exec/safe_exec.py

@@ -1,6 +1,7 @@
 """Capa's specialized use of codejail.safe_exec."""
 
 import codejail.safe_exec
+from . import lazymod
 
 # Establish the Python environment for Capa.
 # Capa assumes float-friendly division always.
@@ -16,23 +17,39 @@ del random_module
 sys.modules['random'] = random
 """
 
+ASSUMED_IMPORTS=[
+    ("numpy", "numpy"),
+    ("math", "math"),
+    ("scipy", "scipy"),
+    ("calc", "calc"),
+    ("eia", "eia"),
+    ("chemcalc", "chem.chemcalc"),
+    ("chemtools", "chem.chemtools"),
+    ("miller", "chem.miller"),
+    ("draganddrop", "verifiers.draganddrop"),
+]
+
+# 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()
+
+LAZY_IMPORTS = [lazymod_py]
+for name, modname in ASSUMED_IMPORTS:
+    LAZY_IMPORTS.append("{} = LazyModule('{}')\n".format(name, modname))
+
+LAZY_IMPORTS = "".join(LAZY_IMPORTS)
+
+
 def safe_exec(code, globals_dict, random_seed=None, python_path=None):
     """Exec python code safely.
 
     """
     code_prolog = CODE_PROLOG % random_seed
+
     codejail.safe_exec.safe_exec(
-        code_prolog + code, globals_dict,
+        code_prolog + LAZY_IMPORTS + code, globals_dict,
         python_path=python_path,
-        assumed_imports=[
-            "numpy",
-            "math",
-            "scipy",
-            "calc",
-            "eia",
-            ("chemcalc", "chem.chemcalc"),
-            ("chemtools", "chem.chemtools"),
-            ("miller", "chem.miller"),
-            ("draganddrop", "verifiers.draganddrop"),
-        ],
     )

common/lib/capa/capa/safe_exec/tests/test_lazymod.py

@@ -0,0 +1,44 @@
+"""Test lazymod.py"""
+
+import sys
+import unittest
+
+from capa.safe_exec.lazymod import LazyModule
+
+
+class ModuleIsolation(object):
+    """
+    Manage changes to sys.modules so that we can roll back imported modules.
+
+    Create this object, it will snapshot the currently imported modules. When
+    you call `clean_up()`, it will delete any module imported since its creation.
+    """
+    def __init__(self):
+        # Save all the names of all the imported modules.
+        self.mods = set(sys.modules)
+
+    def clean_up(self):
+        # Get a list of modules that didn't exist when we were created
+        new_mods = [m for m in sys.modules if m not in self.mods]
+        # and delete them all so another import will run code for real again.
+        for m in new_mods:
+            del sys.modules[m]
+
+
+class TestLazyMod(unittest.TestCase):
+
+    def setUp(self):
+        # Each test will remove modules that it imported.
+        self.addCleanup(ModuleIsolation().clean_up)
+
+    def test_simple(self):
+        # Import some stdlib module that has not been imported before
+        self.assertNotIn("colorsys", sys.modules)
+        colorsys = LazyModule("colorsys")
+        hsv = colorsys.rgb_to_hsv(.3, .4, .2)
+        self.assertEqual(hsv[0], 0.25)
+
+    def test_dotted(self):
+        self.assertNotIn("email.utils", sys.modules)
+        email_utils = LazyModule("email.utils")
+        self.assertEqual(email_utils.quote('"hi"'), r'\"hi\"')

common/lib/codejail/codejail/safe_exec.py

@@ -6,40 +6,17 @@ import shutil
 import sys
 import textwrap
 
-import lazymod
 import jailpy
 
-from util import temp_directory, change_directory, TempDirectory
+from util import temp_directory, change_directory
 
-# 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 names_and_modules(assumed_imports):
-    """Get uniform names and modules from assumed_imports."""
-    for modname in assumed_imports:
-        if isinstance(modname, tuple):
-            yield modname
-        else:
-            yield modname, modname
-
-
-def safe_exec(code, globals_dict, assumed_imports=None, files=None, python_path=None):
+def safe_exec(code, globals_dict, files=None, python_path=None):
     """Execute code as "exec" does, but safely.
 
     `code` is a string of Python code.  `globals_dict` is used as the globals
     during execution.  Modifications the code makes to `globals_dict` are
     reflected in the dictionary on return.
 
-    `assumed_imports` is a list of modules to make available as implicit
-    imports for the code.  Entries are either a name, "mod", which makes
-    "import mod" part of the code, or a pair, ("f", "fooey"), which makes
-    "import fooey as f" part of the code.  The module name can be dotted.
-
     Returns None.  Changes made by `code` are visible in `globals_dict`.
 
     """
@@ -72,11 +49,6 @@ def safe_exec(code, globals_dict, assumed_imports=None, files=None, python_path=
         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):
-            the_code.append("g_dict['{}'] = LazyModule('{}')\n".format(name, modname))
-
     the_code.append(textwrap.dedent(
         # Execute the sandboxed code.
         """
@@ -140,7 +112,7 @@ def json_safe(d):
     return json.loads(json.dumps(jd))
 
 
-def not_safe_exec(code, globals_dict, assumed_imports=None, files=None, python_path=None):
+def not_safe_exec(code, globals_dict, 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.
@@ -151,9 +123,6 @@ def not_safe_exec(code, globals_dict, assumed_imports=None, files=None, python_p
     """
     g_dict = json_safe(globals_dict)
 
-    for name, modname in names_and_modules(assumed_imports or ()):
-        g_dict[name] = lazymod.LazyModule(modname)
-
     with temp_directory(delete_when_done=True) as tmpdir:
         with change_directory(tmpdir):
             # Copy the files here.

common/lib/codejail/codejail/tests/test_lazymod.py

@@ -1,26 +0,0 @@
-"""Test lazymod.py"""
-
-import sys
-import unittest
-
-from codejail.lazymod import LazyModule
-from codejail.util import ModuleIsolation
-
-
-class TestLazyMod(unittest.TestCase):
-
-    def setUp(self):
-        # Each test will remove modules that it imported.
-        self.addCleanup(ModuleIsolation().clean_up)
-
-    def test_simple(self):
-        # Import some stdlib module that has not been imported before
-        self.assertNotIn("colorsys", sys.modules)
-        colorsys = LazyModule("colorsys")
-        hsv = colorsys.rgb_to_hsv(.3, .4, .2)
-        self.assertEqual(hsv[0], 0.25)
-
-    def test_dotted(self):
-        self.assertNotIn("email.utils", sys.modules)
-        email_utils = LazyModule("email.utils")
-        self.assertEqual(email_utils.quote('"hi"'), r'\"hi\"')

common/lib/codejail/codejail/tests/test_safe_exec.py

@@ -14,19 +14,6 @@ class SafeExecTests(object):
         self.safe_exec("a = 17", g)
         self.assertEqual(g['a'], 17)
 
-    def test_assumed_imports(self):
-        g = {}
-        # Using string without importing it is bad.
-        with self.assertRaises(Exception):
-            self.safe_exec("a = string.ascii_lowercase[0]", g)
-        # Using string with an assumed import is fine.
-        self.safe_exec("a = string.ascii_lowercase[0]", g, assumed_imports=["string"])
-        self.assertEqual(g['a'], 'a')
-        # Can also import with a shorthand.
-        self.safe_exec("a = op.join('x', 'y')", g, assumed_imports=[("op", "os.path")])
-        self.assertEqual(g['a'][0], 'x')
-        self.assertEqual(g['a'][-1], 'y')
-
     def test_files_are_copied(self):
         g = {}
         self.safe_exec(

common/lib/codejail/codejail/util.py

@@ -32,25 +32,6 @@ def temp_directory(delete_when_done=True):
         tmp.clean_up()
 
 
-class ModuleIsolation(object):
-    """
-    Manage changes to sys.modules so that we can roll back imported modules.
-
-    Create this object, it will snapshot the currently imported modules. When
-    you call `clean_up()`, it will delete any module imported since its creation.
-    """
-    def __init__(self):
-        # Save all the names of all the imported modules.
-        self.mods = set(sys.modules)
-
-    def clean_up(self):
-        # Get a list of modules that didn't exist when we were created
-        new_mods = [m for m in sys.modules if m not in self.mods]
-        # 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()