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)