Move blacklist from token handler to session info

Author: nsklikas

src/oidcendpoint/jwt_token.py

@@ -14,7 +14,6 @@ class JWTToken(Token):
     def __init__(
         self,
         typ,
-        black_list=None,
         keyjar=None,
         issuer=None,
         aud=None,
@@ -24,7 +23,7 @@ def __init__(
         token_type="Bearer",
         **kwargs
     ):
-        Token.__init__(self, typ, black_list, **kwargs)
+        Token.__init__(self, typ, **kwargs)
         self.token_type = token_type
         self.lifetime = lifetime
         self.args = kwargs
@@ -105,7 +104,6 @@ def info(self, token):
             "type": _payload["ttype"],
             "exp": _payload["exp"],
             "handler": self,
-            "black_listed": self.is_black_listed(token),
         }
         return _res
 

src/oidcendpoint/session.py

@@ -7,6 +7,7 @@
 from oidcmsg.message import OPTIONAL_LIST_OF_STRINGS
 from oidcmsg.message import SINGLE_OPTIONAL_STRING
 from oidcmsg.message import SINGLE_REQUIRED_STRING
+from oidcmsg.message import SINGLE_OPTIONAL_JSON
 from oidcmsg.message import msg_ser
 from oidcmsg.oidc import AuthorizationRequest
 
@@ -89,8 +90,27 @@ class SessionInfo(Message):
         "client_id": SINGLE_REQUIRED_STRING,
         "authn_event": SINGLE_REQUIRED_AUTHN_EVENT,
         "si_redirects": OPTIONAL_LIST_OF_STRINGS,
+        "black_list": SINGLE_OPTIONAL_JSON,
     }
 
