<s> """<STR_LIT>""" <EOL> from grr . lib import communicator <EOL> from grr . lib import config_lib <EOL> from grr . lib import data_store <EOL> from grr . lib import flags <EOL> from grr . lib import flow <EOL> from grr . lib import queue_manager <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 flows as rdf_flows <EOL> from grr . lib . rdfvalues import protodict as rdf_protodict <EOL> class SendingTestFlow ( flow . GRRFlow ) : <EOL> """<STR_LIT>""" <EOL> @ flow . StateHandler ( next_state = "<STR_LIT>" ) <EOL> def Start ( self ) : <EOL> for i in range ( <NUM_LIT:10> ) : <EOL> self . CallClient ( "<STR_LIT>" , <EOL> rdf_protodict . DataBlob ( string = "<STR_LIT>" % i ) , <EOL> data = str ( i ) , <EOL> next_state = "<STR_LIT>" ) <EOL> class GRRFEServerTest ( test_lib . FlowTestsBaseclass ) : <EOL> """<STR_LIT>""" <EOL> string = "<STR_LIT>" <EOL> message_expiry_time = <NUM_LIT:100> <EOL> def setUp ( self ) : <EOL> """<STR_LIT>""" <EOL> super ( GRRFEServerTest , self ) . setUp ( ) <EOL> self . config_overrider = test_lib . ConfigOverrider ( { <EOL> "<STR_LIT>" : [ utils . SmartStr ( <EOL> test_lib . WellKnownSessionTest . well_known_session_id . FlowName ( ) ) ] , <EOL> "<STR_LIT>" : <NUM_LIT:10> } ) <EOL> self . config_overrider . Start ( ) <EOL> prefix = "<STR_LIT>" % self . _testMethodName <EOL> self . server = flow . FrontEndServer ( <EOL> certificate = config_lib . CONFIG [ "<STR_LIT>" ] , <EOL> private_key = config_lib . CONFIG [ "<STR_LIT>" ] , <EOL> message_expiry_time = self . message_expiry_time , <EOL> threadpool_prefix = prefix ) <EOL> def tearDown ( self ) : <EOL> super ( GRRFEServerTest , self ) . tearDown ( ) <EOL> self . config_overrider . Stop ( ) <EOL> def testReceiveMessages ( self ) : <EOL> """<STR_LIT>""" <EOL> flow_obj = self . FlowSetup ( "<STR_LIT>" ) <EOL> session_id = flow_obj . session_id <EOL> messages = [ rdf_flows . GrrMessage ( request_id = <NUM_LIT:1> , <EOL> response_id = i , <EOL> session_id = session_id , <EOL> payload = rdfvalue . RDFInteger ( i ) ) <EOL> for i in range ( <NUM_LIT:1> , <NUM_LIT:10> ) ] <EOL> self . server . ReceiveMessages ( self . client_id , messages ) <EOL> manager = queue_manager . QueueManager ( token = self . token ) <EOL> tasks_on_client_queue = manager . Query ( self . client_id , <NUM_LIT:100> ) <EOL> self . assertEqual ( len ( tasks_on_client_queue ) , <NUM_LIT:1> ) <EOL> for message in messages : <EOL> stored_message , _ = data_store . DB . Resolve ( <EOL> session_id . Add ( "<STR_LIT>" ) , <EOL> manager . FLOW_RESPONSE_TEMPLATE % ( <NUM_LIT:1> , message . response_id ) , <EOL> token = self . token ) <EOL> stored_message = rdf_flows . GrrMessage ( stored_message ) <EOL> self . assertRDFValuesEqual ( stored_message , message ) <EOL> return messages <EOL> def testReceiveMessagesWithStatus ( self ) : <EOL> """<STR_LIT>""" <EOL> flow_obj = self . FlowSetup ( "<STR_LIT>" ) <EOL> session_id = flow_obj . session_id <EOL> messages = [ rdf_flows . GrrMessage ( request_id = <NUM_LIT:1> , <EOL> response_id = i , <EOL> session_id = session_id , <EOL> payload = rdfvalue . RDFInteger ( i ) , <EOL> task_id = <NUM_LIT:15> ) <EOL> for i in range ( <NUM_LIT:1> , <NUM_LIT:10> ) ] <EOL> status = rdf_flows . GrrStatus ( status = rdf_flows . GrrStatus . ReturnedStatus . OK ) <EOL> messages . append ( rdf_flows . GrrMessage ( <EOL> request_id = <NUM_LIT:1> , response_id = len ( messages ) + <NUM_LIT:1> , task_id = <NUM_LIT:15> , <EOL> session_id = messages [ <NUM_LIT:0> ] . session_id , payload = status , <EOL> type = rdf_flows . GrrMessage . Type . STATUS ) ) <EOL> self . server . ReceiveMessages ( self . client_id , messages ) <EOL> manager = queue_manager . QueueManager ( token = self . token ) <EOL> tasks_on_client_queue = manager . Query ( self . client_id , <NUM_LIT:100> ) <EOL> self . assertEqual ( len ( tasks_on_client_queue ) , <NUM_LIT:1> ) <EOL> for message in messages : <EOL> stored_message , _ = data_store . DB . Resolve ( <EOL> session_id . Add ( "<STR_LIT>" ) , <EOL> manager . FLOW_RESPONSE_TEMPLATE % ( <NUM_LIT:1> , message . response_id ) , <EOL> token = self . token ) <EOL> stored_message = rdf_flows . GrrMessage ( stored_message ) <EOL> self . assertRDFValuesEqual ( stored_message , message ) <EOL> def testWellKnownFlows ( self ) : <EOL> """<STR_LIT>""" <EOL> test_lib . WellKnownSessionTest . messages = [ ] <EOL> session_id = test_lib . WellKnownSessionTest . well_known_session_id <EOL> messages = [ rdf_flows . GrrMessage ( request_id = <NUM_LIT:0> , <EOL> response_id = <NUM_LIT:0> , <EOL> session_id = session_id , <EOL> payload = rdfvalue . RDFInteger ( i ) ) <EOL> for i in range ( <NUM_LIT:1> , <NUM_LIT:10> ) ] <EOL> self . server . ReceiveMessages ( self . client_id , messages ) <EOL> self . server . thread_pool . Join ( ) <EOL> test_lib . WellKnownSessionTest . messages . sort ( ) <EOL> self . assertEqual ( test_lib . WellKnownSessionTest . messages , <EOL> list ( range ( <NUM_LIT:1> , <NUM_LIT:10> ) ) ) <EOL> self . assertEqual ( [ ] , data_store . DB . ResolvePrefix ( self . client_id , "<STR_LIT>" , <EOL> token = self . token ) ) <EOL> def testWellKnownFlowsRemote ( self ) : <EOL> """<STR_LIT>""" <EOL> test_lib . WellKnownSessionTest . messages = [ ] <EOL> session_id = test_lib . WellKnownSessionTest . well_known_session_id <EOL> messages = [ rdf_flows . GrrMessage ( request_id = <NUM_LIT:0> , <EOL> response_id = <NUM_LIT:0> , <EOL> session_id = session_id , <EOL> payload = rdfvalue . RDFInteger ( i ) ) <EOL> for i in range ( <NUM_LIT:1> , <NUM_LIT:10> ) ] <EOL> self . server . well_known_flows = { } <EOL> self . server . ReceiveMessages ( self . client_id , messages ) <EOL> self . server . thread_pool . Join ( ) <EOL> self . assertEqual ( test_lib . WellKnownSessionTest . messages , [ ] ) <EOL> self . assertEqual ( [ ] , data_store . DB . ResolvePrefix ( self . client_id , "<STR_LIT>" , <EOL> token = self . token ) ) <EOL> queued_messages = [ ] <EOL> for predicate , _ , _ in data_store . DB . ResolvePrefix ( <EOL> session_id . Add ( "<STR_LIT>" ) , "<STR_LIT>" , token = self . token ) : <EOL> queued_messages . append ( predicate ) <EOL> self . assertEqual ( len ( queued_messages ) , <NUM_LIT:9> ) <EOL> def testWellKnownFlowsNotifications ( self ) : <EOL> test_lib . WellKnownSessionTest . messages = [ ] <EOL> test_lib . WellKnownSessionTest2 . messages = [ ] <EOL> session_id1 = test_lib . WellKnownSessionTest . well_known_session_id <EOL> session_id2 = test_lib . WellKnownSessionTest2 . well_known_session_id <EOL> messages = [ ] <EOL> for i in range ( <NUM_LIT:1> , <NUM_LIT:5> ) : <EOL> messages . append ( rdf_flows . GrrMessage ( request_id = <NUM_LIT:0> , <EOL> response_id = <NUM_LIT:0> , <EOL> session_id = session_id1 , <EOL> payload = rdfvalue . RDFInteger ( i ) ) ) <EOL> messages . append ( rdf_flows . GrrMessage ( request_id = <NUM_LIT:0> , <EOL> response_id = <NUM_LIT:0> , <EOL> session_id = session_id2 , <EOL> payload = rdfvalue . RDFInteger ( i ) ) ) <EOL> self . assertIn ( session_id1 . FlowName ( ) , self . server . well_known_flows ) <EOL> self . assertNotIn ( session_id2 . FlowName ( ) , self . server . well_known_flows ) <EOL> self . server . ReceiveMessages ( self . client_id , messages ) <EOL> self . server . thread_pool . Join ( ) <EOL> test_lib . WellKnownSessionTest . messages . sort ( ) <EOL> self . assertEqual ( test_lib . WellKnownSessionTest . messages , <EOL> list ( range ( <NUM_LIT:1> , <NUM_LIT:5> ) ) ) <EOL> self . assertEqual ( test_lib . WellKnownSessionTest2 . messages , [ ] ) <EOL> manager = queue_manager . WellKnownQueueManager ( token = self . token ) <EOL> notifications = manager . GetNotificationsForAllShards ( session_id1 . Queue ( ) ) <EOL> responses = list ( manager . FetchRequestsAndResponses ( session_id1 ) ) <EOL> self . assertEqual ( responses , [ ] ) <EOL> self . assertNotIn ( session_id1 , [ notification . session_id <EOL> for notification in notifications ] ) <EOL> responses = list ( manager . FetchRequestsAndResponses ( session_id2 ) ) <EOL> self . assertEqual ( len ( responses ) , <NUM_LIT:4> ) <EOL> self . assertIn ( session_id2 , [ notification . session_id <EOL> for notification in notifications ] ) <EOL> def testDrainUpdateSessionRequestStates ( self ) : <EOL> """<STR_LIT>""" <EOL> flow_obj = self . FlowSetup ( "<STR_LIT>" ) <EOL> session_id = flow_obj . session_id <EOL> manager = queue_manager . QueueManager ( token = self . token ) <EOL> tasks = manager . Query ( self . client_id , <NUM_LIT:100> ) <EOL> self . assertEqual ( len ( tasks ) , <NUM_LIT:10> ) <EOL> for task in tasks : <EOL> request_id = task . request_id <EOL> request_state , _ = data_store . DB . Resolve ( <EOL> session_id . Add ( "<STR_LIT:state>" ) , <EOL> manager . FLOW_REQUEST_TEMPLATE % request_id , <EOL> token = self . token ) <EOL> request_state = rdf_flows . RequestState ( request_state ) <EOL> self . assertEqual ( request_state . request . task_id , task . task_id ) <EOL> response = rdf_flows . MessageList ( ) <EOL> self . server . DrainTaskSchedulerQueueForClient ( <EOL> self . client_id , <NUM_LIT:5> , response ) <EOL> self . assertEqual ( len ( response . job ) , <NUM_LIT:5> ) <EOL> for i in range ( <NUM_LIT:4> ) : <EOL> self . assertEqual ( response . job [ i ] . session_id , session_id ) <EOL> self . assertEqual ( response . job [ i ] . name , "<STR_LIT>" ) <EOL> def testUpdateAndCheckIfShouldThrottle ( self ) : <EOL> self . server . SetThrottleBundlesRatio ( <NUM_LIT:1.0> ) <EOL> self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT:0> ) <EOL> self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT:10> ) <EOL> self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT:20> ) <EOL> self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT:30> ) <EOL> self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT:50> ) <EOL> self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . server . SetThrottleBundlesRatio ( <NUM_LIT> ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . assertEqual ( result , True ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . assertEqual ( result , True ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . assertEqual ( result , True ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT:100> ) <EOL> self . assertEqual ( result , False ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . assertEqual ( result , True ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . assertEqual ( result , True ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . assertEqual ( result , True ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . assertEqual ( result , False ) <EOL> self . server . SetThrottleBundlesRatio ( <NUM_LIT:0> ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . assertEqual ( result , True ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . assertEqual ( result , True ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . assertEqual ( result , True ) <EOL> self . server . SetThrottleBundlesRatio ( None ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . assertEqual ( result , False ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . assertEqual ( result , False ) <EOL> result = self . server . UpdateAndCheckIfShouldThrottle ( <NUM_LIT> ) <EOL> self . assertEqual ( result , False ) <EOL> def testHandleMessageBundle ( self ) : <EOL> """<STR_LIT>""" <EOL> client_id = rdf_client . ClientURN ( "<STR_LIT>" + "<STR_LIT:2>" * <NUM_LIT:16> ) <EOL> class MockCommunicator ( object ) : <EOL> """<STR_LIT>""" <EOL> def DecodeMessages ( self , * unused_args ) : <EOL> """<STR_LIT>""" <EOL> return ( [ ] , client_id , <NUM_LIT:100> ) <EOL> def EncodeMessages ( self , * unused_args , ** unused_kw ) : <EOL> """<STR_LIT>""" <EOL> raise communicator . UnknownClientCert ( ) <EOL> self . server . _communicator = MockCommunicator ( ) <EOL> request_comms = rdf_flows . ClientCommunication ( ) <EOL> self . assertRaises ( communicator . UnknownClientCert , <EOL> self . server . HandleMessageBundles , request_comms , <NUM_LIT:2> ) <EOL> flow . GRRFlow . StartFlow ( client_id = client_id , flow_name = "<STR_LIT>" , <EOL> message_count = <NUM_LIT:1> , token = self . token ) <EOL> manager = queue_manager . QueueManager ( token = self . token ) <EOL> tasks = manager . Query ( client_id , limit = <NUM_LIT:100> ) <EOL> self . assertRaises ( communicator . UnknownClientCert , <EOL> self . server . HandleMessageBundles , request_comms , <NUM_LIT:2> ) <EOL> new_tasks = manager . Query ( client_id , limit = <NUM_LIT:100> ) <EOL> lease_time = ( new_tasks [ <NUM_LIT:0> ] . eta - tasks [ <NUM_LIT:0> ] . eta ) / <NUM_LIT> <EOL> self . assertLess ( lease_time , <NUM_LIT:1> ) <EOL> self . assertEqual ( tasks [ <NUM_LIT:0> ] . task_ttl - new_tasks [ <NUM_LIT:0> ] . task_ttl , <NUM_LIT:1> ) <EOL> def _ScheduleResponseAndStatus ( self , client_id , flow_id ) : <EOL> with queue_manager . QueueManager ( token = self . token ) as flow_manager : <EOL> flow_manager . QueueResponse ( flow_id , rdf_flows . GrrMessage ( <EOL> source = client_id , <EOL> session_id = flow_id , <EOL> payload = rdf_protodict . DataBlob ( string = "<STR_LIT>" ) , <EOL> request_id = <NUM_LIT:1> , <EOL> response_id = <NUM_LIT:1> ) ) <EOL> flow_manager . QueueResponse ( flow_id , rdf_flows . GrrMessage ( <EOL> source = client_id , <EOL> session_id = flow_id , <EOL> payload = rdf_flows . GrrStatus ( <EOL> status = rdf_flows . GrrStatus . ReturnedStatus . OK ) , <EOL> request_id = <NUM_LIT:1> , response_id = <NUM_LIT:2> , <EOL> type = rdf_flows . GrrMessage . Type . STATUS ) ) <EOL> def testHandleClientMessageRetransmission ( self ) : <EOL> """<STR_LIT>""" <EOL> client_id = self . SetupClients ( <NUM_LIT:1> ) [ <NUM_LIT:0> ] <EOL> base_time = <NUM_LIT:1000> <EOL> msgs_recvd = [ ] <EOL> default_ttl = rdf_flows . GrrMessage ( ) . task_ttl <EOL> with test_lib . FakeTime ( base_time ) : <EOL> flow . GRRFlow . StartFlow ( client_id = client_id , flow_name = "<STR_LIT>" , <EOL> message_count = <NUM_LIT:1> , token = self . token ) <EOL> for i in range ( default_ttl ) : <EOL> with test_lib . FakeTime ( base_time + i * ( self . message_expiry_time + <NUM_LIT:1> ) ) : <EOL> tasks = self . server . DrainTaskSchedulerQueueForClient ( <EOL> client_id , <NUM_LIT> , rdf_flows . MessageList ( ) ) <EOL> msgs_recvd . append ( tasks ) <EOL> self . assertEqual ( map ( bool , msgs_recvd ) , <EOL> [ True ] * ( rdf_flows . GrrMessage ( ) . task_ttl - <NUM_LIT:1> ) + [ False ] ) <EOL> if default_ttl <= <NUM_LIT:3> : <EOL> self . fail ( "<STR_LIT>" ) <EOL> msgs_recvd = [ ] <EOL> with test_lib . FakeTime ( base_time ) : <EOL> flow_id = flow . GRRFlow . StartFlow ( <EOL> client_id = client_id , flow_name = "<STR_LIT>" , <EOL> message_count = <NUM_LIT:1> , token = self . token ) <EOL> for i in range ( default_ttl ) : <EOL> if i == <NUM_LIT:2> : <EOL> self . _ScheduleResponseAndStatus ( client_id , flow_id ) <EOL> with test_lib . FakeTime ( base_time + i * ( self . message_expiry_time + <NUM_LIT:1> ) ) : <EOL> tasks = self . server . DrainTaskSchedulerQueueForClient ( <EOL> client_id , <NUM_LIT> , rdf_flows . MessageList ( ) ) <EOL> msgs_recvd . append ( tasks ) <EOL> if not tasks : <EOL> new_tasks = queue_manager . QueueManager ( token = self . token ) . Query ( <EOL> queue = rdf_client . ClientURN ( client_id ) . Queue ( ) , limit = <NUM_LIT:1000> ) <EOL> self . assertEqual ( len ( new_tasks ) , <NUM_LIT:0> ) <EOL> self . assertEqual ( <EOL> map ( bool , msgs_recvd ) , <EOL> [ True ] * <NUM_LIT:2> + [ False ] * ( rdf_flows . GrrMessage ( ) . task_ttl - <NUM_LIT:2> ) ) <EOL> def main ( args ) : <EOL> test_lib . main ( args ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> flags . StartMain ( main ) </s>
