Merge pull request #7261 from chrisburr/token-in-pem

[9.0] Include DiracX token in proxy PEM files

Author: fstagni

integration_tests.py

@@ -146,12 +146,13 @@ def create(
     flags: Optional[list[str]] = typer.Argument(None),
     editable: Optional[bool] = None,
     extra_module: Optional[list[str]] = None,
+    diracx_dist_dir: Optional[str] = None,
     release_var: Optional[str] = None,
     run_server_tests: bool = True,
     run_client_tests: bool = True,
 ):
     """Start a local instance of the integration tests"""
-    prepare_environment(flags, editable, extra_module, release_var)
+    prepare_environment(flags, editable, extra_module, diracx_dist_dir, release_var)
     install_server()
     install_client()
     exit_code = 0
@@ -191,6 +192,7 @@ def prepare_environment(
     flags: Optional[list[str]] = typer.Argument(None),
     editable: Optional[bool] = None,
     extra_module: Optional[list[str]] = None,
+    diracx_dist_dir: Optional[str] = None,
     release_var: Optional[str] = None,
 ):
     """Prepare the local environment for installing DIRAC."""
@@ -227,7 +229,7 @@ def prepare_environment(
     extra_services = list(chain(*[config["extra-services"] for config in module_configs.values()]))
 
     typer.secho("Running docker-compose to create containers", fg=c.GREEN)
-    with _gen_docker_compose(modules) as docker_compose_fn:
+    with _gen_docker_compose(modules, diracx_dist_dir=diracx_dist_dir) as docker_compose_fn:
         subprocess.run(
             ["docker-compose", "-f", docker_compose_fn, "up", "-d", "dirac-server", "dirac-client"] + extra_services,
             check=True,
@@ -322,7 +324,7 @@ def prepare_environment(
     typer.secho("Running docker-compose to create DiracX containers", fg=c.GREEN)
     typer.secho(f"Will leave a folder behind: {docker_compose_fn_final}", fg=c.YELLOW)
 
-    with _gen_docker_compose(modules) as docker_compose_fn:
+    with _gen_docker_compose(modules, diracx_dist_dir=diracx_dist_dir) as docker_compose_fn:
         # We cannot use the temporary directory created in the context manager because
         # we don't stay in the contect manager (Popen)
         # So we need something that outlives it.
@@ -545,7 +547,7 @@ class TestExit(typer.Exit):
 
 
 @contextmanager
-def _gen_docker_compose(modules):
+def _gen_docker_compose(modules, *, diracx_dist_dir=None):
     # Load the docker-compose configuration and mount the necessary volumes
     input_fn = Path(__file__).parent / "tests/CI/docker-compose.yml"
     docker_compose = yaml.safe_load(input_fn.read_text())
@@ -560,10 +562,12 @@ def _gen_docker_compose(modules):
     docker_compose["services"]["diracx-wait-for-db"]["volumes"].extend(volumes[:])
 
     module_configs = _load_module_configs(modules)
-    if "diracx" in module_configs:
-        docker_compose["services"]["diracx"]["volumes"].append(
-            f"{modules['diracx']}/src/diracx:{module_configs['diracx']['install-location']}"
-        )
+    if diracx_dist_dir is not None:
+        for container_name in ["dirac-client", "dirac-server", "diracx-init-cs", "diracx-wait-for-db", "diracx"]:
+            docker_compose["services"][container_name]["volumes"].append(f"{diracx_dist_dir}:/diracx_sources")
+            docker_compose["services"][container_name].setdefault("environment", []).append(
+                "DIRACX_CUSTOM_SOURCE_PREFIXES=/diracx_sources"
+            )
 
     # Add any extension services
     for module_name, module_configs in module_configs.items():

setup.cfg

@@ -29,6 +29,8 @@ install_requires =
     cachetools
     certifi
     diraccfg
+    diracx-client
+    diracx-core
     db12
     fts3
     gfal2-python
@@ -160,6 +162,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,85 @@
+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, S_ERROR
+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 addTokenToPEM(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()
+        )
+
+        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: make this code path work with DiracX for Agents and possibly webapp ?
             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 addTokenToPEM
 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 := addTokenToPEM(filename, result["Value"]))["OK"]:  # pylint: disable=unsubscriptable-object
+            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 := addTokenToPEM(filename, result["Value"]))["OK"]:  # pylint: disable=unsubscriptable-object
+                return result
+            return S_OK(filename)
 
         return S_OK(chain)
 

src/DIRAC/FrameworkSystem/Service/ProxyManagerHandler.py

@@ -12,7 +12,7 @@
 from DIRAC.Core.Security import Properties
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
-
+from DIRAC.FrameworkSystem.Utilities.diracx import get_token
 
 DEFAULT_MAIL_FROM = "[email protected]"
 
@@ -412,7 +412,6 @@ def export_getVOMSProxyWithToken(self, userDN, userGroup, requestPem, requiredLi
     @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
 
         credDict = self.getRemoteCredentials()
         return get_token(

src/DIRAC/FrameworkSystem/Utilities/diracx.py

@@ -1,4 +1,3 @@
-# pylint: disable=import-error
 import requests
 
 from cachetools import TTLCache, cached

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 addTokenToPEM
 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 := addTokenToPEM(params.proxyPath, result["Value"]))["OK"]:  # pylint: disable=unsubscriptable-object
+        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()