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 Run(self, args: rdf_timeline.TimelineArgs) -> None:
"""Executes the client action."""
result = rdf_timeline.TimelineResult()
entries = Walk(args.root)
for entry_batch in rdf_timeline.TimelineEntry.SerializeStream(entries):
entry_batch_blob = rdf_protodict.DataBlob(data=entry_batch)
self.SendReply(entry_batch_blob, session_id=self._TRANSFER_STORE_ID)
entry_batch_blob_id = hashlib.sha256(entry_batch).digest()
result.entry_batch_blob_ids.append(entry_batch_blob_id)
self.Progress()
self.SendReply(result)
def Run(self, args: rdf_timeline.TimelineArgs) -> None:
"""Executes the client action."""
entries = iterator.Counted(Walk(args.root))
for entry_batch in rdf_timeline.TimelineEntry.SerializeStream(entries):
entry_batch_blob = rdf_protodict.DataBlob(data=entry_batch)
self.SendReply(entry_batch_blob, session_id=self._TRANSFER_STORE_ID)
entry_batch_blob_id = hashlib.sha256(entry_batch).digest()
result = rdf_timeline.TimelineResult()
result.entry_batch_blob_ids.append(entry_batch_blob_id)
result.entry_count = entries.count
self.SendReply(result)
# Each result should contain information only about the number of entries
# in the current batch, so after the results are sent we simply reset the
# counter.
entries.Reset()
def testNtfsFileReferenceFormatInference(self):
entry = rdf_timeline.TimelineEntry()
entry.path = "/foo/bar/baz".encode("utf-8")
entry.ino = 1688849860339456
client_id = db_test_utils.InitializeClient(data_store.REL_DB)
flow_id = "F00BA542"
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)
blobs = list(rdf_timeline.TimelineEntry.SerializeStream(iter([entry])))
blob_ids = data_store.BLOBS.WriteBlobsWithUnknownHashes(blobs)
result = rdf_timeline.TimelineResult()
result.entry_batch_blob_ids = [
blob_id.AsBytes() for blob_id in blob_ids
]
result.filesystem_type = "NTFS"
flow_result = rdf_flow_objects.FlowResult()
flow_result.client_id = client_id
flow_result.flow_id = flow_id
flow_result.payload = result
data_store.REL_DB.WriteFlowResults([flow_result])
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")
rows = list(csv.reader(io.StringIO(content), delimiter="|"))
self.assertLen(rows, 1)
self.assertEqual(rows[0][1], "/foo/bar/baz")
self.assertEqual(rows[0][2], "75520-6")
def _WriteTimeline(
client_id: Text,
entries: Sequence[rdf_timeline.TimelineEntry],
hunt_id: Optional[Text] = None,
) -> Text:
"""Writes a timeline to the database (as fake flow result).
Args:
client_id: An identifier of the client for which the flow ran.
entries: A sequence of timeline entries produced by the flow run.
hunt_id: An (optional) identifier of a hunt the flows belong to.
Returns:
An identifier of the flow.
"""
flow_id = "".join(random.choice("ABCDEF") for _ in range(8))
flow_obj = rdf_flow_objects.Flow()
flow_obj.flow_id = flow_id
flow_obj.client_id = client_id
flow_obj.flow_class_name = timeline.TimelineFlow.__name__
flow_obj.create_time = rdfvalue.RDFDatetime.Now()
flow_obj.parent_hunt_id = hunt_id
data_store.REL_DB.WriteFlowObject(flow_obj)
blobs = list(rdf_timeline.TimelineEntry.SerializeStream(iter(entries)))
blob_ids = data_store.BLOBS.WriteBlobsWithUnknownHashes(blobs)
result = rdf_timeline.TimelineResult()
result.entry_batch_blob_ids = [blob_id.AsBytes() for blob_id in blob_ids]
flow_result = rdf_flow_objects.FlowResult()
flow_result.client_id = client_id
flow_result.flow_id = flow_id
flow_result.payload = result
data_store.REL_DB.WriteFlowResults([flow_result])
return flow_id
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 testGetCollectedTimelinesGzchunked(self):
client_id = db_test_utils.InitializeClient(data_store.REL_DB)
fqdn = "foo.bar.baz"
snapshot = rdf_objects.ClientSnapshot()
snapshot.client_id = client_id
snapshot.knowledge_base.fqdn = fqdn
data_store.REL_DB.WriteClientSnapshot(snapshot)
hunt_id = "A0B1D2C3"
flow_id = "0A1B2D3C"
hunt_obj = rdf_hunt_objects.Hunt()
hunt_obj.hunt_id = hunt_id
hunt_obj.args.standard.client_ids = [client_id]
hunt_obj.args.standard.flow_name = timeline.TimelineFlow.__name__
hunt_obj.hunt_state = rdf_hunt_objects.Hunt.HuntState.PAUSED
data_store.REL_DB.WriteHuntObject(hunt_obj)
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()
flow_obj.parent_hunt_id = hunt_id
data_store.REL_DB.WriteFlowObject(flow_obj)
entry_1 = rdf_timeline.TimelineEntry()
entry_1.path = "/foo/bar".encode("utf-8")
entry_1.ino = 7890178901
entry_1.size = 4815162342
entry_1.atime_ns = 123 * 10**9
entry_1.mtime_ns = 234 * 10**9
entry_1.ctime_ns = 567 * 10**9
entry_1.mode = 0o654
entry_2 = rdf_timeline.TimelineEntry()
entry_2.path = "/foo/baz".encode("utf-8")
entry_1.ino = 8765487654
entry_2.size = 1337
entry_1.atime_ns = 987 * 10**9
entry_1.mtime_ns = 876 * 10**9
entry_1.ctime_ns = 765 * 10**9
entry_2.mode = 0o757
entries = [entry_1, entry_2]
blobs = list(rdf_timeline.TimelineEntry.SerializeStream(iter(entries)))
blob_ids = data_store.BLOBS.WriteBlobsWithUnknownHashes(blobs)
result = rdf_timeline.TimelineResult()
result.entry_batch_blob_ids = [
blob_id.AsBytes() for blob_id in blob_ids
]
flow_result = rdf_flow_objects.FlowResult()
flow_result.client_id = client_id
flow_result.flow_id = flow_id
flow_result.payload = result
data_store.REL_DB.WriteFlowResults([flow_result])
buffer = io.BytesIO()
fmt = timeline_pb2.ApiGetCollectedTimelineArgs.Format.RAW_GZCHUNKED
self.api.Hunt(hunt_id).GetCollectedTimelines(fmt).WriteToStream(buffer)
with zipfile.ZipFile(buffer, mode="r") as archive:
with archive.open(f"{client_id}_{fqdn}.gzchunked",
mode="r") as file:
chunks = chunked.ReadAll(file)
entries = list(
rdf_timeline.TimelineEntry.DeserializeStream(chunks))
self.assertEqual(entries, [entry_1, entry_2])
def testGetCollectedTimelinesBody(self):
client_id = db_test_utils.InitializeClient(data_store.REL_DB)
fqdn = "foo.bar.quux"
snapshot = rdf_objects.ClientSnapshot()
snapshot.client_id = client_id
snapshot.knowledge_base.fqdn = fqdn
data_store.REL_DB.WriteClientSnapshot(snapshot)
hunt_id = "B1C2E3D4"
flow_id = "1B2C3E4D"
hunt_obj = rdf_hunt_objects.Hunt()
hunt_obj.hunt_id = hunt_id
hunt_obj.args.standard.client_ids = [client_id]
hunt_obj.args.standard.flow_name = timeline.TimelineFlow.__name__
hunt_obj.hunt_state = rdf_hunt_objects.Hunt.HuntState.PAUSED
data_store.REL_DB.WriteHuntObject(hunt_obj)
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()
flow_obj.parent_hunt_id = hunt_id
data_store.REL_DB.WriteFlowObject(flow_obj)
entry_1 = rdf_timeline.TimelineEntry()
entry_1.path = "/bar/baz/quux".encode("utf-8")
entry_1.ino = 5926273453
entry_1.size = 13373
entry_1.atime_ns = 111 * 10**9
entry_1.mtime_ns = 222 * 10**9
entry_1.ctime_ns = 333 * 10**9
entry_1.mode = 0o664
entry_2 = rdf_timeline.TimelineEntry()
entry_2.path = "/bar/baz/quuz".encode("utf-8")
entry_2.ino = 6037384564
entry_2.size = 13374
entry_2.atime_ns = 777 * 10**9
entry_2.mtime_ns = 888 * 10**9
entry_2.ctime_ns = 999 * 10**9
entry_2.mode = 0o777
entries = [entry_1, entry_2]
blobs = list(rdf_timeline.TimelineEntry.SerializeStream(iter(entries)))
blob_ids = data_store.BLOBS.WriteBlobsWithUnknownHashes(blobs)
result = rdf_timeline.TimelineResult()
result.entry_batch_blob_ids = [
blob_id.AsBytes() for blob_id in blob_ids
]
flow_result = rdf_flow_objects.FlowResult()
flow_result.client_id = client_id
flow_result.flow_id = flow_id
flow_result.payload = result
data_store.REL_DB.WriteFlowResults([flow_result])
buffer = io.BytesIO()
self.api.Hunt(hunt_id).GetCollectedTimelines(
timeline_pb2.ApiGetCollectedTimelineArgs.Format.BODY
).WriteToStream(buffer)
with zipfile.ZipFile(buffer, mode="r") as archive:
with archive.open(f"{client_id}_{fqdn}.body", mode="r") as file:
content_file = file.read().decode("utf-8")
rows = list(
csv.reader(io.StringIO(content_file), delimiter="|"))
self.assertLen(rows, 2)
self.assertEqual(rows[0][1], "/bar/baz/quux")
self.assertEqual(rows[0][2], "5926273453")
self.assertEqual(rows[0][3], stat.filemode(0o664))
self.assertEqual(rows[0][6], "13373")
self.assertEqual(rows[0][7], "111")
self.assertEqual(rows[0][8], "222")
self.assertEqual(rows[0][9], "333")
self.assertEqual(rows[1][1], "/bar/baz/quuz")
self.assertEqual(rows[1][2], "6037384564")
self.assertEqual(rows[1][3], stat.filemode(0o777))
self.assertEqual(rows[1][6], "13374")
self.assertEqual(rows[1][7], "777")
self.assertEqual(rows[1][8], "888")
self.assertEqual(rows[1][9], "999")
