diff --git a/common/lib/capa/capa/safe_exec.py b/common/lib/capa/capa/safe_exec.py index e6fa0e7bb5..32f6049453 100644 --- a/common/lib/capa/capa/safe_exec.py +++ b/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", diff --git a/common/lib/capa/capa/tests/test_safe_exec.py b/common/lib/capa/capa/tests/test_safe_exec.py new file mode 100644 index 0000000000..035b15b24d --- /dev/null +++ b/common/lib/capa/capa/tests/test_safe_exec.py @@ -0,0 +1,37 @@ +"""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)