def testDelayedCallState(self): """Tests the ability to delay a CallState invocation.""" with test_lib.FakeTime(10000): client_mock = ClientMock() client_mock = test_lib.MockClient(self.client_id, client_mock, token=self.token) worker_mock = test_lib.MockWorker(check_flow_errors=True, token=self.token) flow.GRRFlow.StartFlow( client_id=self.client_id, flow_name="DelayedCallStateFlow", token=self.token) self.Work(client_mock, worker_mock) # We should have done the first CallState so far. self.assertEqual(DelayedCallStateFlow.flow_ran, 1) with test_lib.FakeTime(10050): # 50 seconds more is not enough. self.Work(client_mock, worker_mock) self.assertEqual(DelayedCallStateFlow.flow_ran, 1) with test_lib.FakeTime(10100): # But 100 is. self.Work(client_mock, worker_mock) self.assertEqual(DelayedCallStateFlow.flow_ran, 2)
def Run(self): with test_lib.FakeTime(42): flow_urn = flow.GRRFlow.StartFlow( flow_name=processes.ListProcesses.__name__, client_id=self.client_id, token=self.token) mock = test_lib.MockClient(self.client_id, None, token=self.token) while mock.Next(): pass replace = {flow_urn.Basename(): "W:ABCDEF"} manager = queue_manager.QueueManager(token=self.token) requests_responses = manager.FetchRequestsAndResponses(flow_urn) for request, responses in requests_responses: replace[str(request.request.task_id)] = "42" for response in responses: replace[str(response.task_id)] = "42" self.Check("ListFlowRequests", args=flow_plugin.ApiListFlowRequestsArgs( client_id=self.client_id.Basename(), flow_id=flow_urn.Basename()), replace=replace)
def Run(self): client_ids = self.SetupClients(1) client_id = client_ids[0] replace = {} with test_lib.FakeTime(42): flow_urn = flow.GRRFlow.StartFlow( client_id=client_id, flow_name=processes.ListProcesses.__name__, token=self.token) replace[flow_urn.Basename()] = "F:123456" # Here we emulate a mock client with no actions (None) that should produce # an error. mock = test_lib.MockClient(client_id, None, token=self.token) while mock.Next(): pass manager = queue_manager.QueueManager(token=self.token) requests_responses = manager.FetchRequestsAndResponses(flow_urn) for request, responses in requests_responses: replace[str(request.request.task_id)] = "42" for response in responses: replace[str(response.task_id)] = "43" self.Check("ListClientActionRequests", args=client_plugin.ApiListClientActionRequestsArgs( client_id=client_id.Basename()), replace=replace) self.Check("ListClientActionRequests", args=client_plugin.ApiListClientActionRequestsArgs( client_id=client_id.Basename(), fetch_responses=True), replace=replace)
def testCPULimitForFlows(self): """This tests that the client actions are limited properly.""" result = {} client_mock = CPULimitClientMock(result) client_mock = test_lib.MockClient(self.client_id, client_mock, token=self.token) client_mock.EnableResourceUsage(user_cpu_usage=[10], system_cpu_usage=[10], network_usage=[1000]) worker_obj = worker.GRRWorker(token=self.token) flow.GRRFlow.StartFlow(client_id=self.client_id, flow_name="CPULimitFlow", cpu_limit=1000, network_bytes_limit=10000, token=self.token) self._Process([client_mock], worker_obj) self.assertEqual(result["cpulimit"], [1000, 980, 960]) self.assertEqual(result["networklimit"], [10000, 9000, 8000]) return result
def setUp(self): super(GRRFuseTest, self).setUp() self.client_id = self.SetupClients(1)[0] self.client_name = str(self.client_id)[len("aff4:/"):] with aff4.FACTORY.Open(self.client_id, token=self.token, mode="rw") as fd: fd.Set(fd.Schema.SYSTEM("Linux")) kb = fd.Schema.KNOWLEDGE_BASE() fd.Set(kb) with aff4.FACTORY.Create(self.client_id.Add("fs/os"), aff4_standard.VFSDirectory, mode="rw", token=self.token) as fd: fd.Set(fd.Schema.PATHSPEC(path="/", pathtype="OS")) # Ignore cache so our tests always get client side updates. self.grr_fuse = fuse_mount.GRRFuse(root="/", token=self.token, ignore_cache=True) self.action_mock = action_mocks.ActionMock( admin.GetClientInfo, admin.GetConfiguration, admin.GetPlatformInfo, file_fingerprint.FingerprintFile, linux.EnumerateFilesystems, linux.EnumerateInterfaces, linux.EnumerateUsers, linux.GetInstallDate, searching.Find, standard.HashBuffer, standard.ListDirectory, standard.StatFile, standard.TransferBuffer, ) self.client_mock = test_lib.MockClient(self.client_id, self.action_mock, token=self.token) self.update_stubber = utils.Stubber(self.grr_fuse, "_RunAndWaitForVFSFileUpdate", self._RunAndWaitForVFSFileUpdate) self.update_stubber.Start() self.start_flow_stubber = utils.Stubber(flow_utils, "StartFlowAndWait", self.StartFlowAndWait) self.start_flow_stubber.Start()
def testWorkerPrioritization(self): """Test that flow priorities work on the worker side.""" result = [] client_mock = PriorityClientMock(result) client_mock = test_lib.MockClient(self.client_id, client_mock, token=self.token) worker_mock = test_lib.MockWorker(check_flow_errors=True, token=self.token) # Start some flows with different priorities. # pyformat: disable args = [ (rdf_flows.GrrMessage.Priority.LOW_PRIORITY, "low priority"), (rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, "medium priority"), (rdf_flows.GrrMessage.Priority.LOW_PRIORITY, "low priority2"), (rdf_flows.GrrMessage.Priority.HIGH_PRIORITY, "high priority"), (rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, "medium priority2") ] # pyformat: enable server_result = [] PriorityFlow.storage = server_result for (priority, msg) in args: flow.GRRFlow.StartFlow(client_id=self.client_id, flow_name="PriorityFlow", msg=msg, priority=priority, token=self.token) while True: # Run all the clients first so workers have messages to choose from. client_processed = 1 while client_processed: client_processed = client_mock.Next() # Now process the results, this should happen in the correct order. flows_run = [] for flow_run in worker_mock.Next(): flows_run.append(flow_run) if not flows_run: break # The flows should be run in order of priority. self.assertEqual(server_result[0:1], [u"high priority"]) self.assertEqual(sorted(server_result[1:3]), [u"medium priority", u"medium priority2"]) self.assertEqual(sorted(server_result[3:5]), [u"low priority", u"low priority2"])
def testClientPrioritization(self): """Test that flow priorities work on the client side.""" result = [] client_mock = PriorityClientMock(result) client_mock = test_lib.MockClient(self.client_id, client_mock, token=self.token) worker_mock = test_lib.MockWorker(check_flow_errors=True, token=self.token) # Start some flows with different priorities. # pyformat: disable args = [ (rdf_flows.GrrMessage.Priority.LOW_PRIORITY, "low priority"), (rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, "medium priority"), (rdf_flows.GrrMessage.Priority.LOW_PRIORITY, "low priority2"), (rdf_flows.GrrMessage.Priority.HIGH_PRIORITY, "high priority"), (rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, "medium priority2") ] # pyformat: enable for (priority, msg) in args: flow.GRRFlow.StartFlow(client_id=self.client_id, flow_name="PriorityFlow", msg=msg, priority=priority, token=self.token) while True: client_processed = client_mock.Next() flows_run = [] for flow_run in worker_mock.Next(): flows_run.append(flow_run) if client_processed == 0 and not flows_run: break # The flows should be run in order of priority. self.assertEqual(result[0:1], [u"high priority"]) self.assertEqual(sorted(result[1:3]), [u"medium priority", u"medium priority2"]) self.assertEqual(sorted(result[3:5]), [u"low priority", u"low priority2"])
def __init__(self, method_name=None): super(GRRFuseTest, self).__init__(method_name) # Set up just once for the whole test suite, since we don't have any # per-test setup to do. super(GRRFuseTest, self).setUp() self.client_name = str(self.client_id)[len("aff4:/"):] with aff4.FACTORY.Open(self.client_id, token=self.token, mode="rw") as fd: fd.Set(fd.Schema.SYSTEM("Linux")) kb = fd.Schema.KNOWLEDGE_BASE() fd.Set(kb) # Ignore cache so our tests always get client side updates. self.grr_fuse = fuse_mount.GRRFuse(root="/", token=self.token, ignore_cache=True) self.action_mock = test_lib.ActionMock("TransferBuffer", "StatFile", "Find", "HashFile", "HashBuffer", "UpdateVFSFile", "EnumerateInterfaces", "EnumerateFilesystems", "GetConfiguration", "GetConfig", "GetClientInfo", "GetInstallDate", "GetPlatformInfo", "EnumerateUsers", "ListDirectory") client_mock = test_lib.MockClient(self.client_id, self.action_mock, token=self.token) worker_mock = test_lib.MockWorker(check_flow_errors=True, token=self.token) # All the flows we've run so far. We'll check them for errors at the end of # each test. self.total_flows = set() # We add the thread as a class variable since we'll be referring to it in # the tearDownClass method, and we want all tests to share it. self.__class__.fake_server_thread = threading.Thread( target=self.RunFakeWorkerAndClient, args=(client_mock, worker_mock)) self.fake_server_thread.start()
def RunFlow(self, flow_name, **kwargs): result = {} client_mock = CPULimitClientMock(result) client_mock = test_lib.MockClient(self.client_id, client_mock, token=self.token) worker_mock = ResourcedWorker(check_flow_errors=True, token=self.token) flow.GRRFlow.StartFlow( client_id=self.client_id, flow_name=flow_name, token=self.token, **kwargs) while True: client_processed = client_mock.Next() flows_run = [] for flow_run in worker_mock.Next(): flows_run.append(flow_run) if client_processed == 0 and not flows_run: break return result
def testInspect(self): """Test the inspect UI.""" client_id = self.SetupClients(1)[0] self.RequestAndGrantClientApproval(client_id) flow.GRRFlow.StartFlow(client_id=client_id, flow_name=flow_discovery.Interrogate.__name__, token=self.token) mock = test_lib.MockClient(client_id, None, token=self.token) while mock.Next(): pass self.Open("/#/clients/%s/debug-requests" % client_id.Basename()) # Check that the we can see both requests and responses. self.WaitUntil(self.IsTextPresent, "GetPlatformInfo") self.WaitUntil(self.IsTextPresent, "GetConfig") self.WaitUntil(self.IsTextPresent, "EnumerateInterfaces") self.WaitUntil(self.IsTextPresent, "GENERIC_ERROR") self.WaitUntil(self.IsTextPresent, "STATUS") self.WaitUntil(self.IsTextPresent, "Task id")
def testInspect(self): """Test the inspect UI.""" with self.ACLChecksDisabled(): self.GrantClientApproval("C.0000000000000001") self.Open("/") self.Type("client_query", "0001") self.Click("client_query_submit") self.WaitUntilEqual(u"C.0000000000000001", self.GetText, "css=span[type=subject]") # Choose client 1 self.Click("css=td:contains('0001')") self.Click("css=a[grrtarget=LaunchFlows]") self.Click("css=#_Administrative ins") self.Click("css=a:contains(Interrogate)") self.Click("css=button.Launch") # Open the "Advanced" dropdown. self.Click("css=a[href='#HostAdvanced']") # Click on the "Debug client requests". self.Click("css=a[grrtarget=DebugClientRequestsView]") self.WaitUntil(self.IsElementPresent, "css=td:contains(GetPlatformInfo)") # Check that the we can see the requests in the table. for request in "GetPlatformInfo GetConfig EnumerateInterfaces".split(): self.assertTrue( self.IsElementPresent("css=td:contains(%s)" % request)) self.Click("css=td:contains(GetPlatformInfo)") # Check that the proto is rendered inside the tab. self.WaitUntil( self.IsElementPresent, "css=.tab-content td.proto_value:contains(GetPlatformInfo)") # Check that the request tab is currently selected. self.assertTrue( self.IsElementPresent("css=li.active:contains(Request)")) # Here we emulate a mock client with no actions (None) this should produce # an error. with self.ACLChecksDisabled(): mock = test_lib.MockClient( rdfvalue.ClientURN("C.0000000000000001"), None, token=self.token) while mock.Next(): pass # Now select the Responses tab: self.Click("css=li a:contains(Responses)") self.WaitUntil(self.IsElementPresent, "css=td:contains('flow:response:')") self.assertTrue( self.IsElementPresent( "css=.tab-content td.proto_value:contains(GENERIC_ERROR)")) self.assertTrue( self.IsElementPresent( "css=.tab-content td.proto_value:contains(STATUS)"))
def testCPULimitForHunts(self): worker_obj = worker.GRRWorker(token=self.token) client_ids = ["C.%016X" % i for i in xrange(10, 20)] result = {} client_mocks = [] for client_id in client_ids: client_mock = CPULimitClientMock(result) client_mock = test_lib.MockClient(rdf_client.ClientURN(client_id), client_mock, token=self.token) client_mock.EnableResourceUsage(user_cpu_usage=[10], system_cpu_usage=[10], network_usage=[1000]) client_mocks.append(client_mock) flow_runner_args = rdf_flows.FlowRunnerArgs(flow_name="CPULimitFlow") with hunts.GRRHunt.StartHunt(hunt_name="GenericHunt", flow_runner_args=flow_runner_args, cpu_limit=5000, network_bytes_limit=1000000, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() hunts.GRRHunt.StartClients(hunt.session_id, client_ids[:1]) self._Process(client_mocks, worker_obj) hunts.GRRHunt.StartClients(hunt.session_id, client_ids[1:2]) self._Process(client_mocks, worker_obj) hunts.GRRHunt.StartClients(hunt.session_id, client_ids[2:3]) self._Process(client_mocks, worker_obj) # The limiting factor here is the overall hunt limit of 5000 cpu # seconds. Clients that finish should decrease the remaining quota # and the following clients should get the reduced quota. self.assertEqual(result["cpulimit"], [ 5000.0, 4980.0, 4960.0, 4940.0, 4920.0, 4900.0, 4880.0, 4860.0, 4840.0 ]) self.assertEqual(result["networklimit"], [ 1000000L, 999000L, 998000L, 997000L, 996000L, 995000L, 994000L, 993000L, 992000L ]) result.clear() with hunts.GRRHunt.StartHunt(hunt_name="GenericHunt", flow_runner_args=flow_runner_args, per_client_cpu_limit=3000, per_client_network_limit_bytes=3000000, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() hunts.GRRHunt.StartClients(hunt.session_id, client_ids[:1]) self._Process(client_mocks, worker_obj) hunts.GRRHunt.StartClients(hunt.session_id, client_ids[1:2]) self._Process(client_mocks, worker_obj) hunts.GRRHunt.StartClients(hunt.session_id, client_ids[2:3]) self._Process(client_mocks, worker_obj) # This time, the per client limit is 3000s / 3000000 bytes. Every # client should get the same limit. self.assertEqual(result["cpulimit"], [ 3000.0, 2980.0, 2960.0, 3000.0, 2980.0, 2960.0, 3000.0, 2980.0, 2960.0 ]) self.assertEqual(result["networklimit"], [ 3000000, 2999000, 2998000, 3000000, 2999000, 2998000, 3000000, 2999000, 2998000 ]) result.clear() for client_mock in client_mocks: client_mock.EnableResourceUsage(user_cpu_usage=[500], system_cpu_usage=[500], network_usage=[1000000]) with hunts.GRRHunt.StartHunt(hunt_name="GenericHunt", flow_runner_args=flow_runner_args, per_client_cpu_limit=3000, cpu_limit=5000, per_client_network_limit_bytes=3000000, network_bytes_limit=5000000, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() hunts.GRRHunt.StartClients(hunt.session_id, client_ids[:1]) self._Process(client_mocks, worker_obj) hunts.GRRHunt.StartClients(hunt.session_id, client_ids[1:2]) self._Process(client_mocks, worker_obj) hunts.GRRHunt.StartClients(hunt.session_id, client_ids[2:3]) self._Process(client_mocks, worker_obj) # The first client gets the full per client limit of 3000s, and # uses all of it. The hunt has a limit of just 5000 total so the # second client gets started with a limit of 2000. It can only run # two of three states, the last client will not be started at all # due to out of quota. self.assertEqual(result["cpulimit"], [3000.0, 2000.0, 1000.0, 2000.0, 1000.0]) self.assertEqual(result["networklimit"], [3000000, 2000000, 1000000, 2000000, 1000000]) errors = list(hunt.GetClientsErrors()) self.assertEqual(len(errors), 2) # Client side out of cpu. self.assertIn("CPU limit exceeded", errors[0].log_message) # Server side out of cpu. self.assertIn("Out of CPU quota", errors[1].backtrace)