def testCountHuntFlowsIgnoresSubflows(self): hunt_obj = rdf_hunt_objects.Hunt(description="foo") self.db.WriteHuntObject(hunt_obj) hunt_id = hunt_obj.hunt_id _, flow_id = self._SetupHuntClientAndFlow( hunt_id=hunt_id, flow_state=rdf_flow_objects.Flow.FlowState.RUNNING) # Whatever state the subflow is in, it should be ignored. self._SetupHuntClientAndFlow( hunt_id=hunt_id, flow_id=flow.RandomFlowId(), parent_flow_id=flow_id, flow_state=rdf_flow_objects.Flow.FlowState.ERROR) self._SetupHuntClientAndFlow( hunt_id=hunt_id, flow_id=flow.RandomFlowId(), parent_flow_id=flow_id, flow_state=rdf_flow_objects.Flow.FlowState.FINISHED) self._SetupHuntClientAndFlow( hunt_id=hunt_id, flow_id=flow.RandomFlowId(), parent_flow_id=flow_id, flow_state=rdf_flow_objects.Flow.FlowState.RUNNING) self.assertEqual(self.db.CountHuntFlows(hunt_id), 1)
def testReadHuntFlowsIgnoresSubflows(self): hunt_obj = rdf_hunt_objects.Hunt(description="foo") self.db.WriteHuntObject(hunt_obj) hunt_id = hunt_obj.hunt_id _, flow_id = self._SetupHuntClientAndFlow( hunt_id=hunt_id, flow_state=rdf_flow_objects.Flow.FlowState.RUNNING) # Whatever state the subflow is in, it should be ignored. self._SetupHuntClientAndFlow( hunt_id=hunt_id, flow_id=flow.RandomFlowId(), parent_flow_id=flow_id, flow_state=rdf_flow_objects.Flow.FlowState.ERROR) self._SetupHuntClientAndFlow( hunt_id=hunt_id, flow_id=flow.RandomFlowId(), parent_flow_id=flow_id, flow_state=rdf_flow_objects.Flow.FlowState.FINISHED) self._SetupHuntClientAndFlow( hunt_id=hunt_id, flow_id=flow.RandomFlowId(), parent_flow_id=flow_id, flow_state=rdf_flow_objects.Flow.FlowState.RUNNING) for state, expceted_results in [ (db.HuntFlowsCondition.COMPLETED_FLOWS_ONLY, 0), (db.HuntFlowsCondition.SUCCEEDED_FLOWS_ONLY, 0), (db.HuntFlowsCondition.FLOWS_IN_PROGRESS_ONLY, 1) ]: results = self.db.ReadHuntFlows(hunt_id, 0, 10, filter_condition=state) self.assertLen(results, expceted_results)
def testHuntOverviewGetsUpdatedWhenHuntChanges(self): hunt_urn = self.CreateSampleHunt() hunt_id = hunt_urn.Basename() if data_store.RelationalDBEnabled(): client_id = self.SetupClient(0).Basename() rdf_flow = rdf_flow_objects.Flow( client_id=client_id, flow_id=flow.RandomFlowId(), parent_hunt_id=hunt_id, create_time=rdfvalue.RDFDatetime.Now()) rdf_flow.cpu_time_used.user_cpu_time = 5000 rdf_flow.network_bytes_sent = 1000000 data_store.REL_DB.WriteFlowObject(rdf_flow) else: with aff4.FACTORY.Open(hunt_urn, mode="rw") as hunt_obj: hunt_stats = hunt_obj.context.usage_stats hunt_stats.user_cpu_stats.sum = 5000 hunt_stats.network_bytes_sent_stats.sum = 1000000 self.Open("/") # Ensure auto-refresh updates happen every second. self.GetJavaScriptValue( "grrUi.hunt.huntOverviewDirective.setAutoRefreshInterval(1000);") self.Click("css=a[grrtarget=hunts]") self.Click("css=td:contains('%s')" % hunt_id) self.WaitUntil(self.IsTextPresent, "1h 23m 20s") self.WaitUntil(self.IsTextPresent, "976.6KiB") if data_store.RelationalDBEnabled(): client_id = self.SetupClient(1).Basename() rdf_flow = rdf_flow_objects.Flow( client_id=client_id, flow_id=flow.RandomFlowId(), parent_hunt_id=hunt_id, create_time=rdfvalue.RDFDatetime.Now()) rdf_flow.cpu_time_used.user_cpu_time = 1000 rdf_flow.network_bytes_sent = 10000000 data_store.REL_DB.WriteFlowObject(rdf_flow) else: with aff4.FACTORY.Open(hunt_urn, mode="rw") as hunt_obj: hunt_stats = hunt_obj.context.usage_stats hunt_stats.user_cpu_stats.sum = 6000 hunt_stats.network_bytes_sent_stats.sum = 11000000 self.WaitUntil(self.IsTextPresent, "1h 40m") self.WaitUntil(self.IsTextPresent, "10.5MiB")
def testHuntOverviewShowsStats(self): """Test the detailed client view works.""" hunt_urn = self.CreateSampleHunt() hunt_id = hunt_urn.Basename() if data_store.RelationalDBEnabled(): client_id = self.SetupClient(0).Basename() rdf_flow = rdf_flow_objects.Flow( client_id=client_id, flow_id=flow.RandomFlowId(), parent_hunt_id=hunt_id, create_time=rdfvalue.RDFDatetime.Now()) rdf_flow.cpu_time_used.user_cpu_time = 5000 rdf_flow.network_bytes_sent = 1000000 data_store.REL_DB.WriteFlowObject(rdf_flow) else: with aff4.FACTORY.Open(hunt_urn, mode="rw") as hunt_obj: hunt_stats = hunt_obj.context.usage_stats hunt_stats.user_cpu_stats.sum = 5000 hunt_stats.network_bytes_sent_stats.sum = 1000000 # Open up and click on View Hunts then the first Hunt. self.Open("/") self.WaitUntil(self.IsElementPresent, "client_query") self.Click("css=a[grrtarget=hunts]") self.WaitUntil(self.IsTextPresent, hunt_id) self.Click("css=td:contains('%s')" % hunt_id) # Click the Overview Tab and check that the stats are present. self.Click("css=li[heading=Overview]") self.WaitUntil(self.IsTextPresent, "1h 23m 20s") self.WaitUntil(self.IsTextPresent, "976.6KiB")
def testHuntOverviewShowsStats(self): """Test the detailed client view works.""" hunt_id = self.CreateSampleHunt() client_id = self.SetupClient(0) rdf_flow = rdf_flow_objects.Flow( client_id=client_id, flow_id=flow.RandomFlowId(), parent_hunt_id=hunt_id, create_time=rdfvalue.RDFDatetime.Now()) rdf_flow.cpu_time_used.user_cpu_time = 5000 rdf_flow.network_bytes_sent = 1000000 data_store.REL_DB.WriteFlowObject(rdf_flow) # Open up and click on View Hunts then the first Hunt. self.Open("/") self.WaitUntil(self.IsElementPresent, "client_query") self.Click("css=a[grrtarget=hunts]") self.WaitUntil(self.IsTextPresent, hunt_id) self.Click("css=td:contains('%s')" % hunt_id) # Click the Overview Tab and check that the stats are present. self.Click("css=li[heading=Overview]") self.WaitUntil(self.IsTextPresent, "1h 23m 20s") self.WaitUntil(self.IsTextPresent, "976.6KiB")
def testStartupTriggersInterrogateWhenPreviousInterrogateIsDone(self): with test_lib.ConfigOverrider({"Source.version_numeric": 3000}): client_id = self.SetupClient(0) self._RunSendStartupInfo(client_id) data_store.REL_DB.WriteFlowObject( rdf_flow_objects.Flow( flow_id=flow.RandomFlowId(), client_id=client_id, flow_class_name=discovery.Interrogate.__name__, flow_state=rdf_flow_objects.Flow.FlowState.FINISHED)) flows = data_store.REL_DB.ReadAllFlowObjects(client_id, include_child_flows=False) orig_count = len([ f for f in flows if f.flow_class_name == discovery.Interrogate.__name__ ]) with mock.patch.multiple(discovery.Interrogate, Start=mock.DEFAULT, End=mock.DEFAULT): with test_lib.ConfigOverrider({"Source.version_numeric": 3001}): self._RunSendStartupInfo(client_id) flows = data_store.REL_DB.ReadAllFlowObjects(client_id, include_child_flows=False) interrogates = [ f for f in flows if f.flow_class_name == discovery.Interrogate.__name__ ] self.assertLen(interrogates, orig_count + 1)
def testStartupDoesNotTriggerInterrogateIfRecentInterrogateIsRunning(self): with test_lib.ConfigOverrider({"Source.version_numeric": 3000}): client_id = self.SetupClient(0) self._RunSendStartupInfo(client_id) data_store.REL_DB.WriteFlowObject( rdf_flow_objects.Flow( flow_id=flow.RandomFlowId(), client_id=client_id, flow_class_name=discovery.Interrogate.__name__, flow_state=rdf_flow_objects.Flow.FlowState.RUNNING)) flows = data_store.REL_DB.ReadAllFlowObjects( client_id, include_child_flows=False) orig_count = len([ f for f in flows if f.flow_class_name == discovery.Interrogate.__name__ ]) self._RunSendStartupInfo(client_id) flows = data_store.REL_DB.ReadAllFlowObjects( client_id, include_child_flows=False) same_ver_count = len([ f for f in flows if f.flow_class_name == discovery.Interrogate.__name__ ]) self.assertEqual(same_ver_count, orig_count)
def testFlowIdGeneration(self): self.assertLen(flow.RandomFlowId(), 16) with mock.patch.object(flow.random, "Id64", return_value=0xF0F1F2F3F4F5F6F7): self.assertEqual(flow.RandomFlowId(), "F0F1F2F3F4F5F6F7") with mock.patch.object(flow.random, "Id64", return_value=0): self.assertEqual(flow.RandomFlowId(), "0000000000000000") with mock.patch.object(flow.random, "Id64", return_value=1): self.assertEqual(flow.RandomFlowId(), "0000000000000001") with mock.patch.object(flow.random, "Id64", return_value=0x0000000100000000): self.assertEqual(flow.RandomFlowId(), "0000000100000000")
def testHuntOverviewGetsUpdatedWhenHuntChanges(self): hunt_id = self.CreateSampleHunt() client_id = self.SetupClient(0) rdf_flow = rdf_flow_objects.Flow( client_id=client_id, flow_id=flow.RandomFlowId(), parent_hunt_id=hunt_id, create_time=rdfvalue.RDFDatetime.Now()) rdf_flow.cpu_time_used.user_cpu_time = 5000 rdf_flow.network_bytes_sent = 1000000 data_store.REL_DB.WriteFlowObject(rdf_flow) self.Open("/") # Ensure auto-refresh updates happen every second. self.GetJavaScriptValue( "grrUi.hunt.huntOverviewDirective.setAutoRefreshInterval(1000);") self.Click("css=a[grrtarget=hunts]") self.Click("css=td:contains('%s')" % hunt_id) self.WaitUntil(self.IsTextPresent, "1h 23m 20s") self.WaitUntil(self.IsTextPresent, "976.6KiB") client_id = self.SetupClient(1) rdf_flow = rdf_flow_objects.Flow( client_id=client_id, flow_id=flow.RandomFlowId(), parent_hunt_id=hunt_id, create_time=rdfvalue.RDFDatetime.Now()) rdf_flow.cpu_time_used.user_cpu_time = 1000 rdf_flow.network_bytes_sent = 10000000 data_store.REL_DB.WriteFlowObject(rdf_flow) self.WaitUntil(self.IsTextPresent, "1h 40m") self.WaitUntil(self.IsTextPresent, "10.5MiB")
def _SetupHuntClientAndFlow(self, client_id=None, hunt_id=None, **additional_flow_args): client_id = self.InitializeClient(client_id=client_id) flow_id = flow.RandomFlowId() self.db.WriteClientMetadata(client_id, fleetspeak_enabled=False) rdf_flow = rdf_flow_objects.Flow( client_id=client_id, flow_id=flow_id, parent_hunt_id=hunt_id, create_time=rdfvalue.RDFDatetime.Now(), **additional_flow_args) self.db.WriteFlowObject(rdf_flow) return client_id, flow_id
def testDrainTaskSchedulerQueue(self): client_id = u"C.1234567890123456" flow_id = flow.RandomFlowId() data_store.REL_DB.WriteClientMetadata(client_id, fleetspeak_enabled=False) rdf_flow = rdf_flow_objects.Flow( client_id=client_id, flow_id=flow_id, create_time=rdfvalue.RDFDatetime.Now()) data_store.REL_DB.WriteFlowObject(rdf_flow) action_requests = [] for i in range(3): data_store.REL_DB.WriteFlowRequests([ rdf_flow_objects.FlowRequest(client_id=client_id, flow_id=flow_id, request_id=i) ]) action_requests.append( rdf_flows.ClientActionRequest(client_id=client_id, flow_id=flow_id, request_id=i, action_identifier="WmiQuery")) data_store.REL_DB.WriteClientActionRequests(action_requests) server = TestServer() res = server.DrainTaskSchedulerQueueForClient( rdfvalue.RDFURN(client_id)) msgs = [ rdf_flow_objects.GRRMessageFromClientActionRequest(r) for r in action_requests ] for r in res: r.task_id = 0 for m in msgs: m.task_id = 0 self.assertItemsEqual(res, msgs)
def _WriteNestedAndNonNestedLogEntries(self, hunt_obj): client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_obj.hunt_id) self.db.WriteFlowLogEntries([ rdf_flow_objects.FlowLogEntry( client_id=client_id, # Top-level hunt-induced flows should have hunt's ids. flow_id=flow_id, hunt_id=hunt_obj.hunt_id, message="blah") ]) for i in range(10): _, nested_flow_id = self._SetupHuntClientAndFlow( client_id=client_id, parent_flow_id=flow_id, hunt_id=hunt_obj.hunt_id, flow_id=flow.RandomFlowId()) self.db.WriteFlowLogEntries([ rdf_flow_objects.FlowLogEntry( client_id=client_id, flow_id=nested_flow_id, hunt_id=hunt_obj.hunt_id, message="blah_%d" % i) ])