def testAddsSingleLabelToSingleClient(self):
for client_id in self.client_ids:
self.assertFalse(
aff4.FACTORY.Open(client_id, token=self.token).GetLabels())
data_store.REL_DB.WriteClientMetadata(client_id.Basename(),
fleetspeak_enabled=False)
self.handler.Handle(client_plugin.ApiAddClientsLabelsArgs(
client_ids=[self.client_ids[0]], labels=[u"foo"]),
token=self.token)
# AFF4 labels.
labels = aff4.FACTORY.Open(self.client_ids[0],
token=self.token).GetLabels()
self.assertLen(labels, 1)
self.assertEqual(labels[0].name, u"foo")
self.assertEqual(labels[0].owner, self.token.username)
for client_id in self.client_ids[1:]:
self.assertFalse(
aff4.FACTORY.Open(client_id, token=self.token).GetLabels())
# Relational DB labels.
cid = self.client_ids[0].Basename()
labels = data_store.REL_DB.ReadClientLabels(cid)
self.assertLen(labels, 1)
self.assertEqual(labels[0].name, u"foo")
self.assertEqual(labels[0].owner, self.token.username)
for client_id in self.client_ids[1:]:
self.assertFalse(
data_store.REL_DB.ReadClientLabels(client_id.Basename()))
def testAddsTwoLabelsToTwoClients(self):
for client_id in self.client_ids:
self.assertFalse(
aff4.FACTORY.Open(client_id, token=self.token).GetLabels())
data_store.REL_DB.WriteClientMetadata(
client_id.Basename(), fleetspeak_enabled=False)
self.handler.Handle(
client_plugin.ApiAddClientsLabelsArgs(
client_ids=[self.client_ids[0], self.client_ids[1]],
labels=["foo", "bar"]),
token=self.token)
# AFF4 labels.
for client_id in self.client_ids[:2]:
labels = aff4.FACTORY.Open(client_id, token=self.token).GetLabels()
self.assertEqual(len(labels), 2)
self.assertEqual(labels[0].name, "foo")
self.assertEqual(labels[0].owner, self.token.username)
self.assertEqual(labels[1].name, "bar")
self.assertEqual(labels[1].owner, self.token.username)
self.assertFalse(
aff4.FACTORY.Open(self.client_ids[2], token=self.token).GetLabels())
# Relational labels.
for client_id in self.client_ids[:2]:
labels = data_store.REL_DB.ReadClientLabels(client_id.Basename())
self.assertEqual(len(labels), 2)
self.assertEqual(labels[0].owner, self.token.username)
self.assertEqual(labels[1].owner, self.token.username)
self.assertItemsEqual([labels[0].name, labels[1].name], ["bar", "foo"])
self.assertFalse(
data_store.REL_DB.ReadClientLabels(self.client_ids[2].Basename()))
def testAddsSingleLabelToSingleClient(self):
self.handler.Handle(client_plugin.ApiAddClientsLabelsArgs(
client_ids=[self.client_ids[0]], labels=[u"foo"]),
token=self.token)
cid = self.client_ids[0]
labels = data_store.REL_DB.ReadClientLabels(cid)
self.assertLen(labels, 1)
self.assertEqual(labels[0].name, u"foo")
self.assertEqual(labels[0].owner, self.token.username)
for client_id in self.client_ids[1:]:
self.assertFalse(data_store.REL_DB.ReadClientLabels(client_id))
def testAddsTwoLabelsToTwoClients(self):
self.handler.Handle(
client_plugin.ApiAddClientsLabelsArgs(
client_ids=[self.client_ids[0], self.client_ids[1]],
labels=[u"foo", u"bar"]),
context=self.context)
for client_id in self.client_ids[:2]:
labels = data_store.REL_DB.ReadClientLabels(client_id)
self.assertLen(labels, 2)
self.assertEqual(labels[0].owner, self.context.username)
self.assertEqual(labels[1].owner, self.context.username)
self.assertCountEqual([labels[0].name, labels[1].name], [u"bar", u"foo"])
self.assertFalse(data_store.REL_DB.ReadClientLabels(self.client_ids[2]))
def testAuditEntryIsCreatedForEveryClient(self):
self.handler.Handle(client_plugin.ApiAddClientsLabelsArgs(
client_ids=self.client_ids, labels=[u"drei", u"ein", u"zwei"]),
token=self.token)
# We need to run .Simulate() so that the appropriate event is fired,
# collected, and finally written to the logs that we inspect.
mock_worker = worker_test_lib.MockWorker(token=self.token)
mock_worker.Simulate()
event = self._FindAuditEvent()
self.assertIsNotNone(event)
self.assertEqual(event.user, self.token.username)
self.assertEqual(
event.description, "%s.drei,%s.ein,%s.zwei" %
(self.token.username, self.token.username, self.token.username))
def _AddLabels(self, client_id, labels): args = client_plugin.ApiAddClientsLabelsArgs() args.client_ids = [client_id] args.labels = labels self.add_labels_handler.Handle(args=args, context=self.context)
