feat (diracx): exchange a proxy for an equivalent token

Author: chaen

docs/source/AdministratorGuide/ServerInstallations/environment_variable_configuration.rst

@@ -21,6 +21,9 @@ DIRAC_DEPRECATED_FAIL
 DIRAC_ENABLE_DIRACX_JOB_MONITORING
   If set, calls the diracx job monitoring service. Off by default.
 
+DIRAC_ENABLE_DIRACX_LOGIN
+  If set, retrieve a DiracX token when calling dirac-proxy-init or dirac-login
+
 DIRAC_FEWER_CFG_LOCKS
   If ``true`` or ``yes`` or ``on`` or ``1`` or ``y`` or ``t``, DIRAC will reduce the number of locks used when accessing the CS for better performance (default, ``no``).
 

src/DIRAC/FrameworkSystem/Service/ProxyManagerHandler.py

@@ -6,6 +6,7 @@
       :dedent: 2
       :caption: ProxyManager options
 """
+
 from DIRAC import gLogger, S_OK, S_ERROR
 from DIRAC.Core.DISET.RequestHandler import RequestHandler, getServiceOption
 from DIRAC.Core.Security import Properties
@@ -406,6 +407,45 @@ def export_getVOMSProxyWithToken(self, userDN, userGroup, requestPem, requiredLi
         self.__proxyDB.logAction("download voms proxy with token", credDict["DN"], credDict["group"], userDN, userGroup)
         return self.__getVOMSProxy(userDN, userGroup, requestPem, requiredLifetime, vomsAttribute, True)
 
+    types_exchangeProxyForToken = []
+
+    def export_exchangeProxyForToken(self):
+        """Exchange a proxy for an equivalent token to be used with diracx"""
+        try:
+            from diracx.routers.auth import (  # pylint: disable=import-error
+                AuthSettings,
+                create_access_token,
+                TokenResponse,
+            )  # pylint: disable=import-error
+
+            authSettings = AuthSettings()
+
+            from uuid import uuid4
+
+            credDict = self.getRemoteCredentials()
+            vo = Registry.getVOForGroup(credDict["group"])
+            payload = {
+                "sub": f"{vo}:{credDict['username']}",
+                "vo": vo,
+                "aud": authSettings.token_audience,
+                "iss": authSettings.token_issuer,
+                "dirac_properties": list(
+                    set(credDict.get("groupProperties", [])) | set(credDict.get("properties", []))
+                ),
+                "jti": str(uuid4()),
+                "preferred_username": credDict["username"],
+                "dirac_group": credDict["group"],
+            }
+            return S_OK(
+                TokenResponse(
+                    access_token=create_access_token(payload, authSettings),
+                    expires_in=authSettings.access_token_expire_minutes * 60,
+                    state="None",
+                ).dict()
+            )
+        except Exception as e:
+            return S_ERROR(f"Could not get token: {e!r}")
+
 
 class ProxyManagerHandler(ProxyManagerHandlerMixin, RequestHandler):
     pass

src/DIRAC/FrameworkSystem/scripts/dirac_login.py

@@ -16,6 +16,9 @@
 import os
 import sys
 import copy
+import datetime
+import json
+from pathlib import Path
 from prompt_toolkit import prompt, print_formatted_text as print, HTML
 
 import DIRAC
@@ -26,6 +29,13 @@
 from DIRAC.Core.Security.ProxyInfo import getProxyInfo, formatProxyInfoAsString
 from DIRAC.Core.Security.X509Chain import X509Chain  # pylint: disable=import-error
 from DIRAC.Core.Base.Script import Script
+from DIRAC.Core.Base.Client import Client
+
+
+# token location
+DIRAC_TOKEN_FILE = Path.home() / ".cache" / "diracx" / "credentials.json"
+EXPIRES_GRACE_SECONDS = 15
+
 
 # At this point, we disable CS synchronization so that an error related
 # to the lack of a proxy certificate does not occur when trying to synchronize.
@@ -304,7 +314,27 @@ def loginWithCertificate(self):
             # Upload proxy to the server if it longer that uploaded one
             if credentials["secondsLeft"] > uploadedProxyLifetime:
                 gLogger.notice("Upload proxy to server.")
-                return gProxyManager.uploadProxy(proxy)
+                res = gProxyManager.uploadProxy(proxy)
+                if not res["OK"]:
+                    return res
+
+            # Get a token for use with diracx
+            if os.getenv("DIRAC_ENABLE_DIRACX_LOGIN", "No").lower() in ("yes", "true"):
+                res = Client(url="Framework/ProxyManager").exchangeProxyForToken()
+                if not res["OK"]:
+                    return res
+                DIRAC_TOKEN_FILE.parent.mkdir(parents=True, exist_ok=True)
+                expires = datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(
+                    seconds=res["Value"]["expires_in"] - EXPIRES_GRACE_SECONDS
+                )
+                credential_data = {
+                    "access_token": res["Value"]["access_token"],
+                    # TODO: "refresh_token":
+                    # TODO: "refresh_token_expires":
+                    "expires": expires.isoformat(),
+                }
+                DIRAC_TOKEN_FILE.write_text(json.dumps(credential_data))
+
         return S_OK()
 
     def __enableCS(self):

src/DIRAC/FrameworkSystem/scripts/dirac_proxy_init.py

@@ -9,6 +9,7 @@
 import os
 import sys
 import glob
+import json
 import time
 import datetime
 
@@ -21,6 +22,12 @@
 from DIRAC.Core.Security.Locations import getCAsLocation
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
 from DIRAC.FrameworkSystem.Client.BundleDeliveryClient import BundleDeliveryClient
+from DIRAC.Core.Base.Client import Client
+from pathlib import Path
+
+
+DIRAC_TOKEN_FILE = Path.home() / ".cache" / "diracx" / "credentials.json"
+EXPIRES_GRACE_SECONDS = 15
 
 
 class Params(ProxyGeneration.CLIParams):
@@ -237,6 +244,22 @@ def doTheMagic(self):
                 if self.__piParams.strict:
                     return resultProxyUpload
 
+        if os.getenv("DIRAC_ENABLE_DIRACX_LOGIN", "No").lower() in ("yes", "true"):
+            res = Client(url="Framework/ProxyManager").exchangeProxyForToken()
+            if not res["OK"]:
+                return res
+            DIRAC_TOKEN_FILE.parent.mkdir(parents=True, exist_ok=True)
+            expires = datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(
+                seconds=res["Value"]["expires_in"] - EXPIRES_GRACE_SECONDS
+            )
+            credential_data = {
+                "access_token": res["Value"]["access_token"],
+                # TODO: "refresh_token":
+                # TODO: "refresh_token_expires":
+                "expires": expires.isoformat(),
+            }
+            DIRAC_TOKEN_FILE.write_text(json.dumps(credential_data))
+
         return S_OK()