Merge pull request DIRACGrid#7243 from chrisburr/diracx

Add /DiracX/EnabledVOs to support partially enabling DiracX

docs/source/AdministratorGuide/ServerInstallations/environment_variable_configuration.rst

@@ -21,9 +21,6 @@ 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``).
 

integration_tests.py

@@ -35,10 +35,7 @@
     "DIRAC_USE_JSON_ENCODE": None,
     "INSTALLATION_BRANCH": "",
 }
-DIRACX_OPTIONS = (
-    "DIRAC_ENABLE_DIRACX_LOGIN",
-    "DIRAC_ENABLE_DIRACX_JOB_MONITORING",
-)
+DIRACX_OPTIONS = ("DIRAC_ENABLE_DIRACX_JOB_MONITORING",)
 DEFAULT_MODULES = {"DIRAC": Path(__file__).parent.absolute()}
 
 # Static configuration

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):

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)

src/DIRAC/FrameworkSystem/scripts/dirac_login.py

@@ -314,7 +314,9 @@ def loginWithCertificate(self):
                     return res
 
             # Get a token for use with diracx
-            if os.getenv("DIRAC_ENABLE_DIRACX_LOGIN", "No").lower() in ("yes", "true"):
+            vo = getVOMSVOForGroup(self.group)
+            enabledVOs = gConfig.getValue("/DiracX/EnabledVOs", [])
+            if vo in enabledVOs:
                 from diracx.core.utils import write_credentials  # pylint: disable=import-error
                 from diracx.core.models import TokenResponse  # pylint: disable=import-error
 

src/DIRAC/FrameworkSystem/scripts/dirac_proxy_init.py

@@ -9,21 +9,19 @@
 import os
 import sys
 import glob
-import json
 import time
 import datetime
 
 import DIRAC
 
-from DIRAC import gLogger, S_OK, S_ERROR
+from DIRAC import gLogger, gConfig, S_OK, S_ERROR
 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.ConfigurationSystem.Client.Helpers import Registry
 from DIRAC.FrameworkSystem.Client.BundleDeliveryClient import BundleDeliveryClient
 from DIRAC.Core.Base.Client import Client
-from pathlib import Path
 
 
 class Params(ProxyGeneration.CLIParams):
@@ -240,15 +238,16 @@ def doTheMagic(self):
                 if self.__piParams.strict:
                     return resultProxyUpload
 
-        if os.getenv("DIRAC_ENABLE_DIRACX_LOGIN", "No").lower() in ("yes", "true"):
+        vo = Registry.getVOMSVOForGroup(self.__piParams.diracGroup)
+        enabledVOs = gConfig.getValue("/DiracX/EnabledVOs", [])
+        if vo in enabledVOs:
             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
-            from DIRAC import gConfig
 
             diracxUrl = gConfig.getValue("/DiracX/URL")
             if not diracxUrl:

src/DIRAC/WorkloadManagementSystem/Service/SandboxStoreHandler.py

@@ -14,7 +14,8 @@
 import threading
 import time
 
-from DIRAC import S_ERROR, S_OK, gLogger
+from DIRAC import S_ERROR, S_OK, gLogger, gConfig
+from DIRAC.ConfigurationSystem.Client.Helpers import Registry
 from DIRAC.Core.DISET.RequestHandler import RequestHandler
 from DIRAC.Core.Security import Locations, Properties, X509Certificate
 from DIRAC.Core.Utilities.File import mkDir
@@ -108,8 +109,10 @@ def _getFromClient(self, fileId, token, fileSize, fileHelper=None, data=""):
         gLogger.info("Upload requested", f"for {aHash} [{extension}]")
 
         credDict = self.getRemoteCredentials()
+        vo = Registry.getVOForGroup(credDict["group"])
 
-        if self._useDiracXBackend:
+        enabledVOs = gConfig.getValue("/DiracX/EnabledVOs", [])
+        if self._useDiracXBackend and vo in enabledVOs:
             from DIRAC.FrameworkSystem.Utilities.diracx import TheImpersonator
             from diracx.client.models import SandboxInfo  # pylint: disable=import-error