feat: Add support for DiracX tokens, and possibly tokens in proxies

fstagni · fstagni · commit 88270c0c8829 · 2025-12-08T16:05:23.000+01:00
diff --git a/.github/workflows/basic.yml b/.github/workflows/basic.yml
@@ -3,17 +3,16 @@ name: Basic tests
 on: [push, pull_request]
 
 jobs:
-
   shellcheck:
     runs-on: ubuntu-24.04
     if: github.event_name != 'push' || github.repository == 'DIRACGrid/DIRAC'
     timeout-minutes: 10
 
     steps:
-    - uses: actions/checkout@v3
-    - name: Run shellcheck
-      run: |
-        find tests/CI -name '*.sh' -print0 | xargs -0 -n1 shellcheck --external-sources;
+      - uses: actions/checkout@v3
+      - name: Run shellcheck
+        run: |
+          find tests/CI -name '*.sh' -print0 | xargs -0 -n1 shellcheck --external-sources;
 
   pycodestyle:
     runs-on: ubuntu-24.04
@@ -23,27 +22,26 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python: 
+        python:
           - 3.6.15
           - 3.9.17
 
     container: python:${{ matrix.python }}-slim
     steps:
-    - uses: actions/checkout@v3
-    - name: Installing dependencies
-      run: |
-        python -m pip install pycodestyle
-    - name: Run pycodestyle
-      run: |
-        if [[ "${REFERENCE_BRANCH}" != "" ]]; then
-            git remote add upstream https://github.com/DIRACGrid/Pilot.git
-            git fetch --no-tags upstream "${REFERENCE_BRANCH}"
-            git branch -vv
-            git diff -U0 "upstream/${REFERENCE_BRANCH}" | pycodestyle --diff
-        fi
-      env:
-        REFERENCE_BRANCH: ${{ github['base_ref'] || github['head_ref'] }}
-
+      - uses: actions/checkout@v3
+      - name: Installing dependencies
+        run: |
+          python -m pip install pycodestyle
+      - name: Run pycodestyle
+        run: |
+          if [[ "${REFERENCE_BRANCH}" != "" ]]; then
+              git remote add upstream https://github.com/DIRACGrid/Pilot.git
+              git fetch --no-tags upstream "${REFERENCE_BRANCH}"
+              git branch -vv
+              git diff -U0 "upstream/${REFERENCE_BRANCH}" | pycodestyle --diff
+          fi
+        env:
+          REFERENCE_BRANCH: ${{ github['base_ref'] || github['head_ref'] }}
 
   pytest:
     runs-on: ubuntu-24.04
@@ -53,23 +51,22 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python: 
+        python:
           - 3.6.15
           - 3.9.17
 
     container: python:${{ matrix.python }}-slim
     steps:
-    - uses: actions/checkout@v3
-    - name: Installing dependencies
-      run: |
-        echo 'deb http://archive.debian.org/debian stretch main' > /etc/apt/sources.list
-        echo 'deb http://archive.debian.org/debian-security stretch/updates main' >> /etc/apt/sources.list
-        apt-get update || true
-        python -m pip install pytest mock
-        apt install -y voms-clients
-    - name: Run pytest
-      run: pytest
-
+      - uses: actions/checkout@v3
+      - name: Installing dependencies
+        run: |
+          echo 'deb http://archive.debian.org/debian stretch main' > /etc/apt/sources.list
+          echo 'deb http://archive.debian.org/debian-security stretch/updates main' >> /etc/apt/sources.list
+          apt-get update || true
+          python -m pip install pytest mock
+          apt install -y voms-clients
+      - name: Run pytest
+        run: pytest
 
   pylint:
     runs-on: ubuntu-24.04
@@ -79,15 +76,15 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python: 
+        python:
           - 3.6.15
           - 3.9.17
 
     container: python:${{ matrix.python }}-slim
     steps:
-    - uses: actions/checkout@v3
-    - name: Installing dependencies
-      run: |
-        python -m pip install pylint
-    - name: Run pylint
-      run: pylint -E Pilot/
+      - uses: actions/checkout@v3
+      - name: Installing dependencies
+        run: |
+          python -m pip install pylint
+      - name: Run pylint
+        run: pylint -E Pilot/
diff --git a/Pilot/dirac-pilot.py b/Pilot/dirac-pilot.py
@@ -31,8 +31,7 @@
     getCommand,
     pythonPathCheck,
 )
