From 08980a034ef2e8de585920ac20991fcc53e44e69 Mon Sep 17 00:00:00 2001 From: Fredrik Thulin Date: Tue, 11 Jun 2013 13:45:11 +0200 Subject: [PATCH] Improve MongoDB support. * Support configuration with MongoDB connection URI strings. * Handle connecting to replica sets. * Use standard notation database+collection rather than collection+sub_collection. Note that these changes are NOT expected to be fully backwards compatible for someone using mongodb with pysaml2 already. --- src/saml2/mongo_store.py | 75 ++++++++++++++++++++++++++++++++-------- src/saml2/server.py | 7 ++-- 2 files changed, 64 insertions(+), 18 deletions(-) diff --git a/src/saml2/mongo_store.py b/src/saml2/mongo_store.py index 6fd3894..607678e 100644 --- a/src/saml2/mongo_store.py +++ b/src/saml2/mongo_store.py @@ -2,6 +2,9 @@ from hashlib import sha1 import logging from pymongo import MongoClient +from pymongo.mongo_replica_set_client import MongoReplicaSetClient +import pymongo.uri_parser +import pymongo.errors from saml2.eptid import Eptid from saml2.mdstore import MetaData from saml2.s_utils import PolicyError @@ -49,10 +52,9 @@ def context_match(cfilter, cntx): class SessionStorageMDB(object): """ Session information is stored in a MongoDB database""" - def __init__(self, collection=""): - connection = MongoClient() - db = connection[collection] - self.assertion = db.assertion + def __init__(self, database="", collection="assertion", **kwargs): + db = _mdb_get_database(database, **kwargs) + self.assertion = db[collection] def store_assertion(self, assertion, to_sign): name_id = assertion.subject.name_id @@ -129,9 +131,9 @@ class SessionStorageMDB(object): class IdentMDB(IdentDB): - def __init__(self, collection="", domain="", name_qualifier=""): + def __init__(self, database="", collection="ident", domain="", name_qualifier=""): IdentDB.__init__(self, None, domain, name_qualifier) - self.mdb = MDB(collection, "ident") + self.mdb = MDB(database=database, collection=collection) self.mdb.primary_key = "user_id" def in_store(self, _id): @@ -195,10 +197,9 @@ class IdentMDB(IdentDB): class MDB(object): primary_key = "mdb" - def __init__(self, collection="", sub_collection=""): - connection = MongoClient() - _db = connection[collection] - self.db = _db[sub_collection] + def __init__(self, database, collection, **kwargs): + _db = _mdb_get_database(database, **kwargs) + self.db = _db[collection] def store(self, value, **kwargs): if value: @@ -250,11 +251,55 @@ class MDB(object): self.db.drop() + +def _mdb_get_database(uri, **kwargs): + """ + Helper-function to connect to MongoDB and return a database object. + + The `uri' argument should be either a full MongoDB connection URI string, + or just a database name in which case a connection to the default mongo + instance at mongodb://localhost:27017 will be made. + + Performs explicit authentication if a username is provided in a connection + string URI, since PyMongo does not always seem to do that as promised. + + :params database: name as string or (uri, name) + :returns: pymongo database object + """ + connection_factory = MongoClient + _parsed_uri = {} + db_name = None + try: + _parsed_uri = pymongo.uri_parser.parse_uri(uri) + except pymongo.errors.InvalidURI: + # assume URI to be just the database name + db_name = uri + pass + else: + if "replicaset" in _parsed_uri["options"]: + connection_factory = MongoReplicaSetClient + db_name = _parsed_uri.get("database", "pysaml2") + + if not "tz_aware" in kwargs: + # default, but not forced + kwargs["tz_aware"] = True + + _conn = connection_factory(uri, **kwargs) + _db = _conn[db_name] + + if "username" in _parsed_uri: + _db.authenticate( + _parsed_uri.get("username", None), + _parsed_uri.get("password", None) + ) + + return _db + #------------------------------------------------------------------------------ class EptidMDB(Eptid): - def __init__(self, secret, collection="", sub_collection="eptid"): + def __init__(self, secret, database="", collection="eptid"): Eptid.__init__(self, secret) - self.mdb = MDB(collection, sub_collection) + self.mdb = MDB(database, collection) self.mdb.primary_key = "eptid_key" def __getitem__(self, key): @@ -330,9 +375,9 @@ def export_mdstore_to_mongo_db(mds, collection, sub_collection=""): class MetadataMDB(MetaData): - def __init__(self, onts, attrc, collection="", sub_collection=""): + def __init__(self, onts, attrc, database="", collection=""): MetaData.__init__(self, onts, attrc) - self.mdb = MDB(collection, sub_collection) + self.mdb = MDB(database, collection) self.mdb.primary_key = "entity_id" def _ext_service(self, entity_id, typ, service, binding): @@ -381,4 +426,4 @@ class MetadataMDB(MetaData): raise CorruptDatabase("More then one document with key %s" % item) def bindings(self, entity_id, typ, service): - pass \ No newline at end of file + pass diff --git a/src/saml2/server.py b/src/saml2/server.py index f28f0be..cb0d086 100644 --- a/src/saml2/server.py +++ b/src/saml2/server.py @@ -89,7 +89,7 @@ class Server(Entity): else: # Should be tuple typ, data = _spec if typ.lower() == "mongodb": - return SessionStorageMDB(data) + return SessionStorageMDB(database=data, collection="session") raise NotImplementedError("No such storage type implemented") @@ -120,7 +120,7 @@ class Server(Entity): elif typ == "dict": # in-memory dictionary idb = {} elif typ == "mongodb": - self.ident = IdentMDB(addr) + self.ident = IdentMDB(database=addr, collection="ident") if typ == "mongodb": pass @@ -140,7 +140,8 @@ class Server(Entity): if typ == "shelve": self.eptid = EptidShelve(secret, addr) elif typ == "mongodb": - self.eptid = EptidMDB(secret, addr, *dbspec[3:]) + self.eptid = EptidMDB(secret, database=addr, + collection="eptid") else: self.eptid = Eptid(secret)