def testExecuteModifiedPython(self):
"""Test that rejects invalid ExecutePython action."""
utils.TEST_VAL = "original"
python_code = "utils.TEST_VAL = 'modified'"
signed_blob = rdfvalue.SignedBlob()
signed_blob.Sign(python_code, self.signing_key)
# Modify the data so the signature does not match.
signed_blob.data = "utils.TEST_VAL = 'notmodified'"
request = rdfvalue.ExecutePythonRequest(python_code=signed_blob)
# Should raise since the code has been modified.
self.assertRaises(rdfvalue.DecodeError, self.RunAction,
"ExecutePython", request)
# Lets also adjust the hash.
signed_blob.digest = hashlib.sha256(signed_blob.data).digest()
request = rdfvalue.ExecutePythonRequest(python_code=signed_blob)
self.assertRaises(rdfvalue.DecodeError, self.RunAction,
"ExecutePython", request)
# Make sure the code never ran.
self.assertEqual(utils.TEST_VAL, "original")
def testExecutePythonEnvironment(self):
"""Test the basic ExecutePython action."""
python_code = """
import StringIO
import uu
def decode(encoded):
# Use the import (uu) inside a function. This will fail if the environment
# for exec is not set up properly.
i = StringIO.StringIO(s)
o = StringIO.StringIO()
uu.decode(i, o)
return o.getvalue()
s = "626567696e20363636202d0a2c3226354c3b265c4035565d523b2630410a200a656e640a"
s = s.decode("hex")
magic_return_str = decode(s)
"""
signed_blob = rdfvalue.SignedBlob()
signed_blob.Sign(python_code, self.signing_key)
request = rdfvalue.ExecutePythonRequest(python_code=signed_blob)
result = self.RunAction("ExecutePython", request)[0]
self.assertTrue(result.time_used > 0)
self.assertEqual(result.return_val, "Hello World!")
def testExecuteBrokenPython(self):
"""Test broken code raises back to the original flow."""
python_code = "raise ValueError"
signed_blob = rdfvalue.SignedBlob()
signed_blob.Sign(python_code, self.signing_key)
request = rdfvalue.ExecutePythonRequest(python_code=signed_blob)
self.assertRaises(ValueError, self.RunAction, "ExecutePython", request)
def testReturnVals(self):
"""Test return values."""
python_code = "magic_return_str = 'return string'"
signed_blob = rdfvalue.SignedBlob()
signed_blob.Sign(python_code, self.signing_key)
request = rdfvalue.ExecutePythonRequest(python_code=signed_blob)
result = self.RunAction("ExecutePython", request)[0]
self.assertEqual(result.return_val, "return string")
def testWrongKey(self):
"""Test return values."""
python_code = "print 'test'"
# Generate a test valid RSA key that isn't the real one.
signing_key = rdfvalue.PEMPrivateKey.GenKey(2048, 65537)
signed_blob = rdfvalue.SignedBlob()
signed_blob.Sign(python_code, signing_key)
request = rdfvalue.ExecutePythonRequest(python_code=signed_blob)
self.assertRaises(rdfvalue.DecodeError, self.RunAction,
"ExecutePython", request)
def testExecutePython(self):
"""Test the basic ExecutePython action."""
utils.TEST_VAL = "original"
python_code = "utils.TEST_VAL = 'modified'"
signed_blob = rdfvalue.SignedBlob()
signed_blob.Sign(python_code, self.signing_key)
request = rdfvalue.ExecutePythonRequest(python_code=signed_blob)
result = self.RunAction("ExecutePython", request)[0]
self.assertTrue(result.time_used > 0)
self.assertEqual(result.return_val, "")
self.assertEqual(utils.TEST_VAL, "modified")
def testArgs(self):
"""Test passing arguments."""
utils.TEST_VAL = "original"
python_code = """
magic_return_str = py_args['test']
utils.TEST_VAL = py_args[43]
"""
signed_blob = rdfvalue.SignedBlob()
signed_blob.Sign(python_code, self.signing_key)
pdict = rdfvalue.Dict({"test": "dict_arg", 43: "dict_arg2"})
request = rdfvalue.ExecutePythonRequest(python_code=signed_blob,
py_args=pdict)
result = self.RunAction("ExecutePython", request)[0]
self.assertEqual(result.return_val, "dict_arg")
self.assertEqual(utils.TEST_VAL, "dict_arg2")
def testStdoutHooking(self):
python_code = """
def f(n):
print "F called:", n
print "Calling f."
f(1)
print "Done."
"""
signed_blob = rdfvalue.SignedBlob()
signed_blob.Sign(python_code, self.signing_key)
request = rdfvalue.ExecutePythonRequest(python_code=signed_blob)
result = self.RunAction("ExecutePython", request)[0]
self.assertTrue(result.time_used > 0)
self.assertEqual(result.return_val, "Calling f.\nF called: 1\nDone.\n")
def testProgress(self):
python_code = """
def f():
# This should also work inside a function.
Progress()
f()
Progress()
print "Done."
"""
signed_blob = rdfvalue.SignedBlob()
signed_blob.Sign(python_code, self.signing_key)
request = rdfvalue.ExecutePythonRequest(python_code=signed_blob)
result = self.RunAction("ExecutePython", request)[0]
self.assertTrue(result.time_used > 0)
self.assertEqual(result.return_val, "Done.\n")