-
-############################
+from proxyTools import revokePilotToken
 
 if __name__ == "__main__":
     pilotStartTime = int(time.time())
@@ -83,7 +82,9 @@
     log.debug("PARAMETER [%s]" % ", ".join(map(str, pilotParams.optList)))
 
     if pilotParams.commandExtensions:
-        log.info("Requested command extensions: %s" % str(pilotParams.commandExtensions))
+        log.info(
+            "Requested command extensions: %s" % str(pilotParams.commandExtensions)
+        )
 
     log.info("Executing commands: %s" % str(pilotParams.commands))
 
@@ -107,3 +108,15 @@
             if remote:
                 log.buffer.flush()
             sys.exit(-1)
+
+    log.info("Pilot tasks finished.")
+
+    if pilotParams.jwt:
+        if not pilotParams.isLegacyPilot:
+            log.info("Revoking pilot token.")
+            revokePilotToken(
+                pilotParams.diracXServer,
+                pilotParams.pilotUUID,
+                pilotParams.jwt,
+                pilotParams.clientID,
+            )
diff --git a/Pilot/pilotCommands.py b/Pilot/pilotCommands.py
@@ -526,7 +526,20 @@ def __init__(self, pilotParams):
 
     @logFinalizer
     def execute(self):
-        """Calls dirac-admin-add-pilot"""
+        """Calls dirac-admin-add-pilot
+        
+        Deprecated in DIRAC V8, new mechanism in V9 and DiracX."""
+
+        if self.pp.jwt:
+            if not self.pp.isLegacyPilot:
+                self.log.warn("Skipping module, normally it is already done via DiracX secret-exchange.")
+                return
+            
+            # If we're here, this is a legacy pilot with a DiracX token embedded in it.
+            # TODO: See if we do a dirac-admin-add-pilot in DiracX for legacy pilots
+        else:
+            # If we're here, this is a DIRAC only pilot without diracX token embedded in it.
+            pass
 
         if not self.pp.pilotReference:
             self.log.warn("Skipping module, no pilot reference found")
@@ -1210,3 +1223,4 @@ def execute(self):
         """Standard entry point to a pilot command"""
         self._setNagiosOptions()
         self._runNagiosProbes()
+
diff --git a/Pilot/pilotTools.py b/Pilot/pilotTools.py
@@ -22,9 +22,7 @@
 from urllib.parse import urlencode
 from urllib.request import urlopen
 
-from proxyTools import getVO
-
-# Utilities functions
+from proxyTools import BaseRequest, extract_diracx_payload, getVO
 
 
 def load_module_from_path(module_name, path_to_module):
@@ -884,10 +882,14 @@ def __init__(self):
         self.setup = ""
         self.configServer = ""
         self.preferredURLPatterns = ""
+        self.diracXServer = ""
         self.ceName = ""
         self.ceType = ""
         self.queueName = ""
         self.gridCEType = ""
+        self.pilotSecret = ""
+        self.clientID = ""
+        self.jwt = {}
         # maxNumberOfProcessors: the number of
         # processors allocated to the pilot which the pilot can allocate to one payload
         # used to set payloadProcessors unless other limits are reached (like the number of processors on the WN)
@@ -922,6 +924,7 @@ def __init__(self):
         self.pilotCFGFile = "pilot.json"
         self.pilotLogging = False
         self.loggerURL = None
+        self.isLegacyPilot = False
         self.loggerTimerInterval = 0
         self.loggerBufsize = 1000
         self.pilotUUID = "unknown"
@@ -978,6 +981,7 @@ def __init__(self):
             ("y:", "CEType=", "CE Type (normally InProcess)"),
             ("z", "pilotLogging", "Activate pilot logging system"),
             ("C:", "configurationServer=", "Configuration servers to use"),
+            ("", "diracx_URL=", "DiracX Server URL to use"),
             ("D:", "disk=", "Require at least <space> MB available"),
             ("E:", "commandExtensions=", "Python modules with extra commands"),
             ("F:", "pilotCFGFile=", "Specify pilot CFG file"),
