<s> """<STR_LIT>""" <EOL> import gzip <EOL> import hashlib <EOL> import os <EOL> import time <EOL> from grr . client . client_actions import standard <EOL> from grr . lib import action_mocks <EOL> from grr . lib import config_lib <EOL> from grr . lib import flags <EOL> from grr . lib import rdfvalue <EOL> from grr . lib import test_lib <EOL> from grr . lib import utils <EOL> from grr . lib . rdfvalues import client as rdf_client <EOL> from grr . lib . rdfvalues import crypto as rdf_crypto <EOL> from grr . lib . rdfvalues import flows as rdf_flows <EOL> from grr . lib . rdfvalues import paths as rdf_paths <EOL> from grr . lib . rdfvalues import protodict as rdf_protodict <EOL> class TestExecutePython ( test_lib . EmptyActionTest ) : <EOL> """<STR_LIT>""" <EOL> def setUp ( self ) : <EOL> super ( TestExecutePython , self ) . setUp ( ) <EOL> self . signing_key = config_lib . CONFIG [ <EOL> "<STR_LIT>" ] <EOL> def testExecutePython ( self ) : <EOL> """<STR_LIT>""" <EOL> utils . TEST_VAL = "<STR_LIT>" <EOL> python_code = "<STR_LIT>" <EOL> signed_blob = rdf_crypto . SignedBlob ( ) <EOL> signed_blob . Sign ( python_code , self . signing_key ) <EOL> request = rdf_client . ExecutePythonRequest ( python_code = signed_blob ) <EOL> result = self . RunAction ( "<STR_LIT>" , request ) [ <NUM_LIT:0> ] <EOL> self . assertTrue ( result . time_used > <NUM_LIT:0> ) <EOL> self . assertEqual ( result . return_val , "<STR_LIT>" ) <EOL> self . assertEqual ( utils . TEST_VAL , "<STR_LIT>" ) <EOL> def testExecutePythonEnvironment ( self ) : <EOL> """<STR_LIT>""" <EOL> python_code = """<STR_LIT>""" <EOL> signed_blob = rdf_crypto . SignedBlob ( ) <EOL> signed_blob . Sign ( python_code , self . signing_key ) <EOL> request = rdf_client . ExecutePythonRequest ( python_code = signed_blob ) <EOL> result = self . RunAction ( "<STR_LIT>" , request ) [ <NUM_LIT:0> ] <EOL> self . assertTrue ( result . time_used > <NUM_LIT:0> ) <EOL> self . assertEqual ( result . return_val , "<STR_LIT>" ) <EOL> def testStdoutHooking ( self ) : <EOL> python_code = """<STR_LIT>""" <EOL> signed_blob = rdf_crypto . SignedBlob ( ) <EOL> signed_blob . Sign ( python_code , self . signing_key ) <EOL> request = rdf_client . ExecutePythonRequest ( python_code = signed_blob ) <EOL> result = self . RunAction ( "<STR_LIT>" , request ) [ <NUM_LIT:0> ] <EOL> self . assertTrue ( result . time_used > <NUM_LIT:0> ) <EOL> self . assertEqual ( result . return_val , "<STR_LIT>" ) <EOL> def testProgress ( self ) : <EOL> python_code = """<STR_LIT>""" <EOL> signed_blob = rdf_crypto . SignedBlob ( ) <EOL> signed_blob . Sign ( python_code , self . signing_key ) <EOL> request = rdf_client . ExecutePythonRequest ( python_code = signed_blob ) <EOL> result = self . RunAction ( "<STR_LIT>" , request ) [ <NUM_LIT:0> ] <EOL> self . assertTrue ( result . time_used > <NUM_LIT:0> ) <EOL> self . assertEqual ( result . return_val , "<STR_LIT>" ) <EOL> def testExecuteModifiedPython ( self ) : <EOL> """<STR_LIT>""" <EOL> utils . TEST_VAL = "<STR_LIT>" <EOL> python_code = "<STR_LIT>" <EOL> signed_blob = rdf_crypto . SignedBlob ( ) <EOL> signed_blob . Sign ( python_code , self . signing_key ) <EOL> signed_blob . data = "<STR_LIT>" <EOL> request = rdf_client . ExecutePythonRequest ( python_code = signed_blob ) <EOL> self . assertRaises ( rdfvalue . DecodeError , <EOL> self . RunAction , "<STR_LIT>" , request ) <EOL> signed_blob . digest = hashlib . sha256 ( signed_blob . data ) . digest ( ) <EOL> request = rdf_client . ExecutePythonRequest ( python_code = signed_blob ) <EOL> self . assertRaises ( rdfvalue . DecodeError , <EOL> self . RunAction , "<STR_LIT>" , request ) <EOL> self . assertEqual ( utils . TEST_VAL , "<STR_LIT>" ) <EOL> def testExecuteBrokenPython ( self ) : <EOL> """<STR_LIT>""" <EOL> python_code = "<STR_LIT>" <EOL> signed_blob = rdf_crypto . SignedBlob ( ) <EOL> signed_blob . Sign ( python_code , self . signing_key ) <EOL> request = rdf_client . ExecutePythonRequest ( python_code = signed_blob ) <EOL> self . assertRaises ( ValueError , <EOL> self . RunAction , "<STR_LIT>" , request ) <EOL> def testExecuteBinary ( self ) : <EOL> """<STR_LIT>""" <EOL> signed_blob = rdf_crypto . SignedBlob ( ) <EOL> signed_blob . Sign ( open ( "<STR_LIT>" ) . read ( ) , self . signing_key ) <EOL> writefile = utils . Join ( self . temp_dir , "<STR_LIT>" , "<STR_LIT>" ) <EOL> os . makedirs ( os . path . dirname ( writefile ) ) <EOL> request = rdf_client . ExecuteBinaryRequest ( executable = signed_blob , <EOL> args = [ __file__ ] , <EOL> write_path = writefile ) <EOL> result = self . RunAction ( "<STR_LIT>" , request ) [ <NUM_LIT:0> ] <EOL> self . assertTrue ( result . time_used > <NUM_LIT:0> ) <EOL> self . assertTrue ( __file__ in result . stdout ) <EOL> def testReturnVals ( self ) : <EOL> """<STR_LIT>""" <EOL> python_code = "<STR_LIT>" <EOL> signed_blob = rdf_crypto . SignedBlob ( ) <EOL> signed_blob . Sign ( python_code , self . signing_key ) <EOL> request = rdf_client . ExecutePythonRequest ( python_code = signed_blob ) <EOL> result = self . RunAction ( "<STR_LIT>" , request ) [ <NUM_LIT:0> ] <EOL> self . assertEqual ( result . return_val , "<STR_LIT>" ) <EOL> def testWrongKey ( self ) : <EOL> """<STR_LIT>""" <EOL> python_code = "<STR_LIT>" <EOL> signing_key = rdf_crypto . PEMPrivateKey . GenKey ( <NUM_LIT> , <NUM_LIT> ) <EOL> signed_blob = rdf_crypto . SignedBlob ( ) <EOL> signed_blob . Sign ( python_code , signing_key ) <EOL> request = rdf_client . ExecutePythonRequest ( python_code = signed_blob ) <EOL> self . assertRaises ( rdfvalue . DecodeError , self . RunAction , <EOL> "<STR_LIT>" , request ) <EOL> def testArgs ( self ) : <EOL> """<STR_LIT>""" <EOL> utils . TEST_VAL = "<STR_LIT>" <EOL> python_code = """<STR_LIT>""" <EOL> signed_blob = rdf_crypto . SignedBlob ( ) <EOL> signed_blob . Sign ( python_code , self . signing_key ) <EOL> pdict = rdf_protodict . Dict ( { "<STR_LIT:test>" : "<STR_LIT>" , <EOL> <NUM_LIT> : "<STR_LIT>" } ) <EOL> request = rdf_client . ExecutePythonRequest ( python_code = signed_blob , <EOL> py_args = pdict ) <EOL> result = self . RunAction ( "<STR_LIT>" , request ) [ <NUM_LIT:0> ] <EOL> self . assertEqual ( result . return_val , "<STR_LIT>" ) <EOL> self . assertEqual ( utils . TEST_VAL , "<STR_LIT>" ) <EOL> class TestCopyPathToFile ( test_lib . EmptyActionTest ) : <EOL> """<STR_LIT>""" <EOL> def setUp ( self ) : <EOL> super ( TestCopyPathToFile , self ) . setUp ( ) <EOL> self . path_in = os . path . join ( self . base_path , "<STR_LIT>" ) <EOL> self . hash_in = hashlib . sha1 ( open ( self . path_in ) . read ( ) ) . hexdigest ( ) <EOL> self . pathspec = rdf_paths . PathSpec ( <EOL> path = self . path_in , pathtype = rdf_paths . PathSpec . PathType . OS ) <EOL> def testCopyPathToFile ( self ) : <EOL> request = rdf_client . CopyPathToFileRequest ( offset = <NUM_LIT:0> , <EOL> length = <NUM_LIT:0> , <EOL> src_path = self . pathspec , <EOL> dest_dir = self . temp_dir , <EOL> gzip_output = False ) <EOL> result = self . RunAction ( "<STR_LIT>" , request ) [ <NUM_LIT:0> ] <EOL> hash_out = hashlib . sha1 ( open ( result . dest_path . path ) . read ( ) ) . hexdigest ( ) <EOL> self . assertEqual ( self . hash_in , hash_out ) <EOL> def testCopyPathToFileLimitLength ( self ) : <EOL> request = rdf_client . CopyPathToFileRequest ( offset = <NUM_LIT:0> , <EOL> length = <NUM_LIT> , <EOL> src_path = self . pathspec , <EOL> dest_dir = self . temp_dir , <EOL> gzip_output = False ) <EOL> result = self . RunAction ( "<STR_LIT>" , request ) [ <NUM_LIT:0> ] <EOL> output = open ( result . dest_path . path ) . read ( ) <EOL> self . assertEqual ( len ( output ) , <NUM_LIT> ) <EOL> def testCopyPathToFileOffsetandLimit ( self ) : <EOL> with open ( self . path_in ) as f : <EOL> f . seek ( <NUM_LIT> ) <EOL> out = f . read ( <NUM_LIT> ) <EOL> hash_in = hashlib . sha1 ( out ) . hexdigest ( ) <EOL> request = rdf_client . CopyPathToFileRequest ( offset = <NUM_LIT> , <EOL> length = <NUM_LIT> , <EOL> src_path = self . pathspec , <EOL> dest_dir = self . temp_dir , <EOL> gzip_output = False ) <EOL> result = self . RunAction ( "<STR_LIT>" , request ) [ <NUM_LIT:0> ] <EOL> output = open ( result . dest_path . path ) . read ( ) <EOL> self . assertEqual ( len ( output ) , <NUM_LIT> ) <EOL> hash_out = hashlib . sha1 ( output ) . hexdigest ( ) <EOL> self . assertEqual ( hash_in , hash_out ) <EOL> def testCopyPathToFileGzip ( self ) : <EOL> request = rdf_client . CopyPathToFileRequest ( offset = <NUM_LIT:0> , <EOL> length = <NUM_LIT:0> , <EOL> src_path = self . pathspec , <EOL> dest_dir = self . temp_dir , <EOL> gzip_output = True ) <EOL> result = self . RunAction ( "<STR_LIT>" , request ) [ <NUM_LIT:0> ] <EOL> self . assertEqual ( hashlib . sha1 ( <EOL> gzip . open ( result . dest_path . path ) . read ( ) ) . hexdigest ( ) , self . hash_in ) <EOL> def testCopyPathToFileLifetimeLimit ( self ) : <EOL> request = rdf_client . CopyPathToFileRequest ( offset = <NUM_LIT:0> , <EOL> length = <NUM_LIT> , <EOL> src_path = self . pathspec , <EOL> dest_dir = self . temp_dir , <EOL> gzip_output = False , <EOL> lifetime = <NUM_LIT:0.1> ) <EOL> result = self . RunAction ( "<STR_LIT>" , request ) [ <NUM_LIT:0> ] <EOL> self . assertTrue ( os . path . exists ( result . dest_path . path ) ) <EOL> time . sleep ( <NUM_LIT:1> ) <EOL> self . assertFalse ( os . path . exists ( result . dest_path . path ) ) <EOL> class TestNetworkByteLimits ( test_lib . EmptyActionTest ) : <EOL> """<STR_LIT>""" <EOL> def setUp ( self ) : <EOL> super ( TestNetworkByteLimits , self ) . setUp ( ) <EOL> pathspec = rdf_paths . PathSpec ( path = "<STR_LIT>" , <EOL> pathtype = rdf_paths . PathSpec . PathType . OS ) <EOL> self . buffer_ref = rdf_client . BufferReference ( pathspec = pathspec , length = <NUM_LIT> ) <EOL> self . data = "<STR_LIT:X>" * <NUM_LIT> <EOL> self . old_read = standard . vfs . ReadVFS <EOL> standard . vfs . ReadVFS = lambda x , y , z , progress_callback = None : self . data <EOL> self . transfer_buf = action_mocks . ActionMock ( "<STR_LIT>" ) <EOL> def testTransferNetworkByteLimitError ( self ) : <EOL> message = rdf_flows . GrrMessage ( name = "<STR_LIT>" , <EOL> payload = self . buffer_ref , <EOL> network_bytes_limit = <NUM_LIT> ) <EOL> responses = self . transfer_buf . HandleMessage ( message ) <EOL> client_alert = responses [ <NUM_LIT:0> ] . payload <EOL> self . assertTrue ( "<STR_LIT>" in str ( client_alert ) ) <EOL> status = responses [ <NUM_LIT:1> ] . payload <EOL> self . assertTrue ( "<STR_LIT>" <EOL> in str ( status . backtrace ) ) <EOL> self . assertEqual ( status . status , <EOL> rdf_flows . GrrStatus . ReturnedStatus . NETWORK_LIMIT_EXCEEDED ) <EOL> def testTransferNetworkByteLimit ( self ) : <EOL> message = rdf_flows . GrrMessage ( name = "<STR_LIT>" , <EOL> payload = self . buffer_ref , <EOL> network_bytes_limit = <NUM_LIT> ) <EOL> responses = self . transfer_buf . HandleMessage ( message ) <EOL> for response in responses : <EOL> if isinstance ( response , rdf_flows . GrrStatus ) : <EOL> self . assertEqual ( response . payload . status , <EOL> rdf_flows . GrrStatus . ReturnedStatus . OK ) <EOL> def tearDown ( self ) : <EOL> super ( TestNetworkByteLimits , self ) . tearDown ( ) <EOL> standard . vfs . ReadVFS = self . old_read <EOL> def main ( argv ) : <EOL> test_lib . main ( argv ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> flags . StartMain ( main ) </s>
