From 93a9bf3a6a68cd15d4e40d45612a03a738ad8c8d Mon Sep 17 00:00:00 2001 From: Fredrik Thulin Date: Tue, 11 Jun 2013 13:45:11 +0200 Subject: [PATCH 1/3] 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 | 71 +++++++++++++++++++++++++++++++--------- src/saml2/server.py | 7 ++-- 2 files changed, 60 insertions(+), 18 deletions(-) diff --git a/src/saml2/mongo_store.py b/src/saml2/mongo_store.py index 6fd3894..315f4a4 100644 --- a/src/saml2/mongo_store.py +++ b/src/saml2/mongo_store.py @@ -2,6 +2,8 @@ from hashlib import sha1 import logging from pymongo import MongoClient +from pymongo.mongo_replica_set_client import MongoReplicaSetClient +import pymongo.uri_parser from saml2.eptid import Eptid from saml2.mdstore import MetaData from saml2.s_utils import PolicyError @@ -49,10 +51,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 +130,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 +196,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 +250,52 @@ class MDB(object): self.db.drop() + +def _mdb_get_database(database, **kwargs): + """ + Helper-function to connect to MongoDB and return a database object. + + Performs explicit authentication if a username is provided in a connection + string URI, since PyMongo does not always seem to do that as promised. + + If database is *not* a connection to mongodb://localhost:27017 will be + made. + + :params database: name as string or (uri, name) + :returns: pymongo database object + """ + _conn = None + if isinstance(database, tuple): + uri, collection, = database + if uri.startswith("mongodb://"): + if "replicaSet=" in uri: + connection_factory = MongoReplicaSetClient + else: + connection_factory = MongoClient + + if not 'tz_aware' in kwargs: + # default, but not forced + kwargs['tz_aware'] = True + + _conn = connection_factory(uri, **kwargs) + if not _conn: + _conn = MongoClient() + + _parsed_uri = pymongo.uri_parser.parse_uri(uri) + _db = _conn[_parsed_uri.get("database")] + if _parsed_uri.get("username", None): + _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 +371,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 +422,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) From c16d9792c6fa1965822f3bfdc31309948ae1f0be Mon Sep 17 00:00:00 2001 From: Fredrik Thulin Date: Tue, 11 Jun 2013 15:04:26 +0200 Subject: [PATCH 2/3] _mdb_get_database: don't require tuple for URI --- src/saml2/mongo_store.py | 46 ++++++++++++++++++++++------------------ 1 file changed, 25 insertions(+), 21 deletions(-) diff --git a/src/saml2/mongo_store.py b/src/saml2/mongo_store.py index 315f4a4..607678e 100644 --- a/src/saml2/mongo_store.py +++ b/src/saml2/mongo_store.py @@ -4,6 +4,7 @@ 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 @@ -251,39 +252,42 @@ class MDB(object): -def _mdb_get_database(database, **kwargs): +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. - If database is *not* a connection to mongodb://localhost:27017 will be - made. - :params database: name as string or (uri, name) :returns: pymongo database object """ - _conn = None - if isinstance(database, tuple): - uri, collection, = database - if uri.startswith("mongodb://"): - if "replicaSet=" in uri: - connection_factory = MongoReplicaSetClient - else: - connection_factory = MongoClient + 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 + if not "tz_aware" in kwargs: + # default, but not forced + kwargs["tz_aware"] = True - _conn = connection_factory(uri, **kwargs) - if not _conn: - _conn = MongoClient() + _conn = connection_factory(uri, **kwargs) + _db = _conn[db_name] - _parsed_uri = pymongo.uri_parser.parse_uri(uri) - _db = _conn[_parsed_uri.get("database")] - if _parsed_uri.get("username", None): + if "username" in _parsed_uri: _db.authenticate( _parsed_uri.get("username", None), _parsed_uri.get("password", None) From 19da1324df5a6def9dd61b05824577e57b928fd5 Mon Sep 17 00:00:00 2001 From: Fredrik Thulin Date: Tue, 11 Jun 2013 16:13:31 +0200 Subject: [PATCH 3/3] _mdb_get_database: Restore 'uri is db name' functionality. Forgot to run the test suite after the last round of changes :(. On the good side, this commit should actually restore backwards compatibility with previous MongoDB support. --- src/saml2/mongo_store.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/saml2/mongo_store.py b/src/saml2/mongo_store.py index 607678e..4d7e204 100644 --- a/src/saml2/mongo_store.py +++ b/src/saml2/mongo_store.py @@ -266,25 +266,28 @@ def _mdb_get_database(uri, **kwargs): :params database: name as string or (uri, name) :returns: pymongo database object """ + if not "tz_aware" in kwargs: + # default, but not forced + kwargs["tz_aware"] = True + connection_factory = MongoClient _parsed_uri = {} db_name = None + _conn = 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 + _conn = MongoClient() pass else: if "replicaset" in _parsed_uri["options"]: connection_factory = MongoReplicaSetClient db_name = _parsed_uri.get("database", "pysaml2") + _conn = connection_factory(uri, **kwargs) - 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: