def testUserFlowsReportPlugin(self): client_id = self.SetupClient(1) with test_lib.FakeTime( rdfvalue.RDFDatetime.FromHumanReadable("2012/12/14")): AddFakeAuditLog(user="******") data_store.REL_DB.WriteFlowObject( rdf_flow_objects.Flow(flow_class_name="GetClientStats", creator="User123", client_id=client_id, flow_id="E0000000", create_time=rdfvalue.RDFDatetime.Now())) with test_lib.FakeTime( rdfvalue.RDFDatetime.FromHumanReadable("2012/12/22")): for i in range(10): data_store.REL_DB.WriteFlowObject( rdf_flow_objects.Flow( flow_class_name="GetClientStats", creator="User123", client_id=client_id, flow_id="{:08X}".format(i), create_time=rdfvalue.RDFDatetime.Now())) AddFakeAuditLog(user="******") data_store.REL_DB.WriteFlowObject( rdf_flow_objects.Flow(flow_class_name="ArtifactCollectorFlow", creator="User456", client_id=client_id, flow_id="F0000000", create_time=rdfvalue.RDFDatetime.Now())) AddFakeAuditLog(user="******") report = report_plugins.GetReportByName( server_report_plugins.UserFlowsReportPlugin.__name__) start = rdfvalue.RDFDatetime.FromHumanReadable("2012/12/15") month_duration = rdfvalue.Duration.From(30, rdfvalue.DAYS) api_report_data = report.GetReportData(stats_api.ApiGetReportArgs( name=report.__class__.__name__, start_time=start, duration=month_duration), token=self.token) self.assertEqual( api_report_data, rdf_report_plugins.ApiReportData( representation_type=RepresentationType.STACK_CHART, stack_chart=rdf_report_plugins.ApiStackChartReportData( x_ticks=[], data=[ rdf_report_plugins.ApiReportDataSeries2D( label=u"GetClientStats\u2003Run By: User123 (10)", points=[ApiReportDataPoint2D(x=0, y=10)]), rdf_report_plugins.ApiReportDataSeries2D( label= u"ArtifactCollectorFlow\u2003Run By: User456 (1)", points=[ApiReportDataPoint2D(x=1, y=1)]) ])))
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 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 testResultMetadataAreCorrectlyUpdatedAfterMultiplePersistStateCalls( self, db: abstract_db.Database): client_id = db_test_utils.InitializeClient(db) flow = rdf_flow_objects.Flow() flow.client_id = client_id flow.flow_id = self._FLOW_ID db.WriteFlowObject(flow) flow_obj = FlowBaseTest.Flow(flow) flow_obj.SendReply(rdf_client.ClientInformation()) flow_obj.PersistState() flow_obj.PersistState() db.WriteFlowObject(flow_obj.rdf_flow) flow_2 = db.ReadFlowObject(client_id, self._FLOW_ID) flow_obj_2 = FlowBaseTest.Flow(flow_2) result_metadata = flow_obj_2.GetResultMetadata() self.assertLen(result_metadata.num_results_per_type_tag, 1) self.assertTrue(result_metadata.is_metadata_set) self.assertEqual(result_metadata.num_results_per_type_tag[0].type, "ClientInformation") self.assertEqual(result_metadata.num_results_per_type_tag[0].tag, "") self.assertEqual(result_metadata.num_results_per_type_tag[0].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 testGetKBDependencies(self): """Test that KB dependencies are calculated correctly.""" artifact_registry.REGISTRY.ClearSources() try: test_artifacts_file = os.path.join(config.CONFIG["Test.data_dir"], "artifacts", "test_artifacts.json") artifact_registry.REGISTRY.AddFileSource(test_artifacts_file) with test_lib.ConfigOverrider({ "Artifacts.knowledge_base": [ "DepsParent", "DepsDesktop", "DepsHomedir", "DepsWindir", "DepsWindirRegex", "DepsControlSet", "FakeArtifact" ], "Artifacts.knowledge_base_additions": ["DepsHomedir2"], "Artifacts.knowledge_base_skip": ["DepsWindir"], "Artifacts.knowledge_base_heavyweight": ["FakeArtifact"] }): args = artifact.KnowledgeBaseInitializationArgs(lightweight=True) kb_init = artifact.KnowledgeBaseInitializationFlow( rdf_flow_objects.Flow(args=args)) kb_init.state["all_deps"] = set() kb_init.state["awaiting_deps_artifacts"] = [] kb_init.state["knowledge_base"] = rdf_client.KnowledgeBase(os="Windows") no_deps = kb_init.GetFirstFlowsForCollection() self.assertCountEqual(no_deps, ["DepsControlSet", "DepsHomedir2"]) self.assertCountEqual(kb_init.state.all_deps, [ "users.homedir", "users.desktop", "users.username", "environ_windir", "current_control_set" ]) self.assertCountEqual( kb_init.state.awaiting_deps_artifacts, ["DepsParent", "DepsDesktop", "DepsHomedir", "DepsWindirRegex"]) finally: artifact.LoadArtifactsOnce()
def _LoadFlows(self, client_id, min_create_time, token): """Yields all flows for the given client_id and time range. Args: client_id: client URN min_create_time: minimum creation time (inclusive) token: acl token Yields: flow_objects.Flow objects """ if data_store.RelationalDBFlowsEnabled(): flow_list = data_store.REL_DB.ReadAllFlowObjects( client_id, min_create_time=min_create_time) for flow_obj in flow_list: if not flow_obj.parent_flow_id: yield flow_obj else: now = rdfvalue.RDFDatetime.Now() client_id_urn = rdf_client.ClientURN(client_id) flows_dir = aff4.FACTORY.Open(client_id_urn.Add("flows"), token=token) # Save DB roundtrips by checking both conditions at once. flow_list = flows_dir.ListChildren( age=(min_create_time.AsMicrosecondsSinceEpoch(), now.AsMicrosecondsSinceEpoch())) for flow_obj in aff4.FACTORY.MultiOpen(flow_list, token=token): yield rdf_flow_objects.Flow( args=flow_obj.args, flow_class_name=flow_obj.runner_args.flow_name, flow_id=flow_obj.urn.Basename(), create_time=flow_obj.context.create_time, creator=flow_obj.creator, )
def testResultMetadataHasGroupedNumberOfReplies(self, db: abstract_db.Database): client_id = db_test_utils.InitializeClient(db) flow = rdf_flow_objects.Flow() flow.client_id = client_id flow.flow_id = self._FLOW_ID db.WriteFlowObject(flow) flow_obj = FlowBaseTest.Flow(flow) flow_obj.SendReply(rdf_client.ClientInformation()) flow_obj.SendReply(rdf_client.StartupInfo()) flow_obj.SendReply(rdf_client.StartupInfo()) flow_obj.SendReply(rdf_client.StartupInfo(), tag="foo") flow_obj.PersistState() db.WriteFlowObject(flow_obj.rdf_flow) flow_2 = db.ReadFlowObject(client_id, self._FLOW_ID) flow_obj_2 = FlowBaseTest.Flow(flow_2) result_metadata = flow_obj_2.GetResultMetadata() self.assertLen(result_metadata.num_results_per_type_tag, 3) sorted_counts = sorted(result_metadata.num_results_per_type_tag, key=lambda v: (v.type, v.tag)) self.assertEqual(sorted_counts[0].type, "ClientInformation") self.assertEqual(sorted_counts[0].tag, "") self.assertEqual(sorted_counts[0].count, 1) self.assertEqual(sorted_counts[1].type, "StartupInfo") self.assertEqual(sorted_counts[1].tag, "") self.assertEqual(sorted_counts[1].count, 2) self.assertEqual(sorted_counts[2].type, "StartupInfo") self.assertEqual(sorted_counts[2].tag, "foo") self.assertEqual(sorted_counts[2].count, 1)
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 testFlowWritingUnknownClient(self): flow_id = u"1234ABCD" client_id = u"C.1234567890123456" rdf_flow = rdf_flow_objects.Flow(client_id=client_id, flow_id=flow_id) with self.assertRaises(db.UnknownClientError): self.db.WriteFlowObject(rdf_flow)
def testClientInfoDefault(self, db: abstract_db.Database): client_id = db_test_utils.InitializeClient(db) flow = rdf_flow_objects.Flow() flow.client_id = client_id flow.flow_id = self._FLOW_ID flow = FlowBaseTest.Flow(flow) self.assertIsInstance(flow.client_info, rdf_client.ClientInformation) self.assertEmpty(flow.client_info.client_name)
def testReceiveMessageList(self): fs_server = fleetspeak_frontend_server.GRRFSServer() client_id = "C.1234567890123456" flow_id = "12345678" data_store.REL_DB.WriteClientMetadata(client_id, fleetspeak_enabled=True) 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) flow_request = rdf_flow_objects.FlowRequest(client_id=client_id, flow_id=flow_id, request_id=1) data_store.REL_DB.WriteFlowRequests([flow_request]) session_id = "%s/%s" % (client_id, flow_id) fs_client_id = fleetspeak_utils.GRRIDToFleetspeakID(client_id) grr_messages = [] for i in range(1, 10): grr_message = rdf_flows.GrrMessage(request_id=1, response_id=i + 1, session_id=session_id, payload=rdfvalue.RDFInteger(i)) grr_messages.append(grr_message) packed_messages = rdf_flows.PackedMessageList() communicator.Communicator.EncodeMessageList( rdf_flows.MessageList(job=grr_messages), packed_messages) fs_message = fs_common_pb2.Message(message_type="MessageList", source=fs_common_pb2.Address( client_id=fs_client_id, service_name=FS_SERVICE_NAME)) fs_message.data.Pack(packed_messages.AsPrimitiveProto()) fs_message.validation_info.tags["foo"] = "bar" with test_lib.FakeTime( rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123)): fs_server.Process(fs_message, None) # Ensure the last-ping timestamp gets updated. client_data = data_store.REL_DB.MultiReadClientMetadata([client_id]) self.assertEqual(client_data[client_id].ping, rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123)) self.assertEqual( client_data[client_id].last_fleetspeak_validation_info. ToStringDict(), {"foo": "bar"}) flow_data = data_store.REL_DB.ReadAllFlowRequestsAndResponses( client_id, flow_id) self.assertLen(flow_data, 1) stored_flow_request, flow_responses = flow_data[0] self.assertEqual(stored_flow_request, flow_request) self.assertLen(flow_responses, 9)
def testClientInfoDefault(self, db: abstract_db.Database): client_id = "C.0123456789ABCDEF" db.WriteClientMetadata(client_id, fleetspeak_enabled=False) flow = rdf_flow_objects.Flow() flow.client_id = client_id flow.flow_id = "FEDCBA9876543210" flow = FlowBaseTest.Flow(flow) self.assertIsInstance(flow.client_info, rdf_client.ClientInformation) self.assertEmpty(flow.client_info.client_name)
def testBodyMultipleResults(self): client_id = db_test_utils.InitializeClient(data_store.REL_DB) flow_id = "ABCDEF42" flow_obj = rdf_flow_objects.Flow() flow_obj.client_id = client_id flow_obj.flow_id = flow_id flow_obj.flow_class_name = timeline.TimelineFlow.__name__ flow_obj.create_time = rdfvalue.RDFDatetime.Now() data_store.REL_DB.WriteFlowObject(flow_obj) entry_1 = rdf_timeline.TimelineEntry() entry_1.path = "/foo".encode("utf-8") blobs_1 = list( rdf_timeline.TimelineEntry.SerializeStream(iter([entry_1]))) (blob_id_1, ) = data_store.BLOBS.WriteBlobsWithUnknownHashes(blobs_1) result_1 = rdf_timeline.TimelineResult() result_1.entry_batch_blob_ids = [blob_id_1.AsBytes()] entry_2 = rdf_timeline.TimelineEntry() entry_2.path = "/bar".encode("utf-8") blobs_2 = list( rdf_timeline.TimelineEntry.SerializeStream(iter([entry_2]))) (blob_id_2, ) = data_store.BLOBS.WriteBlobsWithUnknownHashes(blobs_2) result_2 = rdf_timeline.TimelineResult() result_2.entry_batch_blob_ids = [blob_id_2.AsBytes()] flow_result_1 = rdf_flow_objects.FlowResult() flow_result_1.client_id = client_id flow_result_1.flow_id = flow_id flow_result_1.payload = result_1 flow_result_2 = rdf_flow_objects.FlowResult() flow_result_2.client_id = client_id flow_result_2.flow_id = flow_id flow_result_2.payload = result_2 data_store.REL_DB.WriteFlowResults([flow_result_1, flow_result_2]) args = api_timeline.ApiGetCollectedTimelineArgs() args.client_id = client_id args.flow_id = flow_id args.format = api_timeline.ApiGetCollectedTimelineArgs.Format.BODY result = self.handler.Handle(args) content = b"".join(result.GenerateContent()).decode("utf-8") self.assertIn("|/foo|", content) self.assertIn("|/bar|", content)
def testGrepRegexCombination(self): args = rdf_artifacts.ArtifactCollectorFlowArgs() collect_flow = collectors.ArtifactCollectorFlow( rdf_flow_objects.Flow(args=args)) self.assertEqual(collect_flow._CombineRegex([b"simple"]), b"simple") self.assertEqual(collect_flow._CombineRegex([b"a", b"b"]), b"(a)|(b)") self.assertEqual( collect_flow._CombineRegex([b"a", b"b", b"c"]), b"(a)|(b)|(c)") self.assertEqual( collect_flow._CombineRegex([b"a|b", b"[^_]b", b"c|d"]), b"(a|b)|([^_]b)|(c|d)")
def testReturnsDefaultFlowProgressForEmptyFlow(self, db: abstract_db.Database): client_id = db_test_utils.InitializeClient(db) flow = rdf_flow_objects.Flow() flow.client_id = client_id flow.flow_id = self._FLOW_ID db.WriteFlowObject(flow) flow_obj = FlowBaseTest.Flow(flow) progress = flow_obj.GetProgress() self.assertIsInstance(progress, rdf_flow_objects.DefaultFlowProgress)
def _SetupClientAndFlow(self, **additional_flow_args): flow_id = u"1234ABCD" client_id = u"C.1234567890123456" self.db.WriteClientMetadata(client_id, fleetspeak_enabled=False) rdf_flow = rdf_flow_objects.Flow(client_id=client_id, flow_id=flow_id, **additional_flow_args) self.db.WriteFlowObject(rdf_flow) return client_id, flow_id
def setUp(self): super(SplunkOutputPluginTest, self).setUp() self.client_id = self.SetupClient(0) self.flow_id = '12345678' data_store.REL_DB.WriteFlowObject( rdf_flow_objects.Flow( flow_id=self.flow_id, client_id=self.client_id, flow_class_name='ClientFileFinder', create_time=rdfvalue.RDFDatetime.Now(), ))
def testReadChildFlows(self): client_id = u"C.1234567890123456" self.db.WriteClientMetadata(client_id, fleetspeak_enabled=False) self.db.WriteFlowObject( rdf_flow_objects.Flow(flow_id=u"00000001", client_id=client_id)) self.db.WriteFlowObject( rdf_flow_objects.Flow(flow_id=u"00000002", client_id=client_id, parent_flow_id=u"00000001")) self.db.WriteFlowObject( rdf_flow_objects.Flow(flow_id=u"00000003", client_id=client_id, parent_flow_id=u"00000002")) self.db.WriteFlowObject( rdf_flow_objects.Flow(flow_id=u"00000004", client_id=client_id, parent_flow_id=u"00000001")) # This one is completely unrelated (different client id). self.db.WriteClientMetadata(u"C.1234567890123457", fleetspeak_enabled=False) self.db.WriteFlowObject( rdf_flow_objects.Flow(flow_id=u"00000002", client_id=u"C.1234567890123457", parent_flow_id=u"00000001")) children = self.db.ReadChildFlowObjects(client_id, u"00000001") self.assertEqual(len(children), 2) for c in children: self.assertEqual(c.parent_flow_id, u"00000001") children = self.db.ReadChildFlowObjects(client_id, u"00000002") self.assertEqual(len(children), 1) self.assertEqual(children[0].parent_flow_id, u"00000002") self.assertEqual(children[0].flow_id, u"00000003") children = self.db.ReadChildFlowObjects(client_id, u"00000003") self.assertEqual(len(children), 0)
def testReceiveMessages_Relational(self): if not data_store.RelationalDBEnabled(): self.skipTest("Rel-db-only test.") fs_server = fs_frontend_tool.GRRFSServer() client_id = "C.1234567890123456" flow_id = "12345678" data_store.REL_DB.WriteClientMetadata(client_id, fleetspeak_enabled=True) 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) flow_request = rdf_flow_objects.FlowRequest(client_id=client_id, flow_id=flow_id, request_id=1) data_store.REL_DB.WriteFlowRequests([flow_request]) session_id = "%s/%s" % (client_id, flow_id) fs_client_id = fleetspeak_utils.GRRIDToFleetspeakID(client_id) fs_messages = [] for i in range(1, 10): grr_message = rdf_flows.GrrMessage(request_id=1, response_id=i + 1, session_id=session_id, payload=rdfvalue.RDFInteger(i)) fs_message = fs_common_pb2.Message( message_type="GrrMessage", source=fs_common_pb2.Address(client_id=fs_client_id, service_name=FS_SERVICE_NAME)) fs_message.data.Pack(grr_message.AsPrimitiveProto()) fs_messages.append(fs_message) with test_lib.FakeTime( rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123)): for fs_message in fs_messages: fs_server.Process(fs_message, None) # Ensure the last-ping timestamp gets updated. client_data = data_store.REL_DB.MultiReadClientMetadata([client_id]) self.assertEqual(client_data[client_id].ping, rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123)) flow_data = data_store.REL_DB.ReadAllFlowRequestsAndResponses( client_id, flow_id) self.assertLen(flow_data, 1) stored_flow_request, flow_responses = flow_data[0] self.assertEqual(stored_flow_request, flow_request) self.assertLen(flow_responses, 9)
def testFlowWithNoResult(self, db: abstract_db.Database) -> None: client_id = "C.1234567890123456" flow_id = "ABCDEF92" db.WriteClientMetadata(client_id, last_ping=rdfvalue.RDFDatetime.Now()) flow_obj = rdf_flow_objects.Flow() flow_obj.client_id = client_id flow_obj.flow_id = flow_id flow_obj.flow_class_name = timeline_flow.TimelineFlow.__name__ flow_obj.create_time = rdfvalue.RDFDatetime.Now() db.WriteFlowObject(flow_obj) self.assertIsNone(timeline_flow.FilesystemType(client_id, flow_id))
def testLogWithoutFormatArgs(self, db: abstract_db.Database) -> None: client_id = db_test_utils.InitializeClient(db) flow = rdf_flow_objects.Flow() flow.client_id = client_id flow.flow_id = self._FLOW_ID db.WriteFlowObject(flow) flow = FlowBaseTest.Flow(flow) flow.Log("foo %s %s") logs = db.ReadFlowLogEntries(client_id, self._FLOW_ID, offset=0, count=1024) self.assertLen(logs, 1) self.assertEqual(logs[0].message, "foo %s %s")
def testReturnsEmptyResultMetadataForEmptyFlow(self, db: abstract_db.Database): client_id = db_test_utils.InitializeClient(db) flow = rdf_flow_objects.Flow() flow.client_id = client_id flow.flow_id = self._FLOW_ID db.WriteFlowObject(flow) flow_obj = FlowBaseTest.Flow(flow) result_metadata = flow_obj.GetResultMetadata() self.assertIsInstance(result_metadata, rdf_flow_objects.FlowResultMetadata) self.assertEmpty(result_metadata.num_results_per_type_tag)
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 _FlowSetup(self, client_id, flow_id): rdf_flow = rdf_flow_objects.Flow( flow_class_name=compatibility.GetName( administrative.OnlineNotification), client_id=client_id, flow_id=flow_id, create_time=rdfvalue.RDFDatetime.Now()) data_store.REL_DB.WriteFlowObject(rdf_flow) req = rdf_flow_objects.FlowRequest( client_id=client_id, flow_id=flow_id, request_id=1) data_store.REL_DB.WriteFlowRequests([req]) return rdf_flow, req
def testLogWithFormatArgs(self, db: abstract_db.Database) -> None: client_id = db_test_utils.InitializeClient(db) flow_id = "FEDCBA9876543210" flow = rdf_flow_objects.Flow() flow.client_id = client_id flow.flow_id = flow_id db.WriteFlowObject(flow) flow = FlowBaseTest.Flow(flow) flow.Log("foo %s %s", "bar", 42) logs = db.ReadFlowLogEntries(client_id, flow_id, offset=0, count=1024) self.assertLen(logs, 1) self.assertEqual(logs[0].message, "foo bar 42")
def testPathSpecCasingIsCorrected(self): flow = memory.DumpProcessMemory(rdf_flow_objects.Flow()) flow.SendReply = mock.Mock(spec=flow.SendReply) request = rdf_flow_objects.FlowRequest( request_data={ "YaraProcessDumpResponse": rdf_memory.YaraProcessDumpResponse(dumped_processes=[ rdf_memory.YaraProcessDumpInformation(memory_regions=[ rdf_memory.ProcessMemoryRegion( start=1, size=1, file=rdf_paths.PathSpec.Temp( path="/C:/grr/x_1_0_1.tmp")), rdf_memory.ProcessMemoryRegion( start=1, size=1, file=rdf_paths.PathSpec.Temp( path="/C:/GRR/x_1_1_2.tmp")) ]) ]) }) pathspecs = [ rdf_paths.PathSpec.Temp(path="/C:/Grr/x_1_0_1.tmp"), rdf_paths.PathSpec.Temp(path="/C:/Grr/x_1_1_2.tmp") ] responses = flow_responses.Responses.FromResponses( request, [ rdf_flow_objects.FlowResponse(payload=rdf_client_fs.StatEntry( pathspec=pathspec)) for pathspec in pathspecs ]) flow.ProcessMemoryRegions(responses) flow.SendReply.assert_any_call( rdf_memory.YaraProcessDumpResponse(dumped_processes=[ rdf_memory.YaraProcessDumpInformation(memory_regions=[ rdf_memory.ProcessMemoryRegion( start=1, size=1, file=rdf_paths.PathSpec.Temp( path="/C:/Grr/x_1_0_1.tmp")), rdf_memory.ProcessMemoryRegion( start=1, size=1, file=rdf_paths.PathSpec.Temp( path="/C:/Grr/x_1_1_2.tmp")) ]) ]))
def testInterpolateArgs(self): args = rdf_artifacts.ArtifactCollectorFlowArgs() collect_flow = collectors.ArtifactCollectorFlow( rdf_flow_objects.Flow(args=args)) kb = rdf_client.KnowledgeBase() kb.MergeOrAddUser(rdf_client.User(username="******")) kb.MergeOrAddUser(rdf_client.User(username="******")) collect_flow.state["knowledge_base"] = kb collect_flow.current_artifact_name = "blah" test_rdf = rdf_client.KnowledgeBase() action_args = { "usernames": ["%%users.username%%", "%%users.username%%"], "nointerp": "asdfsdf", "notastring": test_rdf } kwargs = collect_flow.InterpolateDict(action_args) self.assertCountEqual(kwargs["usernames"], ["test1", "test2", "test1", "test2"]) self.assertEqual(kwargs["nointerp"], "asdfsdf") self.assertEqual(kwargs["notastring"], test_rdf) # We should be using an array since users.username will expand to multiple # values. self.assertRaises(ValueError, collect_flow.InterpolateDict, {"bad": "%%users.username%%"}) list_args = collect_flow.InterpolateList( ["%%users.username%%", r"%%users.username%%\aa"]) self.assertCountEqual(list_args, ["test1", "test2", r"test1\aa", r"test2\aa"]) list_args = collect_flow.InterpolateList(["one"]) self.assertEqual(list_args, ["one"]) # Ignore the failure in users.desktop, report the others. collect_flow.args.ignore_interpolation_errors = True list_args = collect_flow.InterpolateList( ["%%users.desktop%%", r"%%users.username%%\aa"]) self.assertCountEqual(list_args, [r"test1\aa", r"test2\aa"]) # Both fail. list_args = collect_flow.InterpolateList( [r"%%users.desktop%%\aa", r"%%users.sid%%\aa"]) self.assertCountEqual(list_args, [])
def testClientInfo(self, db: abstract_db.Database): client_id = db_test_utils.InitializeClient(db) startup_info = rdf_client.StartupInfo() startup_info.client_info.client_name = "rrg" startup_info.client_info.client_version = 1337 db.WriteClientStartupInfo(client_id, startup_info) flow = rdf_flow_objects.Flow() flow.client_id = client_id flow.flow_id = self._FLOW_ID flow = FlowBaseTest.Flow(flow) self.assertIsInstance(flow.client_info, rdf_client.ClientInformation) self.assertEqual(flow.client_info.client_name, "rrg") self.assertEqual(flow.client_info.client_version, 1337)