Merge pull request #16 from peppelinux/custom_sessiondb

Custom sessiondb

Author: rohe

src/oidcendpoint/endpoint_context.py

@@ -21,7 +21,7 @@
 from oidcendpoint.util import build_endpoints
 from oidcendpoint.util import importer
 
-LOGGER = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
 
 
 def add_path(url, path):
@@ -67,13 +67,36 @@ def init_service(conf, endpoint_context=None):
     return conf["class"](**kwargs)
 
 
+def get_token_handlers(conf):
+    th_args = conf.get("token_handler_args", None)
+    if not th_args:
+        # create 3 keys
+        keydef = [
+            {"type": "oct", "bytes": "24", "use": ["enc"], "kid": "code"},
+            {"type": "oct", "bytes": "24", "use": ["enc"], "kid": "token"},
+            {"type": "oct", "bytes": "24", "use": ["enc"], "kid": "refresh"},
+        ]
+
+        jwks_def = {
+            "private_path": "private/token_jwks.json",
+            "key_defs": keydef,
+            "read_only": False,
+        }
+        th_args = {"jwks_def": jwks_def}
+        for typ, tid in [("code", 600), ("token", 3600), ("refresh", 86400)]:
+            th_args[typ] = {"lifetime": tid}
+
+    return th_args
+
+
 class EndpointContext:
     def __init__(
             self,
             conf,
             keyjar=None,
             client_db=None,
             session_db=None,
+            sso_db=None,
             cwd="",
             cookie_dealer=None,
             httpc=None,
@@ -112,7 +135,13 @@ def __init__(
         # arguments for endpoints add-ons
         self.args = {}
 
+        # session db
         self._sub_func = None
+        self.sdb = session_db
+        if not self.sdb:
+            self.set_session_db(conf, sso_db)
+        #
+
         self.scope2claims = SCOPE2CLAIMS
 
         if cookie_name:
@@ -175,11 +204,6 @@ def __init__(
             if _func:
                 _func(self.conf)
 
-        if session_db:
-            self.sdb = session_db
-        else:
-            self.do_session_db(conf)
-
         _cap = self.do_endpoints(conf)
 
         for item in ["userinfo", "login_hint_lookup", "login_hint2acrs",
@@ -199,6 +223,14 @@ def __init__(
         # client registration access tokens
         self.registration_access_token = {}
 
+    def set_session_db(self, conf, sso_db=None):
+        # this populate self.sdb
+        sso_db = sso_db if sso_db else SSODb()
+        self.do_session_db(conf, sso_db)
+        # this append useinfo db to the session db
+        self.do_userinfo(conf)
+        logger.debug('Session DB: {}'.format(self.sdb.__dict__))
+
     def do_add_on(self, conf):
         if 'add_on' in conf:
             for spec in conf["add_on"].values():
@@ -233,8 +265,10 @@ def do_userinfo(self, conf):
         except KeyError:
             pass
         else:
-            self.userinfo = init_user_info(_conf, self.cwd)
-            self.sdb.userinfo = self.userinfo
+            if self.sdb:
+                self.userinfo = init_user_info(_conf, self.cwd)
+                self.sdb.userinfo = self.userinfo
+
 
     def do_id_token(self, conf):
         try:
@@ -262,44 +296,23 @@ def do_cookie_dealer(self, conf):
                 self.cookie_dealer = init_service(_conf)
 
     def do_sub_func(self, conf):
-        try:
-            _conf = conf["sub_func"]
-        except KeyError:
-            self._sub_func = None
-        else:
-            self._sub_func = {}
-            for key, args in _conf.items():
-                if "class" in args:
-                    self._sub_func[key] = init_service(args)
-                elif "function" in args:
-                    if isinstance(args["function"], str):
-                        self._sub_func[key] = util.importer(args["function"])
-                    else:
-                        self._sub_func[key] = args["function"]
-
-    def do_session_db(self, conf):
-        try:
-            _th_args = conf["token_handler_args"]
-        except KeyError:
-            # create 3 keys
-            keydef = [
-                {"type": "oct", "bytes": "24", "use": ["enc"], "kid": "code"},
-                {"type": "oct", "bytes": "24", "use": ["enc"], "kid": "token"},
-                {"type": "oct", "bytes": "24", "use": ["enc"], "kid": "refresh"},
-            ]
-
-            jwks_def = {
-                "private_path": "private/token_jwks.json",
-                "key_defs": keydef,
-                "read_only": False,
-            }
-
-            _th_args = {"jwks_def": jwks_def}
-            for typ, tid in [("code", 600), ("token", 3600), ("refresh", 86400)]:
-                _th_args[typ] = {"lifetime": tid}
+        _conf = conf.get("sub_func", {})
+        self._sub_func = {}
+        for key, args in _conf.items():
+            if "class" in args:
+                self._sub_func[key] = init_service(args)
+            elif "function" in args:
+                if isinstance(args["function"], str):
+                    self._sub_func[key] = util.importer(args["function"])
+                else:
+                    self._sub_func[key] = args["function"]
 
+    def do_session_db(self, conf, sso_db):
+        th_args = get_token_handlers(conf)
         self.sdb = create_session_db(
-            self, _th_args, db=None, sso_db=SSODb(), sub_func=self._sub_func
+            self, th_args, db=None,
+            sso_db=sso_db,
+            sub_func=self._sub_func
         )
 
     def do_endpoints(self, conf):

src/oidcendpoint/session.py

@@ -565,7 +565,8 @@ def get_authentication_event(self, sid):
                 raise ValueError("No Authn event info")
 
 
-def create_session_db(ec, token_handler_args, db=None, sso_db=SSODb(), sub_func=None):
+def create_session_db(ec, token_handler_args, db=None,
+                      sso_db=SSODb(), sub_func=None):
     _token_handler = token_handler.factory(ec, **token_handler_args)
 
     if not db: