feat: Include DiracX token in proxy PEM files

setup.cfg

@@ -159,6 +159,7 @@ console_scripts =
     # FrameworkSystem
     dirac-login = DIRAC.FrameworkSystem.scripts.dirac_login:main
     dirac-logout = DIRAC.FrameworkSystem.scripts.dirac_logout:main
+    dirac-diracx-whoami = DIRAC.FrameworkSystem.scripts.dirac_diracx_whoami:main
     dirac-admin-get-CAs = DIRAC.FrameworkSystem.scripts.dirac_admin_get_CAs:main [server]
     dirac-admin-get-proxy = DIRAC.FrameworkSystem.scripts.dirac_admin_get_proxy:main [admin]
     dirac-admin-proxy-upload = DIRAC.FrameworkSystem.scripts.dirac_admin_proxy_upload:main [admin]

src/DIRAC/Core/Security/DiracX.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+__all__ = (
+    "DiracXClient",
+    "diracxTokenFromPEM",
+)
+
+import base64
+import json
+import re
+import textwrap
+from contextlib import contextmanager
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from typing import Any
+
+from diracx.client import DiracClient as _DiracClient
+from diracx.core.models import TokenResponse
+from diracx.core.preferences import DiracxPreferences
+from diracx.core.utils import serialize_credentials
+
+from DIRAC import gConfig
+from DIRAC.ConfigurationSystem.Client.Helpers import Registry
+from DIRAC.Core.Security.Locations import getDefaultProxyLocation
+from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue, returnValueOrRaise
+
+
+PEM_BEGIN = "-----BEGIN DIRACX-----"
+PEM_END = "-----END DIRACX-----"
+RE_DIRACX_PEM = re.compile(rf"{PEM_BEGIN}\n(.*)\n{PEM_END}", re.MULTILINE | re.DOTALL)
+
+
+@convertToReturnValue
+def addProxyToPEM(pemPath, group):
+    from DIRAC.Core.Base.Client import Client
+
+    vo = Registry.getVOMSVOForGroup(group)
+    disabledVOs = gConfig.getValue("/DiracX/DisabledVOs", [])
+    if vo and vo not in disabledVOs:
+        token_content = returnValueOrRaise(
+            Client(url="Framework/ProxyManager", proxyLocation=pemPath).exchangeProxyForToken()
+        )
+
+        diracxUrl = gConfig.getValue("/DiracX/URL")
+        if not diracxUrl:
+            return S_ERROR("Missing mandatory /DiracX/URL configuration")
+
+        token = TokenResponse(
+            access_token=token_content["access_token"],
+            expires_in=token_content["expires_in"],
+            token_type=token_content.get("token_type"),
+            refresh_token=token_content.get("refresh_token"),
+        )
+
+        token_pem = f"{PEM_BEGIN}\n"
+        data = base64.b64encode(serialize_credentials(token).encode("utf-8")).decode()
+        token_pem += textwrap.fill(data, width=64)
+        token_pem += f"\n{PEM_END}\n"
+
+        with open(pemPath, "a") as f:
+            f.write(token_pem)
+
+
+def diracxTokenFromPEM(pemPath) -> dict[str, Any] | None:
+    """Extract the DiracX token from the proxy PEM file"""
+    pem = Path(pemPath).read_text()
+    if match := RE_DIRACX_PEM.search(pem):
+        match = match.group(1)
+        return json.loads(base64.b64decode(match).decode("utf-8"))
+
+
+@contextmanager
+def DiracXClient() -> _DiracClient:
+    """Get a DiracX client instance with the current user's credentials"""
+    diracxUrl = gConfig.getValue("/DiracX/URL")
+    if not diracxUrl:
+        raise ValueError("Missing mandatory /DiracX/URL configuration")
+
+    proxyLocation = getDefaultProxyLocation()
+    diracxToken = diracxTokenFromPEM(proxyLocation)
+
+    with NamedTemporaryFile(mode="wt") as token_file:
+        token_file.write(json.dumps(diracxToken))
+        token_file.flush()
+        token_file.seek(0)
+
+        pref = DiracxPreferences(url=diracxUrl, credentials_path=token_file.name)
+        with _DiracClient(diracx_preferences=pref) as api:
+            yield api

src/DIRAC/Core/Security/ProxyInfo.py

@@ -8,6 +8,7 @@
 from DIRAC.Core.Security.X509Chain import X509Chain  # pylint: disable=import-error
 from DIRAC.Core.Security.VOMS import VOMS
 from DIRAC.Core.Security import Locations
+from DIRAC.Core.Security.DiracX import diracxTokenFromPEM
 
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
 
@@ -25,6 +26,7 @@ def getProxyInfo(proxy=False, disableVOMS=False):
           * 'validDN' : Valid DN in DIRAC
           * 'validGroup' : Valid Group in DIRAC
           * 'secondsLeft' : Seconds left
+          * 'hasDiracxToken'
       * values that can be there
           * 'path' : path to the file,
           * 'group' : DIRAC group
@@ -67,6 +69,11 @@ def getProxyInfo(proxy=False, disableVOMS=False):
             infoDict["VOMS"] = retVal["Value"]
         else:
             infoDict["VOMSError"] = retVal["Message"].strip()
+
+    infoDict["hasDiracxToken"] = False
+    if proxyLocation:
+        infoDict["hasDiracxToken"] = bool(diracxTokenFromPEM(proxyLocation))
+
     return S_OK(infoDict)
 
 
@@ -94,6 +101,7 @@ def formatProxyInfoAsString(infoDict):
         "subproxyUser",
         ("secondsLeft", "timeleft"),
         ("group", "DIRAC group"),
+        ("hasDiracxToken", "DiracX"),
         "rfc",
         "path",
         "username",

src/DIRAC/Core/Tornado/Client/private/TornadoBaseClient.py

@@ -511,10 +511,12 @@ def _request(self, retry=0, outputFile=None, **kwargs):
         # getting certificate
         # Do we use the server certificate ?
         if self.kwargs[self.KW_USE_CERTIFICATES]:
+            # TODO: Does this code path need to work with DiracX?
             auth = {"cert": Locations.getHostCertificateAndKeyLocation()}
 
         # Use access token?
         elif self.__useAccessToken:
