refactor: Unify code legacy token exchange with DiracX

chrisburr · Oct 12, 2023 · fcaba13 · fcaba13
1 parent 6f73f3c
commit fcaba13
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 47 deletions.
diff --git a/src/DIRAC/FrameworkSystem/Service/ProxyManagerHandler.py b/src/DIRAC/FrameworkSystem/Service/ProxyManagerHandler.py
@@ -6,10 +6,8 @@
       :dedent: 2
       :caption: ProxyManager options
 """
-
-import os
-import requests
-from DIRAC import gLogger, S_OK, S_ERROR, gConfig
+from DIRAC import gLogger, S_OK, S_ERROR
+from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue
 from DIRAC.Core.DISET.RequestHandler import RequestHandler, getServiceOption
 from DIRAC.Core.Security import Properties
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
@@ -411,39 +409,12 @@ def export_getVOMSProxyWithToken(self, userDN, userGroup, requestPem, requiredLi
 
     types_exchangeProxyForToken = []
 
+    @convertToReturnValue
     def export_exchangeProxyForToken(self):
         """Exchange a proxy for an equivalent token to be used with diracx"""
+        from DIRAC.FrameworkSystem.Utilities.diracx import get_token
 
-        apiKey = gConfig.getValue("/DiracX/LegacyExchangeApiKey")
-        if not apiKey:
-            return S_ERROR("Missing mandatory /DiracX/LegacyExchangeApiKey configuration")
-
-        diracxUrl = gConfig.getValue("/DiracX/URL")
-        if not diracxUrl:
-            return S_ERROR("Missing mandatory /DiracX/URL configuration")
-
-        credDict = self.getRemoteCredentials()
-        vo = Registry.getVOForGroup(credDict["group"])
-        dirac_properties = list(set(credDict.get("groupProperties", [])) | set(credDict.get("properties", [])))
-        group = credDict["group"]
-        scopes = [f"vo:{vo}", f"group:{group}"] + [f"property:{prop}" for prop in dirac_properties]
-
-        try:
-            r = requests.get(
-                f"{diracxUrl}/api/auth/legacy-exchange",
-                params={
-                    "preferred_username": credDict["username"],
-                    "scope": " ".join(scopes),
-                },
-                headers={"Authorization": f"Bearer {apiKey}"},
-            )
-        except requests.exceptions.RequestException as exc:
-            return S_ERROR(f"Failed to contact DiracX: {exc}")
-        else:
-            if not r.ok:
-                return S_ERROR(f"Failed to contact DiracX: {r.status_code} {r.text}")
-
-        return S_OK(r.json())
+        return get_token(self.getRemoteCredentials())
 
 
 class ProxyManagerHandler(ProxyManagerHandlerMixin, RequestHandler):

diff --git a/src/DIRAC/FrameworkSystem/Utilities/diracx.py b/src/DIRAC/FrameworkSystem/Utilities/diracx.py
@@ -24,13 +24,11 @@
 _token_cache = TTLCache(maxsize=100, ttl=DEFAULT_TOKEN_CACHE_TTL)
 
 
-@cached(_token_cache, key=lambda x, y: repr(x))
-def _get_token(credDict, diracxUrl, /) -> Path:
-    """
-    Write token to a temporary file and return the path to that file
-
-    """
-
+def get_token(credDict, *, expires_minutes=None):
+    """Do a legacy exchange to get a DiracX access_token+refresh_token"""
+    diracxUrl = gConfig.getValue("/DiracX/URL")
+    if not diracxUrl:
+        raise ValueError("Missing mandatory /DiracX/URL configuration")
     apiKey = gConfig.getValue("/DiracX/LegacyExchangeApiKey")
     if not apiKey:
         raise ValueError("Missing mandatory /DiracX/LegacyExchangeApiKey configuration")
@@ -46,17 +44,23 @@ def _get_token(credDict, diracxUrl, /) -> Path:
         params={
             "preferred_username": credDict["username"],
             "scope": " ".join(scopes),
+            "expires_minutes": expires_minutes,
         },
         headers={"Authorization": f"Bearer {apiKey}"},
         timeout=10,
     )
+    if not r.ok:
+        raise RuntimeError(f"Error getting token from DiracX: {r.status_code} {r.text}")
 
-    r.raise_for_status()
-
-    token_location = Path(NamedTemporaryFile().name)
+    return r.json()
 
-    write_credentials(TokenResponse(**r.json()), location=token_location)
 
+@cached(_token_cache, key=lambda x, y: repr(x))
+def _get_token_file(credDict) -> Path:
+    """Write token to a temporary file and return the path to that file"""
+    data = get_token(credDict)
+    token_location = Path(NamedTemporaryFile().name)
+    write_credentials(TokenResponse(**data), location=token_location)
     return token_location
 
 
@@ -69,11 +73,10 @@ def TheImpersonator(credDict: dict[str, Any]) -> DiracClient:
 
     Use as a context manager
     """
-
     diracxUrl = gConfig.getValue("/DiracX/URL")
     if not diracxUrl:
         raise ValueError("Missing mandatory /DiracX/URL configuration")
-    token_location = _get_token(credDict, diracxUrl)
+    token_location = _get_token_file(credDict)
     pref = DiracxPreferences(url=diracxUrl, credentials_path=token_location)
 
     return DiracClient(diracx_preferences=pref)