Add claims per client in access token

Author: nsklikas

src/oidcendpoint/jwt_token.py

@@ -32,10 +32,14 @@ def __init__(
 
         self.key_jar = keyjar or ec.keyjar
         self.issuer = issuer or ec.issuer
+        self.cdb = ec.cdb
 
         self.def_aud = aud or []
         self.alg = alg
         self.scope_claims_map = kwargs.get("scope_claims_map", ec.scope2claims)
+        self.enable_claims_per_client = kwargs.get(
+            "enable_claims_per_client", False
+        )
 
     def add_claims(self, payload, uinfo, claims):
         for attr in claims:
@@ -47,7 +51,8 @@ def add_claims(self, payload, uinfo, claims):
                 pass
 
     def __call__(
-        self, sid: str, uinfo: Dict, sinfo: Dict, *args, aud: Optional[Any], **kwargs
+        self, sid: str, uinfo: Dict, sinfo: Dict, *args, aud: Optional[Any],
+        client_id: Optional[str], **kwargs
     ):
         """
         Return a token.
@@ -70,6 +75,12 @@ def __call__(
                     sinfo["authn_req"]["scope"], map=self.scope_claims_map
                 ).keys(),
             )
+        # Add claims if is access token
+        if self.type == 'T' and self.enable_claims_per_client:
+            client = self.cdb.get(client_id, {})
+            client_claims = client.get("access_token_claims")
+            if client_claims:
+                self.add_claims(payload, uinfo, client_claims)
 
         payload.update(kwargs)
         signer = JWT(

src/oidcendpoint/session.py

@@ -302,14 +302,15 @@ def replace_refresh_token(self, sid, sinfo):
 
     def _make_at(self, sid, session_info, aud=None, client_id_aud=True):
         uid = self.sso_db.get_uid_by_sid(sid)
-
-        uinfo = self.userinfo(uid, session_info["client_id"]) or {}
+        client_id = session_info["client_id"]
+        uinfo = self.userinfo(uid, client_id) or {}
         at_aud = aud or []
 
         if client_id_aud:
-            at_aud.append(session_info["client_id"])
+            at_aud.append(client_id)
         return self.handler["access_token"](
-            sid=sid, sinfo=session_info, uinfo=uinfo, aud=at_aud
+            sid=sid, sinfo=session_info, uinfo=uinfo, aud=at_aud,
+            client_id=client_id
         )
 
     def upgrade_to_token(

tests/test_27_jwt_token.py

@@ -189,6 +189,24 @@ def test_info(self):
         assert _info["type"] == "T"
         assert _info["sid"] == session_id
 
+    @pytest.mark.parametrize("enable_claims_per_client", [True, False])
+    def test_client_claims(self, enable_claims_per_client):
+        ec = self.endpoint.endpoint_context
+        handler = ec.sdb.handler.handler["access_token"]
+        session_id = setup_session(
+            ec, AUTH_REQ, uid="diana"
+        )
+        ec.cdb["client_1"]['access_token_claims'] = {
+            "address": None
+        }
+        handler.enable_claims_per_client = enable_claims_per_client
+        _dic = ec.sdb.upgrade_to_token(key=session_id)
+
+        token = _dic["access_token"]
+        _jwt = JWT(key_jar=KEYJAR, iss="client_1")
+        res = _jwt.unpack(token)
+        assert enable_claims_per_client is ("address" in res)
+
     def test_is_expired(self):
         session_id = setup_session(
             self.endpoint.endpoint_context, AUTH_REQ, uid="diana"