def testClientLabels(self):
d = self.db
client_id = "C.0000000000000001"
self._InitializeClient(client_id)
self.assertEqual(d.GetClientLabels(client_id), [])
d.AddClientLabels(client_id, "owner1", ["label1"])
d.AddClientLabels(client_id, "owner2", ["label2", "label3"])
all_labels = [
objects.ClientLabel(name="label1", owner="owner1"),
objects.ClientLabel(name="label2", owner="owner2"),
objects.ClientLabel(name="label3", owner="owner2")
]
self.assertEqual(d.GetClientLabels(client_id), all_labels)
self.assertEqual(d.GetClientLabels("C.0000000000000002"), [])
# Can't hurt to insert this one again.
d.AddClientLabels(client_id, "owner1", ["label1"])
self.assertEqual(d.GetClientLabels(client_id), all_labels)
d.RemoveClientLabels(client_id, "owner1", ["does not exist"])
self.assertEqual(d.GetClientLabels(client_id), all_labels)
# Label3 is actually owned by owner2.
d.RemoveClientLabels(client_id, "owner1", ["label3"])
self.assertEqual(d.GetClientLabels(client_id), all_labels)
d.RemoveClientLabels(client_id, "owner2", ["label3"])
self.assertEqual(d.GetClientLabels(client_id), [
objects.ClientLabel(name="label1", owner="owner1"),
objects.ClientLabel(name="label2", owner="owner2"),
])
def testClientLabelsUnicode(self):
d = self.db
client_id = "C.0000000000000001"
self._InitializeClient(client_id)
self.assertEqual(d.GetClientLabels(client_id), [])
d.AddClientLabels(client_id, "owner1", [u"⛄࿄1"])
d.AddClientLabels(client_id, "owner2", [u"⛄࿄2"])
d.AddClientLabels(client_id, "owner2", [utils.SmartStr(u"⛄࿄3")])
all_labels = [
objects.ClientLabel(name=u"⛄࿄1", owner="owner1"),
objects.ClientLabel(name=u"⛄࿄2", owner="owner2"),
objects.ClientLabel(name=u"⛄࿄3", owner="owner2")
]
self.assertEqual(d.GetClientLabels(client_id), all_labels)
d.RemoveClientLabels(client_id, "owner1", ["does not exist"])
self.assertEqual(d.GetClientLabels(client_id), all_labels)
# This label is actually owned by owner2.
d.RemoveClientLabels(client_id, "owner1", [u"⛄࿄3"])
self.assertEqual(d.GetClientLabels(client_id), all_labels)
d.RemoveClientLabels(client_id, "owner2", [u"⛄࿄3"])
self.assertEqual(d.GetClientLabels(client_id), [
objects.ClientLabel(name=u"⛄࿄1", owner="owner1"),
objects.ClientLabel(name=u"⛄࿄2", owner="owner2")
])
def testClientLabelsUnicode(self):
d = self.CreateDatabase()
client_id = "C.0000000000000001"
d.WriteClientMetadata(client_id, fleetspeak_enabled=True)
self.assertEqual(d.GetClientLabels(client_id), [])
d.AddClientLabels(client_id, "owner1", [u"⛄࿄1"])
d.AddClientLabels(client_id, "owner2", [u"⛄࿄2"])
d.AddClientLabels(client_id, "owner2", [utils.SmartStr(u"⛄࿄3")])
all_labels = [
objects.ClientLabel(name=u"⛄࿄1", owner="owner1"),
objects.ClientLabel(name=u"⛄࿄2", owner="owner2"),
objects.ClientLabel(name=u"⛄࿄3", owner="owner2")
]
self.assertEqual(d.GetClientLabels(client_id), all_labels)
d.RemoveClientLabels(client_id, "owner1", ["does not exist"])
self.assertEqual(d.GetClientLabels(client_id), all_labels)
# This label is actually owned by owner2.
d.RemoveClientLabels(client_id, "owner1", [u"⛄࿄3"])
self.assertEqual(d.GetClientLabels(client_id), all_labels)
d.RemoveClientLabels(client_id, "owner2", [u"⛄࿄3"])
self.assertEqual(d.GetClientLabels(client_id), [
objects.ClientLabel(name=u"⛄࿄1", owner="owner1"),
objects.ClientLabel(name=u"⛄࿄2", owner="owner2")
])
def testReadAllClientsFullInfoReadsMultipleClientsWithMultipleLabels(self):
d = self.db
for i in range(10):
client_id = "C.000000005000000%d" % i
self._InitializeClient(client_id)
cl = objects.ClientSnapshot(
client_id=client_id,
knowledge_base=rdf_client.KnowledgeBase(
fqdn="test%d.examples.com" % i),
kernel="12.3.%d" % i)
d.WriteClientSnapshot(cl)
d.WriteClientMetadata(client_id, certificate=CERT)
si = rdf_client.StartupInfo(boot_time=i)
d.WriteClientStartupInfo(client_id, si)
d.AddClientLabels(
client_id, "test_owner",
["test_label-a-%d" % i, "test_label-b-%d" % i])
c_infos = sorted(
d.ReadAllClientsFullInfo(), key=lambda c: c.last_snapshot.client_id)
for i, full_info in enumerate(c_infos):
self.assertEqual(full_info.last_snapshot.client_id,
"C.000000005000000%d" % i)
self.assertEqual(full_info.metadata.certificate, CERT)
self.assertEqual(full_info.last_startup_info.boot_time, i)
self.assertEqual(
sorted(full_info.labels, key=lambda l: l.name), [
objects.ClientLabel(
owner="test_owner", name="test_label-a-%d" % i),
objects.ClientLabel(
owner="test_owner", name="test_label-b-%d" % i)
])
def testClientLabels(self):
d = self.CreateDatabase()
client_id = "C.0000000000000001"
d.WriteClientMetadata(client_id, fleetspeak_enabled=True)
self.assertEqual(d.GetClientLabels(client_id), [])
d.AddClientLabels(client_id, "owner1", ["label1"])
d.AddClientLabels(client_id, "owner2", ["label2", "label3"])
all_labels = [
objects.ClientLabel(name="label1", owner="owner1"),
objects.ClientLabel(name="label2", owner="owner2"),
objects.ClientLabel(name="label3", owner="owner2")
]
self.assertEqual(d.GetClientLabels(client_id), all_labels)
self.assertEqual(d.GetClientLabels("C.0000000000000002"), [])
# Can't hurt to insert this one again.
d.AddClientLabels(client_id, "owner1", ["label1"])
self.assertEqual(d.GetClientLabels(client_id), all_labels)
d.RemoveClientLabels(client_id, "owner1", ["does not exist"])
self.assertEqual(d.GetClientLabels(client_id), all_labels)
# Label3 is actually owned by owner2.
d.RemoveClientLabels(client_id, "owner1", ["label3"])
self.assertEqual(d.GetClientLabels(client_id), all_labels)
d.RemoveClientLabels(client_id, "owner2", ["label3"])
self.assertEqual(d.GetClientLabels(client_id), [
objects.ClientLabel(name="label1", owner="owner1"),
objects.ClientLabel(name="label2", owner="owner2"),
])
def _VerifyFullInfos(self, c_infos):
c_infos = sorted(c_infos, key=lambda c: c.last_snapshot.client_id)
for i, full_info in enumerate(c_infos):
self.assertEqual(full_info.last_snapshot.client_id,
"C.000000005000000%d" % i)
self.assertEqual(full_info.metadata.certificate, CERT)
self.assertEqual(full_info.last_startup_info.boot_time, i)
self.assertEqual(sorted(full_info.labels, key=lambda l: l.name), [
objects.ClientLabel(owner="test_owner",
name="test_label-a-%d" % i),
objects.ClientLabel(owner="test_owner",
name="test_label-b-%d" % i)
])
def testReadAllLabelsReturnsLabelsFromMultipleClients(self):
d = self.db
client_id_1 = self.InitializeClient()
client_id_2 = self.InitializeClient()
d.AddClientLabels(client_id_1, "owner1", [u"foo"])
d.AddClientLabels(client_id_2, "owner1", [u"foo"])
d.AddClientLabels(client_id_1, "owner2", [u"bar"])
d.AddClientLabels(client_id_2, "owner2", [u"bar"])
all_labels = sorted(d.ReadAllClientLabels(), key=lambda l: l.name)
self.assertEqual(all_labels, [
objects.ClientLabel(name="bar", owner="owner2"),
objects.ClientLabel(name="foo", owner="owner1")
])
def GetClientLabels(self, client_id):
res = []
owner_dict = self.labels.get(client_id, {})
for owner, labels in owner_dict.items():
for l in labels:
res.append(objects.ClientLabel(owner=owner, name=l))
return sorted(res, key=lambda label: (label.owner, label.name))
def ReadAllClientLabels(self):
result = set()
for labels_dict in self.labels.values():
for owner, names in labels_dict.items():
for name in names:
result.add(objects.ClientLabel(owner=owner, name=name))
return list(result)
def HandleRelationalDB(self, args, token=None):
labels = data_store.REL_DB.ReadAllClientLabels()
label_objects = []
for name in set(l.name for l in labels):
label_objects.append(objects.ClientLabel(name=name))
return ApiListClientsLabelsResult(
items=sorted(label_objects, key=lambda l: l.name))
def _ResponseToClientsFullInfo(self, response):
c_full_info = None
prev_cid = None
for row in response:
(cid, fs, crt, ping, clk, ip, foreman, first, last_client_ts,
last_crash_ts, last_startup_ts, client_obj, client_startup_obj,
last_startup_obj, label_owner, label_name) = row
if cid != prev_cid:
if c_full_info:
yield _IntToClientID(prev_cid), c_full_info
metadata = objects.ClientMetadata(
certificate=crt,
fleetspeak_enabled=fs,
first_seen=_MysqlToRDFDatetime(first),
ping=_MysqlToRDFDatetime(ping),
clock=_MysqlToRDFDatetime(clk),
ip=_StringToRDFProto(rdf_client.NetworkAddress, ip),
last_foreman_time=_MysqlToRDFDatetime(foreman),
startup_info_timestamp=_MysqlToRDFDatetime(
last_startup_ts),
last_crash_timestamp=_MysqlToRDFDatetime(last_crash_ts))
if client_obj is not None:
l_snapshot = objects.ClientSnapshot.FromSerializedString(
client_obj)
l_snapshot.timestamp = _MysqlToRDFDatetime(last_client_ts)
l_snapshot.startup_info = rdf_client.StartupInfo.FromSerializedString(
client_startup_obj)
l_snapshot.startup_info.timestamp = l_snapshot.timestamp
else:
l_snapshot = objects.ClientSnapshot(
client_id=_IntToClientID(cid))
if last_startup_obj is not None:
startup_info = rdf_client.StartupInfo.FromSerializedString(
last_startup_obj)
startup_info.timestamp = _MysqlToRDFDatetime(
last_startup_ts)
else:
startup_info = None
prev_cid = cid
c_full_info = objects.ClientFullInfo(
metadata=metadata,
labels=[],
last_snapshot=l_snapshot,
last_startup_info=startup_info)
if label_owner and label_name:
c_full_info.labels.append(
objects.ClientLabel(name=label_name, owner=label_owner))
if c_full_info:
yield _IntToClientID(prev_cid), c_full_info
def testReadAllLabelsReturnsLabelsFromSingleClient(self):
d = self.db
client_id = self.InitializeClient()
d.AddClientLabels(client_id, "owner1", [u"foo"])
all_labels = d.ReadAllClientLabels()
self.assertEqual(all_labels,
[objects.ClientLabel(name="foo", owner="owner1")])
def Handle(self, args, token=None):
labels_index = aff4.FACTORY.Create(standard.LabelSet.CLIENT_LABELS_URN,
standard.LabelSet,
mode="r",
token=token)
label_objects = []
for label in labels_index.ListLabels():
label_objects.append(objects.ClientLabel(name=label))
return ApiListClientsLabelsResult(items=label_objects)
def MultiReadClientLabels(self, client_ids):
res = {}
for client_id in client_ids:
res[client_id] = []
owner_dict = self.labels.get(client_id, {})
for owner, labels in owner_dict.items():
for l in labels:
res[client_id].append(
objects.ClientLabel(owner=owner, name=l))
res[client_id].sort(key=lambda label: (label.owner, label.name))
return res
def ReadAllClientLabels(self, cursor=None):
"""Reads the user labels for a list of clients."""
cursor.execute("SELECT DISTINCT owner, label FROM client_labels")
result = []
for owner, label in cursor.fetchall():
result.append(
objects.ClientLabel(name=utils.SmartUnicode(label), owner=owner))
result.sort(key=lambda label: (label.owner, label.name))
return result
def InitFromAff4Object(self, client_obj):
self.urn = client_obj.urn
self.client_id = self.urn.Basename()
self.agent_info = client_obj.Get(client_obj.Schema.CLIENT_INFO)
self.hardware_info = client_obj.Get(client_obj.Schema.HARDWARE_INFO)
self.os_info = rdf_client.Uname(
system=client_obj.Get(client_obj.Schema.SYSTEM),
release=client_obj.Get(client_obj.Schema.OS_RELEASE),
# TODO(user): Check if ProtoString.Validate should be fixed
# to do an isinstance() check on a value. Is simple type
# equality check used there for performance reasons?
version=utils.SmartStr(
client_obj.Get(client_obj.Schema.OS_VERSION, "")),
kernel=client_obj.Get(client_obj.Schema.KERNEL),
machine=client_obj.Get(client_obj.Schema.ARCH),
fqdn=(client_obj.Get(client_obj.Schema.FQDN)
or client_obj.Get(client_obj.Schema.HOSTNAME)),
install_date=client_obj.Get(client_obj.Schema.INSTALL_DATE))
self.knowledge_base = client_obj.Get(client_obj.Schema.KNOWLEDGE_BASE)
self.memory_size = client_obj.Get(client_obj.Schema.MEMORY_SIZE)
self.first_seen_at = client_obj.Get(client_obj.Schema.FIRST_SEEN)
self.last_seen_at = client_obj.Get(client_obj.Schema.PING)
self.last_booted_at = client_obj.Get(client_obj.Schema.LAST_BOOT_TIME)
self.last_clock = client_obj.Get(client_obj.Schema.CLOCK)
last_crash = client_obj.Get(client_obj.Schema.LAST_CRASH)
if last_crash is not None:
self.last_crash_at = last_crash.timestamp
self.labels = [
objects.ClientLabel(name=l.name, owner=l.owner)
for l in client_obj.GetLabels()
]
self.interfaces = client_obj.Get(client_obj.Schema.INTERFACES)
kb = client_obj.Get(client_obj.Schema.KNOWLEDGE_BASE)
self.users = kb and kb.users or []
self.volumes = client_obj.Get(client_obj.Schema.VOLUMES)
type_obj = client_obj.Get(client_obj.Schema.TYPE)
if type_obj:
# Without self.Set self.age would reference "age" attribute instead of a
# protobuf field.
self.Set("age", type_obj.age)
self.cloud_instance = client_obj.Get(client_obj.Schema.CLOUD_INSTANCE)
return self
def MultiReadClientLabels(self, client_ids, cursor=None):
"""Reads the user labels for a list of clients."""
int_ids = [_ClientIDToInt(cid) for cid in client_ids]
query = ("SELECT client_id, owner, label "
"FROM client_labels "
"WHERE client_id IN ({})").format(", ".join(
["%s"] * len(client_ids)))
ret = {client_id: [] for client_id in client_ids}
cursor.execute(query, int_ids)
for client_id, owner, label in cursor.fetchall():
ret[_IntToClientID(client_id)].append(
objects.ClientLabel(name=utils.SmartUnicode(label), owner=owner))
for r in ret.values():
r.sort(key=lambda label: (label.owner, label.name))
return ret
def testReadClientFullInfoReturnsCorrectResult(self):
d = self.db
client_id = self.InitializeClient()
cl = 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,
[objects.ClientLabel(owner="test_owner", name="test_label")])
def testFullInfo(self):
d = self.db
client_id = "C.0000000050000001"
self._InitializeClient(client_id)
cl = objects.Client(client_id=client_id,
knowledge_base=rdf_client.KnowledgeBase(
fqdn="test1234.examples.com"),
kernel="12.3")
d.WriteClient(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.ReadFullInfoClient(client_id)
self.assertEqual(full_info["client"], cl)
self.assertEqual(full_info["metadata"].certificate, CERT)
self.assertEqual(full_info["last_startup_info"], si)
self.assertEqual(
full_info["labels"],
[objects.ClientLabel(owner="test_owner", name="test_label")])