safe_exec seeds the random module, and now we have tests for it.

common/lib/capa/capa/safe_exec.py

@@ -2,11 +2,17 @@
 
 import codejail.safe_exec
 
-def safe_exec(code, globals_dict, locals_dict):
+CODE_PROLOG = """\
+import random as random_module
+random = random_module.Random(%r)
+random.Random = random_module.Random
+"""
+
+def safe_exec(code, globals_dict, locals_dict, random_seed=None):
+    code_prolog = CODE_PROLOG % random_seed
     codejail.safe_exec.safe_exec(
-        code, globals_dict, locals_dict, future_division=True,
+        code_prolog + code, globals_dict, locals_dict, future_division=True,
         assumed_imports=[
-            "random",
             "numpy",
             "math",
             "scipy",

common/lib/capa/capa/tests/test_safe_exec.py

+"""Test safe_exec.py"""
+
+import random
+import unittest
+
+from capa.safe_exec import safe_exec
+
+class TestSafeExec(unittest.TestCase):
+    def test_set_values(self):
+        g, l = {}, {}
+        safe_exec("a = 17", g, l)
+        self.assertEqual(l['a'], 17)
+
+    def test_division(self):
+        g, l = {}, {}
+        # Future division: 1/2 is 0.5.
+        safe_exec("a = 1/2", g, l)
+        self.assertEqual(l['a'], 0.5)
+
+    def test_assumed_imports(self):
+        g, l = {}, {}
+        # Math is always available.
+        safe_exec("a = int(math.pi)", g, l)
+        self.assertEqual(l['a'], 3)
+
+    def test_random_seeding(self):
+        g, l = {}, {}
+        r = random.Random(17)
+        rnums = [r.randint(0, 999) for _ in xrange(100)]
+
+        # Without a seed, the results are unpredictable
+        safe_exec("rnums = [random.randint(0, 999) for _ in xrange(100)]", g, l)
+        self.assertNotEqual(l['rnums'], rnums)
+
+        # With a seed, the results are predictable
+        safe_exec("rnums = [random.randint(0, 999) for _ in xrange(100)]", g, l, random_seed=17)
+        self.assertEqual(l['rnums'], rnums)