+    def __init__(self, *args, **kwargs):
+        super(SessionInfo, self).__init__(*args, **kwargs)
+        self["black_list"] = {}
+
+    def is_black_listed(self, typ, token):
+        # If session is revoked
+        if "revoked" in self:
+            return True
+
+        return typ in self["black_list"] and token in self["black_list"][typ]
+
+    def black_list(self, typ):
+        if typ in self:
+            if typ in self["black_list"]:
+                self["black_list"][typ].append(self[typ])
+            else:
+                self["black_list"][typ] = [self[typ]]
+
 
 def pairwise_id(uid, sector_identifier, salt, **kwargs):
     return hashlib.sha256(
@@ -268,9 +288,9 @@ def do_sub(
 
         return sub
 
-    def is_valid(self, item):
+    def is_valid(self, typ, item):
         try:
-            return not self.handler.is_black_listed(item)
+            return not self[item].is_black_listed(typ, item)
         except KeyError:
             return False
 
@@ -295,9 +315,8 @@ def replace_token(self, sid, sinfo, token_type):
 
         if token_type in self.handler:
             refresh_token = self.handler[token_type](sid, sinfo=sinfo)
-            # blacklist the old is there is one
-            if sinfo.get(token_type):
-                self.handler[token_type].black_list(sinfo[token_type])
+            # blacklist the old
+            sinfo.black_list(token_type)
 
             sinfo[token_type] = refresh_token
         return sinfo
@@ -336,22 +355,20 @@ def upgrade_to_token(
             _tinfo = self.handler["code"].info(grant)
 
             session_info = self[_tinfo["sid"]]
+            key = _tinfo["sid"]
 
-            if self.handler["code"].is_black_listed(grant):
+            if session_info.is_black_listed("code", grant):
                 # invalidate the released access token and refresh token
                 for item in ["access_token", "refresh_token"]:
-                    try:
-                        self.handler[item].black_list(session_info[item])
-                    except KeyError:
-                        pass
+                    session_info.black_list(item)
+                    self[key] = session_info
                 raise AccessCodeUsed(grant)
 
             # mint a new access token
             _at = self._make_at(_tinfo["sid"], session_info)
 
             # make sure the code can't be used again
-            self.handler["code"].black_list(grant)
-            key = _tinfo["sid"]
+            session_info.black_list("code")
         else:
             session_info = self[key]
             _at = self._make_at(key, session_info)
@@ -392,11 +409,11 @@ def refresh_token(self, token, new_refresh=False):
         except KeyError:
             return False
 
-        if is_expired(int(_tinfo["exp"])) or _tinfo["black_listed"]:
-            raise ExpiredToken()
-
         _sid = _tinfo["sid"]
         session_info = self[_sid]
+        if is_expired(int(_tinfo["exp"])) or \
+           session_info.is_black_listed("refresh_token", token):
+            raise ExpiredToken()
 
         session_info = self.replace_token(_sid, session_info, "access_token")
 
@@ -420,11 +437,11 @@ def is_token_valid(self, token):
         except KeyError:
             return False
 
-        if is_expired(int(_tinfo["exp"])) or _tinfo["black_listed"]:
-            return False
-
         # Dependent on what state the session is in.
         session_info = self[_tinfo["sid"]]
+        if is_expired(int(_tinfo["exp"])) or \
+           session_info.is_black_listed("access_token", token):
+            return False
 
         if session_info["oauth_state"] == "authz":
             if _tinfo["handler"] != self.handler["code"]:
@@ -435,22 +452,24 @@ def is_token_valid(self, token):
 
         return True
 
-    def revoke_token(self, token, token_type=""):
+    def revoke_token(self, sid, token_type):
         """
-        Revokes access token
+        Revokes token
 
-        :param token: access token
+        :param sid: session id
+        :param token_type: token type, one of "code", "access_token" or
+            "refresh_token"
         """
-        if token_type:
-            self.handler[token_type].black_list(token)
-        else:
-            self.handler.black_list(token)
+        _sinfo = self[sid]
+        _sinfo.black_list(token_type)
+        self[sid] = _sinfo
 
     def revoke_all_tokens(self, token):
-        _sinfo = self[token]
+        sid = self.handler.sid(token)
+        _sinfo = self[sid]
         for typ in self.handler.keys():
-            if _sinfo.get(typ):
-                self.revoke_token(_sinfo[typ], typ)
+            _sinfo.black_list(typ)
+        self[sid] = _sinfo
 
     def revoke_session(self, sid="", token=""):
         """
@@ -465,11 +484,9 @@ def revoke_session(self, sid="", token=""):
             else:
                 raise ValueError('Need one of "sid" or "token"')
 
-        for typ in ["access_token", "refresh_token", "code"]:
-            try:
-                self.revoke_token(self[sid][typ], typ)
-            except KeyError:  # If no such token has been issued
-                pass
+        _sinfo = self[sid]
+        for typ in self.handler.keys():
+            _sinfo.black_list(typ)
 
         self.update(sid, revoked=True)
 

src/oidcendpoint/token_handler.py

@@ -69,11 +69,10 @@ def decrypt(self, ciphertext):
 
 
 class Token(object):
-    def __init__(self, typ, black_list=None, lifetime=300, **kwargs):
+    def __init__(self, typ, lifetime=300, **kwargs):
         self.type = typ
         self.lifetime = lifetime
         self.args = kwargs
-        self.blist = [] if black_list is None else black_list
 
     def __call__(self, sid, *args, **kwargs):
         """
@@ -113,19 +112,12 @@ def is_expired(self, token, when=0):
     def gather_args(self, *args, **kwargs):
         return {}
 
-    def black_list(self, token):
-        if token:
-            self.blist.append(token)
-
-    def is_black_listed(self, token):
-        return token in self.blist
-
 
 class DefaultToken(Token):
     def __init__(
-        self, password, typ="", black_list=None, token_type="Bearer", **kwargs
+        self, password, typ="", token_type="Bearer", **kwargs
     ):
-        Token.__init__(self, typ, black_list, **kwargs)
+        Token.__init__(self, typ, **kwargs)
         self.crypt = Crypt(password)
         self.token_type = token_type
 
@@ -189,7 +181,6 @@ def info(self, token):
             raise WrongTokenType(_res["type"])
         else:
             _res["handler"] = self
-            _res["black_listed"] = self.is_black_listed(token)
             return _res
 
     def is_expired(self, token, when=0):
@@ -239,10 +230,6 @@ def sid(self, token, order=None):
     def type(self, token, order=None):
         return self.info(token, order)["type"]
 
-    def is_black_listed(self, token, order=None):
-        _handler = self.get_handler(token, order)
-        return _handler.is_black_listed(token)
-
     def get_handler(self, token, order=None):
         if order is None:
             order = self.handler_order
@@ -257,10 +244,6 @@ def get_handler(self, token, order=None):
 
         return None
 
-    def black_list(self, token, order=None):
-        _handler = self.get_handler(token, order)
-        _handler.black_list(token)
-
     def keys(self):
         return self.handler.keys()
 

tests/test_04_token_handler.py

@@ -83,10 +83,8 @@ def test_default_token_info(self):
             "sid",
             "exp",
             "handler",
-            "black_listed",
         }
         assert _info["handler"] == self.th
-        assert _info["black_listed"] is False
 
     def test_is_expired(self):
         _token = self.th("another_id")
@@ -95,14 +93,6 @@ def test_is_expired(self):
         when = time.time() + 900
         assert self.th.is_expired(_token, when)
 
-    def test_blacklist(self):
-        _token = self.th("another_id")
-        self.th.black_list(_token)
-        assert self.th.is_black_listed(_token)
-
-        _info = self.th.info(_token)
-        assert _info["black_listed"] is True
-
 
 class TestTokenHandler(object):
     @pytest.fixture(autouse=True)
@@ -160,11 +150,5 @@ def test_get_handler(self):
         th = self.handler.get_handler(_token)
         assert th.type == "A"
 
-    def test_blacklist(self):
-        _token = self.handler["code"]("another_id")
-        self.handler.black_list(_token)
-        _info = self.handler.info(_token)
-        assert _info["black_listed"]
-
     def test_keys(self):
         assert set(self.handler.keys()) == {"access_token", "code", "refresh_token"}