diff --git a/src/saml2/mdstore.py b/src/saml2/mdstore.py index 30b69a2..d354c61 100644 --- a/src/saml2/mdstore.py +++ b/src/saml2/mdstore.py @@ -195,7 +195,7 @@ class MetaData(object): binding)) try: srvs = [] - for t in self.entity[entity_id][typ]: + for t in self[entity_id][typ]: try: srvs.extend(t[service]) except KeyError: @@ -223,7 +223,7 @@ class MetaData(object): def _ext_service(self, entity_id, typ, service, binding): try: - srvs = self.entity[entity_id][typ] + srvs = self[entity_id][typ] except KeyError: return None @@ -250,7 +250,7 @@ class MetaData(object): :return: """ res = {} - for ent in self.entity.keys(): + for ent in self.keys(): bind = self._service(ent, typ, service, binding) if bind: res[ent] = bind @@ -280,7 +280,7 @@ class MetaData(object): res = {"required": [], "optional": []} try: - for sp in self.entity[entity_id]["spsso_descriptor"]: + for sp in self[entity_id]["spsso_descriptor"]: _res = attribute_requirement(sp) res["required"].extend(_res["required"]) res["optional"].extend(_res["optional"]) @@ -290,22 +290,22 @@ class MetaData(object): return res def dumps(self): - return json.dumps(self.entity, indent=2) + return json.dumps(self.items(), indent=2) def with_descriptor(self, descriptor): res = {} desc = "%s_descriptor" % descriptor - for eid, ent in self.entity.items(): + for eid, ent in self.items(): if desc in ent: res[eid] = ent return res def __str__(self): - return "%s" % (self.entity,) + return "%s" % self.items() def construct_source_id(self): res = {} - for eid, ent in self.entity.items(): + for eid, ent in self.items(): for desc in ["spsso_descriptor", "idpsso_descriptor"]: try: for srv in ent[desc]: diff --git a/src/saml2/mongo_store.py b/src/saml2/mongo_store.py index c60d888..ce76d34 100644 --- a/src/saml2/mongo_store.py +++ b/src/saml2/mongo_store.py @@ -3,7 +3,7 @@ import logging from pymongo import MongoClient from saml2.eptid import Eptid -from saml2.mdstore import MetaData +from saml2.mdstore import MetaData, attribute_requirement from saml2.s_utils import PolicyError from saml2.ident import code, IdentDB, Unknown @@ -249,7 +249,6 @@ class MDB(object): #------------------------------------------------------------------------------ class EptidMDB(Eptid): - def __init__(self, secret, collection="", sub_collection="eptid"): Eptid.__init__(self, secret) self.mdb = MDB(collection, sub_collection) @@ -269,6 +268,16 @@ class EptidMDB(Eptid): #------------------------------------------------------------------------------ +def export_mdstore_to_mongo_db(mds, onts, collection, sub_collection=""): + mdb = MDB(collection, sub_collection) + mdb.primary_key = "entity_id" + for key, desc in mds.items(): + kwargs = { + "entity_description": to_dict(desc, onts.values(), True), + } + mdb.store(key, **kwargs) + + class MetadataMDB(MetaData): def __init__(self, onts, attrc, collection="", sub_collection=""): MetaData.__init__(self, onts, attrc) @@ -291,7 +300,7 @@ class MetadataMDB(MetaData): def _ext_service(self, entity_id, typ, service, binding): try: - srvs = self.entity[entity_id][typ] + srvs = self[entity_id][typ] except KeyError: return None @@ -312,22 +321,23 @@ class MetadataMDB(MetaData): pass def items(self): - return self.mdb.items() + for key, item in self.mdb.items(): + yield key, from_dict(item["entity_description"], self.onts, True) def keys(self): return self.mdb.keys() def __contains__(self, item): - pass + return item in self.mdb - def attribute_requirement(self): - pass - - def with_descriptor(self): - pass - - def construct_source_id(self): - pass + def __getitem__(self, item): + res = self.mdb.get(item) + if not res: + raise KeyError(item) + elif len(res) == 1: + return from_dict(res[0]["entity_description"], self.onts, True) + else: + raise CorruptDatabase("More then one document with key %s" % item) def bindings(self, entity_id, typ, service): pass \ No newline at end of file