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")])