+            # TODO: Remove this code path?
             from DIRAC.FrameworkSystem.private.authorization.utils.Tokens import (
                 getLocalTokenDict,
                 writeTokenDictToTokenFile,
@@ -543,13 +545,13 @@ def _request(self, retry=0, outputFile=None, **kwargs):
 
             auth = {"headers": {"Authorization": f"Bearer {token['access_token']}"}}
         elif self.kwargs.get(self.KW_PROXY_STRING):
+            # TODO: This code path cannot work with DiracX
             tmpHandle, cert = tempfile.mkstemp()
             fp = os.fdopen(tmpHandle, "w")
             fp.write(self.kwargs[self.KW_PROXY_STRING])
             fp.close()
-
-        # CHRIS 04.02.21
-        # TODO: add proxyLocation check ?
+        elif self.kwargs.get(self.KW_PROXY_LOCATION):
+            auth = {"cert": self.kwargs[self.KW_PROXY_LOCATION]}
         else:
             auth = {"cert": Locations.getProxyLocation()}
             if not auth["cert"]:

src/DIRAC/FrameworkSystem/Client/ProxyManagerClient.py

@@ -10,6 +10,7 @@
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
 from DIRAC.Core.Utilities import ThreadSafe, DIRACSingleton
 from DIRAC.Core.Utilities.DictCache import DictCache
+from DIRAC.Core.Security.DiracX import addProxyToPEM
 from DIRAC.Core.Security.ProxyFile import multiProxyArgument, deleteMultiProxy
 from DIRAC.Core.Security.X509Chain import X509Chain  # pylint: disable=import-error
 from DIRAC.Core.Security.X509Request import X509Request  # pylint: disable=import-error
@@ -547,6 +548,10 @@ def dumpProxyToFile(self, chain, destinationFile=None, requiredTimeLeft=600):
         if not retVal["OK"]:
             return retVal
         filename = retVal["Value"]
+        if not (result := chain.getDIRACGroup())["OK"]:
+            return result
+        if not (result := addProxyToPEM(filename, result["Value"]))["OK"]:
+            return result
         self.__filesCache.add(cHash, chain.getRemainingSecs()["Value"], filename)
         return S_OK(filename)
 
@@ -655,7 +660,14 @@ def renewProxy(self, proxyToBeRenewed=None, minLifeTime=3600, newProxyLifeTime=4
         chain = retVal["Value"]
 
         if not proxyToRenewDict["tempFile"]:
-            return chain.dumpAllToFile(proxyToRenewDict["file"])
+            filename = proxyToRenewDict["file"]
+            if not (result := chain.dumpAllToFile(filename))["OK"]:
+                return result
+            if not (result := chain.getDIRACGroup())["OK"]:
+                return result
+            if not (result := addProxyToPEM(filename, result["Value"]))["OK"]:
+                return result
+            return S_OK(filename)
 
         return S_OK(chain)
 

src/DIRAC/FrameworkSystem/scripts/dirac_admin_get_proxy.py

@@ -15,6 +15,7 @@
 import DIRAC
 from DIRAC import gLogger, S_OK, S_ERROR
 from DIRAC.Core.Base.Script import Script
+from DIRAC.Core.Security.DiracX import addProxyToPEM
 from DIRAC.FrameworkSystem.Client.ProxyManagerClient import gProxyManager
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
 
@@ -159,6 +160,10 @@ def main():
     if not result["OK"]:
         gLogger.notice(f"Proxy file cannot be written to {params.proxyPath}: {result['Message']}")
         DIRAC.exit(2)
+    if not (result := chain.getDIRACGroup())["OK"]:
+        return result
+    if not (result := addProxyToPEM(params.proxyPath, result["Value"]))["OK"]:
+        return result
     gLogger.notice(f"Proxy downloaded to {params.proxyPath}")
     DIRAC.exit(0)
 

src/DIRAC/FrameworkSystem/scripts/dirac_diracx_whoami.py

@@ -0,0 +1,22 @@
+"""Query DiracX for information about the current user
+
+This is a stripped down version of the "dirac whoami" script from DiracX.
+It primarily exists as a method of validating the current user's credentials are functional.
+"""
+import json
+
+from DIRAC.Core.Base.Script import Script
+from DIRAC.Core.Security.DiracX import DiracXClient
+
+
+@Script()
+def main():
+    Script.parseCommandLine()
+
+    with DiracXClient() as api:
+        user_info = api.auth.userinfo()
+        print(json.dumps(user_info.as_dict(), indent=2))
+
+
+if __name__ == "__main__":
+    main()

src/DIRAC/FrameworkSystem/scripts/dirac_login.py

@@ -25,6 +25,7 @@
 from DIRAC import gConfig, gLogger, S_OK, S_ERROR
 from DIRAC.Core.Security.Locations import getDefaultProxyLocation, getCertificateAndKeyLocation
 from DIRAC.Core.Security.VOMS import VOMS
+from DIRAC.Core.Security.DiracX import addProxyToPEM
 from DIRAC.Core.Security.ProxyFile import writeToProxyFile
 from DIRAC.Core.Security.ProxyInfo import getProxyInfo, formatProxyInfoAsString
 from DIRAC.Core.Security.X509Chain import X509Chain  # pylint: disable=import-error
@@ -314,32 +315,8 @@ def loginWithCertificate(self):
                     return res
 
             # Get a token for use with diracx
-            vo = getVOMSVOForGroup(self.group)
-            disabledVOs = gConfig.getValue("/DiracX/DisabledVOs", [])
-            if vo not in disabledVOs:
-                from diracx.core.utils import write_credentials  # pylint: disable=import-error
-                from diracx.core.models import TokenResponse  # pylint: disable=import-error
-                from diracx.core.preferences import DiracxPreferences  # pylint: disable=import-error
-
-                res = Client(url="Framework/ProxyManager").exchangeProxyForToken()
-                if not res["OK"]:
-                    return res
-                token_content = res["Value"]
-
-                diracxUrl = gConfig.getValue("/DiracX/URL")
-                if not diracxUrl:
-                    return S_ERROR("Missing mandatory /DiracX/URL configuration")
-
-                preferences = DiracxPreferences(url=diracxUrl)
-                write_credentials(
-                    TokenResponse(
-                        access_token=token_content["access_token"],
-                        expires_in=token_content["expires_in"],
-                        token_type=token_content.get("token_type"),
-                        refresh_token=token_content.get("refresh_token"),
-                    ),
-                    location=preferences.credentials_path,
-                )
+            if not (result := addProxyToPEM(self.outputFile, self.group))["OK"]:
+                return result
 
         return S_OK()
 

src/DIRAC/FrameworkSystem/scripts/dirac_proxy_info.py

@@ -77,6 +77,7 @@ def main():
     from DIRAC.Core.Security import VOMS
     from DIRAC.FrameworkSystem.Client.ProxyManagerClient import gProxyManager
     from DIRAC.ConfigurationSystem.Client.Helpers import Registry
+    from DIRAC.Core.Security.DiracX import DiracXClient
 
     if params.csEnabled:
         retVal = Script.enableCS()
@@ -151,6 +152,12 @@ def invalidProxy(msg):
                 invalidProxy(f"Cannot determine life time of VOMS attributes: {result['Message']}")
             if int(result["Value"].strip()) == 0:
                 invalidProxy("VOMS attributes are expired")
+        # Ensure the proxy is working with DiracX
+        try:
+            with DiracXClient() as api:
+                api.auth.userinfo()
+        except Exception as e:
+            invalidProxy(f"Failed to access DiracX: {e}")
 
     sys.exit(0)
 

src/DIRAC/FrameworkSystem/scripts/dirac_proxy_init.py

@@ -18,10 +18,10 @@
 from DIRAC.Core.Base.Script import Script
 from DIRAC.FrameworkSystem.Client import ProxyGeneration, ProxyUpload
 from DIRAC.Core.Security import X509Chain, ProxyInfo, VOMS
-from DIRAC.Core.Security.Locations import getCAsLocation
+from DIRAC.Core.Security.DiracX import addProxyToPEM
+from DIRAC.Core.Security.Locations import getCAsLocation, getDefaultProxyLocation
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
 from DIRAC.FrameworkSystem.Client.BundleDeliveryClient import BundleDeliveryClient
-from DIRAC.Core.Base.Client import Client
 
 
 class Params(ProxyGeneration.CLIParams):
@@ -221,6 +221,11 @@ def doTheMagic(self):
         self.checkCAs()
         pI.certLifeTimeCheck()
         resultProxyWithVOMS = pI.addVOMSExtIfNeeded()
+
+        proxyLoc = self.__piParams.proxyLoc or getDefaultProxyLocation()
+        if not (result := addProxyToPEM(proxyLoc, self.__piParams.diracGroup))["OK"]:
+            return result
+
         if not resultProxyWithVOMS["OK"]:
             if "returning a valid AC for the user" in resultProxyWithVOMS["Message"]:
                 gLogger.error(resultProxyWithVOMS["Message"])
@@ -238,33 +243,6 @@ def doTheMagic(self):
                 if self.__piParams.strict:
                     return resultProxyUpload
 
-        vo = Registry.getVOMSVOForGroup(self.__piParams.diracGroup)
-        disabledVOs = gConfig.getValue("/DiracX/DisabledVOs", [])
-        if vo not in disabledVOs:
-            from diracx.core.utils import write_credentials  # pylint: disable=import-error
-            from diracx.core.models import TokenResponse  # pylint: disable=import-error
-            from diracx.core.preferences import DiracxPreferences  # pylint: disable=import-error
-
-            res = Client(url="Framework/ProxyManager").exchangeProxyForToken()
-            if not res["OK"]:
-                return res
-
-            diracxUrl = gConfig.getValue("/DiracX/URL")
-            if not diracxUrl:
-                return S_ERROR("Missing mandatory /DiracX/URL configuration")
-
-            token_content = res["Value"]
-            preferences = DiracxPreferences(url=diracxUrl)
-            write_credentials(
-                TokenResponse(
-                    access_token=token_content["access_token"],
-                    expires_in=token_content["expires_in"],
-                    token_type=token_content.get("token_type"),
-                    refresh_token=token_content.get("refresh_token"),
-                ),
-                location=preferences.credentials_path,
-            )
-
         return S_OK()
 
 

src/DIRAC/Interfaces/Utilities/DCommands.py

@@ -495,16 +495,21 @@ def proxyInit(self):
         params.diracGroup = retVal["Value"]
 
         result = ProxyGeneration.generateProxy(params)
-
         if not result["OK"]:
             raise Exception(result["Message"])
+        filename = result["Value"]
+
         self.checkCAs()
+
         try:
-            self.addVomsExt(result["Value"])
+            self.addVomsExt(filename)
         except:
             # silently skip VOMS errors
             pass
 
+        if not (result := addProxyToPEM(filename, params.diracGroup))["OK"]:
+            raise Exception(result["Message"])
+
     def addVomsExt(self, proxy):
         retVal = self.getEnv("group_name")
         if not retVal["OK"]: