fix: Fixed to match the current PR endpoint

Author: Robin-Van-de-Merghel

Pilot/pilotCommands.py

@@ -593,13 +593,19 @@ def __init__(self, pilotParams):
     def execute(self):
         """Calls diracX api"""
 
-        if not self.pp.pilotReference:
-            self.log.warn("Skipping module, no pilot reference found")
-            return
+        if not self.pp.pilotUUID:
+            self.log.error("PilotUUID not given, exiting...")
+            sys.exit(-1)
 
         if not self.pp.pilotSecret:
-            self.log.warn("Skipping module, no pilot secret found")
-            return
+            self.log.error("PilotSecret not given, exiting...")
+            sys.exit(-1)
+
+        if not self.pp.diracXServer:
+            self.log.error("DiracXServer (url) not given, exiting...")
+            sys.exit(-1)
+
+        self.log.info("Fetching JWT in DiracX (URL: %s)" % self.pp.diracXServer)
 
         config = BaseRequest(
             "%s/api/auth/pilot-login" % (
@@ -610,11 +616,16 @@ def execute(self):
 
         config.generateUserAgent(self.pp.pilotUUID)
 
-        self.pp.jwt = config.executeRequest({
-            "pilot_stamp": self.pp.pilotStamp,
-            "pilot_secret": self.pp.pilotSecret
-        }, insecure=True)
-        
+        try:
+            self.pp.jwt = config.executeRequest({
+                "pilot_stamp": self.pp.pilotUUID,
+                "pilot_secret": self.pp.pilotSecret
+            }, insecure=True)
+        except HTTPError as e:
+            self.log.error("Request failed: %s" % str(e))
+
+        self.log.info("Fetched the pilot token with the pilot secret.")
+
 class CheckCECapabilities(CommandBase):
     """Used to get CE tags and other relevant parameters."""
 

Pilot/pilotTools.py

@@ -729,26 +729,11 @@ def sendMessage(url, pilotUUID, wnVO, method, rawMessage, jwt={}):
     config.generateUserAgent(pilotUUID=pilotUUID)
 
     # Do the request
-    _res = config.executeRequest(raw_data=raw_data)
-
-
-
-def retrieveJWT(diracXURL, pilotUUID, pilotSecret):
-    
-    caPath = os.getenv("X509_CERT_DIR")
-    
-    data_in_url = "?pilot_job_reference=%s&pilot_secret=%s" % (pilotUUID, pilotSecret)
-    
-    config = BaseRequest(
-        "%s/api/auth/pilot-login%s" % (diracXURL, data_in_url),
-        caPath=caPath
+    _res = config.executeRequest(
+        raw_data=raw_data,
+        content_type="x-www-form-urlencoded"
     )
-    
-    config.addHeader("Content-Type", "application/json")
-    
-    config.generateUserAgent(pilotUUID=pilotUUID)
 
-    return config.executeRequest(raw_data={}, insecure=True)
 
 class CommandBase(object):
     """CommandBase is the base class for every command in the pilot commands toolbox"""

Pilot/proxyTools.py

@@ -17,12 +17,12 @@
 
 try:
     from urllib.parse import urlencode
+    from urllib.error import HTTPError
     from urllib.request import Request, urlopen
 except ImportError:
     from urllib import urlencode
 
-    from urllib2 import Request, urlopen
-
+    from urllib2 import Request, urlopen, HTTPError
 
 VOMS_FQANS_OID = b"1.3.6.1.4.1.8005.100.100.4"
 VOMS_EXTENSION_OID = b"1.3.6.1.4.1.8005.100.100.5"
@@ -81,7 +81,7 @@ def getVO(proxy_data):
 
 
 class BaseRequest(object):
-    """This class helps supporting multiple kinds of requests that requires connections"""
+    """This class helps supporting multiple kinds of requests that require connections"""
 
     def __init__(self, url, caPath, name="unknown"):
         self.name = name
@@ -114,21 +114,32 @@ def addHeader(self, key, value):
         """Add a header (key, value) into the request header"""
         self.headers[key] = value
 
-    def executeRequest(self, raw_data, insecure=False):
+    def executeRequest(self, raw_data, insecure=False, content_type="json"):
         """Execute a HTTP request with the data, headers, and the pre-defined data (SSL + auth)
 
         :param raw_data: Data to send
         :type raw_data: dict
         :param insecure: Deactivate proxy verification WARNING Debug ONLY
         :type insecure: bool
+        :param content_type: Data format to send, either "json" or "x-www-form-urlencoded"
+        :type content_type: str
         :return: Parsed JSON response
         :rtype: dict
-        """    
-        if sys.version_info.major == 3:
-            data = urlencode(raw_data).encode("utf-8")  # encode to bytes ! for python3
+        """
+        if content_type == "json":
+            data = json.dumps(raw_data).encode("utf-8")
+            self.headers["Content-Type"] = "application/json"
+        elif content_type == "x-www-form-urlencoded":
+            if sys.version_info.major == 3:
+                data = urlencode(raw_data).encode("utf-8")  # encode to bytes ! for python3
+            else:
+                # Python2
+                data = urlencode(raw_data)
+            self.headers["Content-Type"] = "application/x-www-form-urlencoded"
         else:
-            # Python2
-            data = urlencode(raw_data)
+            raise ValueError("Invalid content_type. Use 'json' or 'x-www-form-urlencoded'.")
+
+        self.headers["Content-Length"] = str(len(data))
 
         request = Request(self.url, data=data, headers=self.headers, method="POST")
 
@@ -190,8 +201,12 @@ def __init__(self, url, caPath, certEnv):
             )
             self._hasExtraCredentials = True
 
-    def executeRequest(self, raw_data):
+    def executeRequest(self, raw_data, insecure=False, content_type="json"):
         # Adds a flag if the passed cert is a Directory
         if self._hasExtraCredentials:
             raw_data["extraCredentials"] = '"hosts"'
-        return super(X509BasedRequest, self).executeRequest(raw_data)
+        return super(X509BasedRequest, self).executeRequest(
+            raw_data,
+            insecure=insecure,
+            content_type=content_type
+        )