fix: perform pilot operations with proxies or tokens

Author: atsareg

src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py

@@ -14,7 +14,7 @@
 from DIRAC.WorkloadManagementSystem.Client import PilotStatus
 from DIRAC.WorkloadManagementSystem.Service.WMSUtilities import (
     getPilotCE,
-    getPilotProxy,
+    setPilotCredentials,
     getPilotRef,
     killPilotsInQueues,
 )
@@ -135,13 +135,10 @@ def export_getPilotOutput(self, pilotReference):
         if not hasattr(ce, "getJobOutput"):
             return S_ERROR(f"Pilot output not available for {pilotDict['GridType']} CEs")
 
-        result = getPilotProxy(pilotDict)
+        result = setPilotCredentials(ce, pilotDict)
         if not result["OK"]:
             return result
 
-        proxy = result["Value"]
-        ce.setProxy(proxy)
-
         result = getPilotRef(pilotReference, pilotDict)
         if not result["OK"]:
             return result
@@ -213,13 +210,11 @@ def export_getPilotLoggingInfo(self, pilotReference):
             self.log.info("Pilot logging not available for", f"{pilotDict['GridType']} CEs")
             return S_ERROR(f"Pilot logging not available for {pilotDict['GridType']} CEs")
 
-        result = getPilotProxy(pilotDict)
+        # Set proxy or token for the CE
+        result = setPilotCredentials(ce, pilotDict)
         if not result["OK"]:
             return result
 
-        proxy = result["Value"]
-        ce.setProxy(proxy)
-
         result = getPilotRef(pilotReference, pilotDict)
         if not result["OK"]:
             return result

src/DIRAC/WorkloadManagementSystem/Service/WMSUtilities.py

@@ -6,8 +6,9 @@
 
 from DIRAC import S_OK, S_ERROR, gLogger, gConfig
 from DIRAC.ConfigurationSystem.Client.Helpers.Resources import getQueue
-from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getGroupOption
+from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getGroupOption, getUsernameForDN
 from DIRAC.FrameworkSystem.Client.ProxyManagerClient import gProxyManager
+from DIRAC.FrameworkSystem.Client.TokenManagerClient import gTokenManager
 from DIRAC.Resources.Computing.ComputingElementFactory import ComputingElementFactory
 
 
@@ -49,7 +50,11 @@ def getPilotCE(pilotDict):
 
 
 def getPilotProxy(pilotDict):
-    """Get a proxy bound to a pilot"""
+    """Get a proxy bound to a pilot
+
+    :param dict pilotDict: pilot parameters
+    :return: S_OK/S_ERROR with proxy as Value
+    """
     owner = pilotDict["OwnerDN"]
     group = pilotDict["OwnerGroup"]
 
@@ -62,6 +67,47 @@ def getPilotProxy(pilotDict):
     return S_OK(proxy)
 
 
+def getPilotToken(pilotDict):
+    """Get a token corresponding to the pilot
+
+    :param dict pilotDict: pilot parameters
+    :return: S_OK/S_ERROR with token as Value
+    """
+    ownerDN = pilotDict["OwnerDN"]
+    group = pilotDict["OwnerGroup"]
+
+    result = getUsernameForDN(ownerDN)
+    if not result["OK"]:
+        return result
+    username = result["Value"]
+    result = gTokenManager.getToken(
+        username=username,
+        userGroup=group,
+        requiredTimeLeft=3600,
+    )
+    return result
+
+
+def setPilotCredentials(ce, pilotDict):
+    """Instrument the given CE with proxy or token
+
+    :param obj ce:  CE object
+    :param pilotDict: pilot parameter dictionary
+    :return: S_OK/S_ERROR
+    """
+    if "Token" in ce.ceParameters.get("Tag", []):
+        result = getPilotToken(pilotDict)
+        if not result["OK"]:
+            return result
+        ce.setToken(result["Value"], 3500)
+    else:
+        result = getPilotProxy(pilotDict)
+        if not result["OK"]:
+            return result
+        ce.setProxy(result["Value"])
+    return S_OK()
+
+
 def getPilotRef(pilotReference, pilotDict):
     """Add the pilotStamp to the pilotReference, if the pilotStamp is in the dictionary,
     otherwise return unchanged pilotReference.