From 2b5e897ac4e9c92cd3a4073db0e4e05a592270b1 Mon Sep 17 00:00:00 2001
From: YuviPanda <yuvipanda@gmail.com>
Date: Thu, 18 Jan 2024 14:17:51 -0800
Subject: [PATCH] Make username_claim callable in CILogon

Companion to https://github.com/jupyterhub/oauthenticator/pull/717

Fixes https://github.com/jupyterhub/oauthenticator/issues/712
---
 oauthenticator/cilogon.py                  |  6 +++++-
 oauthenticator/schemas/cilogon-schema.yaml |  4 ++--
 oauthenticator/tests/test_cilogon.py       | 23 ++++++++++++++++++++++
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/oauthenticator/cilogon.py b/oauthenticator/cilogon.py
index b11739fc..832c0179 100644
--- a/oauthenticator/cilogon.py
+++ b/oauthenticator/cilogon.py
@@ -363,7 +363,11 @@ def _user_info_to_unprocessed_username(self, user_info):
         username_derivation = self.allowed_idps[user_idp]["username_derivation"]
         username_claim = username_derivation["username_claim"]
 
-        username = user_info.get(username_claim)
+        if callable(username_claim):
+            username = username_claim(user_info)
+        else:
+            username = user_info.get(self.username_claim, None)
+
         if not username:
             message = f"Configured username_claim {username_claim} for {user_idp} was not found in the response {user_info.keys()}"
             self.log.error(message)
diff --git a/oauthenticator/schemas/cilogon-schema.yaml b/oauthenticator/schemas/cilogon-schema.yaml
index 2f8436ff..59f2d701 100644
--- a/oauthenticator/schemas/cilogon-schema.yaml
+++ b/oauthenticator/schemas/cilogon-schema.yaml
@@ -23,8 +23,8 @@ properties:
     required:
       - username_claim
     properties:
-      username_claim:
-        type: string
+      # FIXME: This needs to take a string or a callable
+      username_claim: {}
       action:
         type: string
         enum:
diff --git a/oauthenticator/tests/test_cilogon.py b/oauthenticator/tests/test_cilogon.py
index 4dac4927..a448a932 100644
--- a/oauthenticator/tests/test_cilogon.py
+++ b/oauthenticator/tests/test_cilogon.py
@@ -100,6 +100,29 @@ async def test_cilogon(
     else:
         assert auth_model == None
 
+async def test_username_claim_callable(
+    cilogon_client,
+):
+    c = Config()
+    c.CILogonOAuthenticator = Config()
+
+    c.CILogonOAuthenticator.allowed_idps = {
+        "https://some-idp.com/login/oauth/authorize": {
+            "username_derivation": {
+                "username_claim": lambda user_info: f"prefixed-{user_info['username']}",
+            },
+        },
+    }
+
+
+    authenticator = CILogonOAuthenticator(config=c)
+
+    handled_user_model = user_model("user1", "username")
+    handler = cilogon_client.handler_for_user(handled_user_model)
+    auth_model = await authenticator.get_authenticated_user(handler, None)
+
+    assert auth_model["name"] == f"prefixed-user1"
+
 
 @mark.parametrize(
     "test_variation_id,idp_config,class_config,test_user_name,expect_allowed,expect_admin",