@@ -1015,6 +1019,8 @@ def __init__(self):
             ),
             ("", "architectureScript=", "architecture script to use"),
             ("", "CVMFS_locations=", "comma-separated list of CVMS locations"),
+            ("", "pilotSecret=", "secret that the pilot uses with DiracX"),
+            ("", "clientID=", "client id used by DiracX to revoke a token"),
         )
 
         # Possibly get Setup and JSON URL/filename from command line
@@ -1041,6 +1047,74 @@ def __init__(self):
             self.installEnv["X509_USER_PROXY"] = self.certsLocation
             os.environ["X509_USER_PROXY"] = self.certsLocation
 
+        try:
+            self.__get_diracx_jwt()
+        except Exception as e:
+            self.log.error("Error setting DiracX: %s" % e)
+            # Remove all settings to prevent using it.
+            self.diracXServer = None
+            self.pilotSecret = None
+            self.loggerURL = None
+            self.jwt = {}
+            self.log.error("Won't use DiracX.")
+
+    def __get_diracx_jwt(self):
+        # Pilot auth: two cases
+        # 1. Has a secret (DiracX Pilot), exchange for a token
+        # 2. Legacy Pilot, has a proxy with a DiracX section in it (extract the jwt from it)
+        if self.pilotUUID and self.pilotSecret and self.diracXServer:
+            self.log.info("Fetching JWT in DiracX (URL: %s)" % self.diracXServer)
+
+            config = BaseRequest(
+                "%s/api/auth/secret-exchange" % (self.diracXServer),
+                os.getenv("X509_CERT_DIR"),
+                self.pilotUUID,
+            )
+
+            try:
+                self.jwt = config.executeRequest(
+                    {"pilot_stamp": self.pilotUUID, "pilot_secret": self.pilotSecret}
+                )
+            except HTTPError as e:
+                self.log.error("Request failed: %s" % str(e))
+                self.log.error("Could not fetch pilot tokens.")
+                if e.code == 401:
+                    # First test if the error occurred because of "bad pilot_stamp"
+                    # If so, this pilot is in the vacuum case
+                    # So we redo auth, but this time with the right data for vacuum cases
+                    self.log.error("Retrying with vacuum case data...")
+                    self.jwt = config.executeRequest(
+                        {
+                            "pilot_stamp": self.pilotUUID,
+                            "pilot_secret": self.pilotSecret,
+                            "vo": self.wnVO,
+                            "grid_type": self.gridCEType,
+                            "grid_site": self.site,
+                            "status": "Running",
+                        }
+                    )
+                else:
+                    raise RuntimeError("Can't be a vacuum case.")
+
+            self.log.info("Fetched the pilot token with the pilot secret.")
+            self.isLegacyPilot = False
+        elif self.pilotUUID and self.diracXServer:
+            # Try to extract a token for proxy
+            self.log.info("Trying to extract diracx token from proxy.")
+
+            cert = os.getenv("X509_USER_PROXY")
+            if cert:
+                with open(cert, "rb") as fp:
+                    self.jwt = extract_diracx_payload(fp.read())
+                    self.isLegacyPilot = True
+                    self.log.info("Successfully extracted token from proxy.")
+            else:
+                raise RuntimeError("Could not locate a proxy via X509_USER_PROXY")
+        else:
+            self.log.info(
+                "PilotUUID, pilotSecret, and diracXServer are needed to support DiracX."
+            )
+
     def __setSecurityDir(self, envName, dirLocation):
         """Set the environment variable of the `envName`, and add it also to the Pilot Parameters
 
@@ -1152,6 +1226,8 @@ def __initCommandLine2(self):
                 self.keepPythonPath = True
             elif o in ("-C", "--configurationServer"):
                 self.configServer = v
+            elif o == "--diracx_URL":
+                self.diracXServer = v
             elif o in ("-G", "--Group"):
                 self.userGroup = v
             elif o in ("-x", "--execute"):
@@ -1225,6 +1301,10 @@ def __initCommandLine2(self):
                 self.architectureScript = v
             elif o == "--CVMFS_locations":
                 self.CVMFS_locations = v.split(",")
+            elif o == "--pilotSecret":
+                self.pilotSecret = v
+            elif o == "--clientID":
+                self.clientID = v
 
     def __loadJSON(self):
         """
diff --git a/Pilot/proxyTools.py b/Pilot/proxyTools.py
