Using shelve based db for sso and session db.

Author: rohe

src/oidcendpoint/endpoint_context.py

@@ -241,9 +241,21 @@ def __init__(
             self.httpc_params = {"verify": conf.get("verify_ssl")}
 
     def set_session_db(self, sso_db=None, db=None):
-        sso_db = sso_db or SSODb()
+        if sso_db is None and self.conf.get("sso_db"):
+            _spec = self.conf.get("sso_db")
+            _kwargs = _spec.get("kwargs", {})
+            _db = importer(_spec["class"])(**_kwargs)
+            sso_db = SSODb(_db)
+        else:
+            sso_db = sso_db or SSODb()
+
+        if db is None and self.conf.get("session_db"):
+            _spec = self.conf.get("session_db")
+            _kwargs = _spec.get("kwargs", {})
+            db = importer(_spec["class"])(**_kwargs)
+
         self.do_session_db(sso_db, db)
-        # append useinfo db to the session db
+        # append userinfo db to the session db
         self.do_userinfo()
         logger.debug("Session DB: {}".format(self.sdb.__dict__))
 

tests/test_05_sso_db.py

@@ -1,4 +1,6 @@
 import pytest
+
+from oidcendpoint.shelve_db import ShelveDataBase
 from oidcendpoint.sso_db import SSODb
 
 
@@ -115,3 +117,120 @@ def test_get_sub_by_uid_different_sub(self):
         res = self.sso_db.get_subs_by_uid("Lizz")
 
         assert set(res) == {"abcdefgh", "012346789"}
+
+
+class TestSessionShelveDB(object):
+    @pytest.fixture(autouse=True)
+    def create_sdb(self):
+        # Create fresh database each time
+        _db = ShelveDataBase(filename='shelf', flag='n', writeback=True)
+        self.sso_db = SSODb(_db)
+
+    def test_map_sid2uid(self):
+        self.sso_db.map_sid2uid("session id 1", "Lizz")
+        assert self.sso_db.get_sids_by_uid("Lizz") == ["session id 1"]
+
+    def test_missing_map(self):
+        assert self.sso_db.get_sids_by_uid("Lizz") is None
+
+    def test_multiple_map_sid2uid(self):
+        self.sso_db.map_sid2uid("session id 1", "Lizz")
+        self.sso_db.map_sid2uid("session id 2", "Lizz")
+        assert set(self.sso_db.get_sids_by_uid("Lizz")) == {
+            "session id 1",
+            "session id 2",
+        }
+
+    def test_map_unmap_sid2uid(self):
+        self.sso_db.map_sid2uid("session id 1", "Lizz")
+        self.sso_db.map_sid2uid("session id 2", "Lizz")
+        assert set(self.sso_db.get_sids_by_uid("Lizz")) == {
+            "session id 1",
+            "session id 2",
+        }
+
+        self.sso_db.remove_sid2uid("session id 1", "Lizz")
+        assert self.sso_db.get_sids_by_uid("Lizz") == ["session id 2"]
+
+    def test_get_uid_by_sid(self):
+        self.sso_db.map_sid2uid("session id 1", "Lizz")
+        self.sso_db.map_sid2uid("session id 2", "Lizz")
+
+        assert self.sso_db.get_uid_by_sid("session id 1") == "Lizz"
+        assert self.sso_db.get_uid_by_sid("session id 2") == "Lizz"
+
+    def test_remove_uid(self):
+        self.sso_db.map_sid2uid("session id 1", "Lizz")
+        self.sso_db.map_sid2uid("session id 2", "Diana")
+
+        self.sso_db.remove_uid("Lizz")
+        assert self.sso_db.get_uid_by_sid("session id 1") is None
+        assert self.sso_db.get_sids_by_uid("Lizz") is None
+
+    def test_map_sid2sub(self):
+        self.sso_db.map_sid2sub("session id 1", "abcdefgh")
+        assert self.sso_db.get_sids_by_sub("abcdefgh") == ["session id 1"]
+
+    def test_missing_sid2sub_map(self):
+        assert self.sso_db.get_sids_by_sub("abcdefgh") is None
+
+    def test_multiple_map_sid2sub(self):
+        self.sso_db.map_sid2sub("session id 1", "abcdefgh")
+        self.sso_db.map_sid2sub("session id 2", "abcdefgh")
+        assert set(self.sso_db.get_sids_by_sub("abcdefgh")) == {
+            "session id 1",
+            "session id 2",
+        }
+
+    def test_map_unmap_sid2sub(self):
+        self.sso_db.map_sid2sub("session id 1", "abcdefgh")
+        self.sso_db.map_sid2sub("session id 2", "abcdefgh")
+        assert set(self.sso_db.get_sids_by_sub("abcdefgh")) == {
+            "session id 1",
+            "session id 2",
+        }
+
+        self.sso_db.remove_sid2sub("session id 1", "abcdefgh")
+        assert self.sso_db.get_sids_by_sub("abcdefgh") == ["session id 2"]
+
+    def test_get_sub_by_sid(self):
+        self.sso_db.map_sid2sub("session id 1", "abcdefgh")
+        self.sso_db.map_sid2sub("session id 2", "abcdefgh")
+
+        assert set(self.sso_db.get_sids_by_sub("abcdefgh")) == {
+            "session id 1",
+            "session id 2",
+        }
+
+    def test_remove_sub(self):
+        self.sso_db.map_sid2sub("session id 1", "abcdefgh")
+        self.sso_db.map_sid2sub("session id 2", "012346789")
+
+        self.sso_db.remove_sub("abcdefgh")
+        assert self.sso_db.get_sub_by_sid("session id 1") is None
+        assert self.sso_db.get_sids_by_sub("abcdefgh") is None
+        # have not touched the others
+        assert self.sso_db.get_sub_by_sid("session id 2") == "012346789"
+        assert self.sso_db.get_sids_by_sub("012346789") == ["session id 2"]
+
+    def test_get_sub_by_uid_same_sub(self):
+        self.sso_db.map_sid2sub("session id 1", "abcdefgh")
+        self.sso_db.map_sid2sub("session id 2", "abcdefgh")
+
+        self.sso_db.map_sid2uid("session id 1", "Lizz")
+        self.sso_db.map_sid2uid("session id 2", "Lizz")
+
+        res = self.sso_db.get_subs_by_uid("Lizz")
+
+        assert set(res) == {"abcdefgh"}
+
+    def test_get_sub_by_uid_different_sub(self):
+        self.sso_db.map_sid2sub("session id 1", "abcdefgh")
+        self.sso_db.map_sid2sub("session id 2", "012346789")
+
+        self.sso_db.map_sid2uid("session id 1", "Lizz")
+        self.sso_db.map_sid2uid("session id 2", "Lizz")
+
+        res = self.sso_db.get_subs_by_uid("Lizz")
+
+        assert set(res) == {"abcdefgh", "012346789"}