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 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 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 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 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 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 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 testIncorrectFlowType(self, db: abstract_db.Database) -> None: client_id = db_test_utils.InitializeClient(db) flow_id = "4815162342ABCDEF" flow_obj = rdf_flow_objects.Flow() flow_obj.client_id = client_id flow_obj.flow_id = flow_id flow_obj.flow_class_name = "NotArtifactCollector" db.WriteFlowObject(flow_obj) args = flow_plugin.ApiListFlowApplicableParsersArgs() args.client_id = client_id args.flow_id = flow_id with self.assertRaisesRegex(ValueError, "Not an artifact-collector flow"): self.handler.Handle(args)
def testNotAppliedParsers(self, db: abstract_db.Database) -> None: client_id = db_test_utils.InitializeClient(db) flow_id = "4815162342ABCDEF" flow_obj = rdf_flow_objects.Flow() flow_obj.client_id = client_id flow_obj.flow_id = flow_id flow_obj.flow_class_name = collectors.ArtifactCollectorFlow.__name__ flow_obj.args = rdf_artifacts.ArtifactCollectorFlowArgs( apply_parsers=False) db.WriteFlowObject(flow_obj) flow_result = rdf_flow_objects.FlowResult() flow_result.client_id = client_id flow_result.flow_id = flow_id flow_result.tag = "artifact:Fake" flow_result.payload = rdfvalue.RDFString("foobar") db.WriteFlowResults([flow_result]) args = flow_plugin.ApiListFlowApplicableParsersArgs() args.client_id = client_id args.flow_id = flow_id result = self.handler.Handle(args) self.assertCountEqual(result.parsers, [ flow_plugin.ApiParserDescriptor( type=flow_plugin.ApiParserDescriptor.Type.SINGLE_RESPONSE, name="FakeSingleResponse", ), flow_plugin.ApiParserDescriptor( type=flow_plugin.ApiParserDescriptor.Type.MULTI_RESPONSE, name="FakeMultiResponse", ), flow_plugin.ApiParserDescriptor( type=flow_plugin.ApiParserDescriptor.Type.SINGLE_FILE, name="FakeSingleFile", ), flow_plugin.ApiParserDescriptor( type=flow_plugin.ApiParserDescriptor.Type.MULTI_FILE, name="FakeMultiFile", ), ])
def testFlowWithResult(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) flow_result = rdf_flow_objects.FlowResult() flow_result.client_id = client_id flow_result.flow_id = flow_id flow_result.payload = rdf_timeline.TimelineResult( filesystem_type="ntfs") db.WriteFlowResults([flow_result]) self.assertEqual(timeline_flow.FilesystemType(client_id, flow_id), "ntfs")
def testAlreadyAppliedParsers(self, db: abstract_db.Database) -> None: client_id = db_test_utils.InitializeClient(db) flow_id = "4815162342ABCDEF" flow_obj = rdf_flow_objects.Flow() flow_obj.client_id = client_id flow_obj.flow_id = flow_id flow_obj.flow_class_name = collectors.ArtifactCollectorFlow.__name__ flow_obj.args = rdf_artifacts.ArtifactCollectorFlowArgs( apply_parsers=True) db.WriteFlowObject(flow_obj) flow_result = rdf_flow_objects.FlowResult() flow_result.client_id = client_id flow_result.flow_id = flow_id flow_result.tag = "artifact:Fake" db.WriteFlowResults([flow_result]) args = flow_plugin.ApiListFlowApplicableParsersArgs() args.client_id = client_id args.flow_id = flow_id result = self.handler.Handle(args) self.assertEmpty(result.parsers)