def testClientStartupInfo(self):
"""StartupInfo is written to a separate table, make sure the merge works."""
d = self.db
client_id = self.InitializeClient()
client = rdf_objects.ClientSnapshot(client_id=client_id, kernel="12.3")
client.startup_info = rdf_client.StartupInfo(boot_time=123)
client.knowledge_base.fqdn = "test1234.examples.com"
d.WriteClientSnapshot(client)
client = d.ReadClientSnapshot(client_id)
self.assertEqual(client.startup_info.boot_time, 123)
client.kernel = "12.4"
client.startup_info = rdf_client.StartupInfo(boot_time=124)
d.WriteClientSnapshot(client)
client.kernel = "12.5"
client.startup_info = rdf_client.StartupInfo(boot_time=125)
d.WriteClientSnapshot(client)
hist = d.ReadClientSnapshotHistory(client_id)
self.assertEqual(len(hist), 3)
startup_infos = [cl.startup_info for cl in hist]
self.assertEqual([si.boot_time for si in startup_infos],
[125, 124, 123])
# StartupInfos written using WriteClient show up in the StartupInfoHistory.
history = d.ReadClientStartupInfoHistory(client_id)
self.assertEqual(len(history), 3)
self.assertEqual(startup_infos, history)
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 testReadClientStartupInfo(self):
d = self.db
client_id = self.InitializeClient()
d.WriteClientStartupInfo(client_id, rdf_client.StartupInfo(boot_time=1337))
d.WriteClientStartupInfo(client_id, rdf_client.StartupInfo(boot_time=2000))
last_is = d.ReadClientStartupInfo(client_id)
self.assertIsInstance(last_is, rdf_client.StartupInfo)
self.assertEqual(last_is.boot_time, 2000)
self.assertIsInstance(last_is.timestamp, rdfvalue.RDFDatetime)
md = self.db.ReadClientMetadata(client_id)
self.assertEqual(md.startup_info_timestamp, last_is.timestamp)
def testWriteClientSnapshotHistoryUpdatesOnlyLastClientTimestamp(self):
client_id = self.InitializeClient()
client_old = rdf_objects.ClientSnapshot(client_id=client_id)
client_old.kernel = "1.0.0"
client_old.startup_info.client_info.client_name = "foo"
self.db.WriteClientSnapshot(client_old)
old_timestamp = self.db.ReadClientSnapshot(client_id).timestamp
startup_info = rdf_client.StartupInfo()
startup_info.client_info.client_name = "bar"
self.db.WriteClientStartupInfo(client_id, startup_info)
startup_timestamp = self.db.ReadClientStartupInfo(client_id).timestamp
client_new = rdf_objects.ClientSnapshot(client_id=client_id)
client_new.kernel = "2.0.0"
client_new.startup_info.client_info.client_name = "baz"
client_new.timestamp = rdfvalue.RDFDatetime.Lerp(
0.5, start_time=old_timestamp, end_time=startup_timestamp)
self.db.WriteClientSnapshotHistory([client_new])
info = self.db.ReadClientFullInfo(client_id)
last_snapshot = info.last_snapshot
last_startup_info = info.last_startup_info
self.assertEqual(last_snapshot.kernel, "2.0.0")
self.assertEqual(last_snapshot.startup_info.client_info.client_name,
"baz")
self.assertEqual(last_snapshot.timestamp, client_new.timestamp)
self.assertEqual(last_startup_info.client_info.client_name, "bar")
self.assertEqual(last_startup_info.timestamp, startup_timestamp)
def testReadClientStartupInfoHistory(self):
d = self.db
client_id = self.InitializeClient()
d.WriteClientStartupInfo(client_id, rdf_client.StartupInfo(boot_time=1))
d.WriteClientStartupInfo(client_id, rdf_client.StartupInfo(boot_time=2))
d.WriteClientStartupInfo(client_id, rdf_client.StartupInfo(boot_time=3))
hist = d.ReadClientStartupInfoHistory(client_id)
self.assertEqual(len(hist), 3)
self.assertEqual([si.boot_time for si in hist], [3, 2, 1])
self.assertIsInstance(hist[0].timestamp, rdfvalue.RDFDatetime)
self.assertGreater(hist[0].timestamp, hist[1].timestamp)
self.assertGreater(hist[1].timestamp, hist[2].timestamp)
md = self.db.ReadClientMetadata(client_id)
self.assertEqual(md.startup_info_timestamp, hist[0].timestamp)
def _SetUpReadClientStartupInfoHistoryTest(self):
d = self.db
self.client_id = self.InitializeClient()
timestamps = [rdfvalue.RDFDatetime.Now()]
si = rdf_client.StartupInfo(boot_time=1)
d.WriteClientStartupInfo(self.client_id, si)
timestamps.append(d.ReadClientStartupInfo(self.client_id).timestamp)
timestamps.append(rdfvalue.RDFDatetime.Now())
si = rdf_client.StartupInfo(boot_time=2)
d.WriteClientStartupInfo(self.client_id, si)
timestamps.append(d.ReadClientStartupInfo(self.client_id).timestamp)
timestamps.append(rdfvalue.RDFDatetime.Now())
return timestamps
def Run(self, unused_arg, ttl=None):
"""Returns the startup information."""
logging.debug("Sending startup information.")
boot_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(psutil.boot_time())
response = rdf_client.StartupInfo(
boot_time=boot_time, client_info=GetClientInformation())
self.grr_worker.SendReply(
response,
session_id=self.well_known_session_id,
response_id=0,
request_id=0,
message_type=rdf_flows.GrrMessage.Type.MESSAGE,
require_fastpoll=False,
ttl=ttl)
def _SetupFullInfoClients(self):
for i in range(10):
client_id = self.InitializeClient("C.000000005000000%d" % i)
cl = rdf_objects.ClientSnapshot(
client_id=client_id,
knowledge_base=rdf_client.KnowledgeBase(fqdn="test%d.examples.com" %
i),
kernel="12.3.%d" % i)
self.db.WriteClientSnapshot(cl)
self.db.WriteClientMetadata(client_id, certificate=CERT)
si = rdf_client.StartupInfo(boot_time=i)
self.db.WriteClientStartupInfo(client_id, si)
self.db.AddClientLabels(
client_id, "test_owner",
["test_label-a-%d" % i, "test_label-b-%d" % i])
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)
def testClientInfo(self, db: abstract_db.Database):
client_id = "C.0123456789ABCDEF"
db.WriteClientMetadata(client_id, fleetspeak_enabled=False)
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 = "FEDCBA9876543210"
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)
def testReadClientFullInfoReturnsCorrectResult(self):
d = self.db
client_id = self.InitializeClient()
cl = rdf_objects.ClientSnapshot(
client_id=client_id,
knowledge_base=rdf_client.KnowledgeBase(fqdn="test1234.examples.com"),
kernel="12.3")
d.WriteClientSnapshot(cl)
d.WriteClientMetadata(client_id, certificate=CERT)
si = rdf_client.StartupInfo(boot_time=1)
d.WriteClientStartupInfo(client_id, si)
d.AddClientLabels(client_id, "test_owner", ["test_label"])
full_info = d.ReadClientFullInfo(client_id)
self.assertEqual(full_info.last_snapshot, cl)
self.assertEqual(full_info.metadata.certificate, CERT)
self.assertEqual(full_info.last_startup_info, si)
self.assertEqual(
full_info.labels,
[rdf_objects.ClientLabel(owner="test_owner", name="test_label")])
