feat: Adding JWT support alongside X509 auth

Author: Robin Van de Merghel

Pilot/pilotTools.py

@@ -18,6 +18,15 @@
 from functools import partial, wraps
 from threading import RLock
 
+try:
+    from Pilot.proxyTools import (
+        X509BasedRequest
+    )
+except ImportError:
+    from proxyTools import (
+        X509BasedRequest
+    )
+
 ############################
 # python 2 -> 3 "hacks"
 try:
@@ -700,26 +709,20 @@ def sendMessage(url, pilotUUID, wnVO, method, rawMessage):
     caPath = os.getenv("X509_CERT_DIR")
     cert = os.getenv("X509_USER_PROXY")
 
-    context = ssl.create_default_context()
-    context.load_verify_locations(capath=caPath)
-
+    
     message = json.dumps((json.dumps(rawMessage), pilotUUID, wnVO))
 
-    try:
-        context.load_cert_chain(cert)  # this is a proxy
-        raw_data = {"method": method, "args": message}
-    except IsADirectoryError:  # assuming it'a dir containing cert and key
-        context.load_cert_chain(os.path.join(cert, "hostcert.pem"), os.path.join(cert, "hostkey.pem"))
-        raw_data = {"method": method, "args": message, "extraCredentials": '"hosts"'}
-
-    if sys.version_info[0] == 3:
-        data = urlencode(raw_data).encode("utf-8")  # encode to bytes ! for python3
-    else:
-        # Python2
-        data = urlencode(raw_data)
-
-    res = urlopen(url, data, context=context)
-    res.close()
+    raw_data = {"method": method, "args": message}
+
+    X509config = X509BasedRequest(url=url,
+                                  caPath=caPath,
+                                  certEnv=cert)
+    
+    # Do the request
+    _res = X509config.executeRequest(raw_data=raw_data,
+                              headers={
+                                "User-Agent": X509config.generateUserAgent(pilotUUID=pilotUUID)
+                              })
 
 
 class CommandBase(object):

Pilot/proxyTools.py

@@ -1,10 +1,27 @@
-"""few functions for dealing with proxies"""
+"""few functions for dealing with proxies and authentication"""
 
 from __future__ import absolute_import, division, print_function
 
 import re
 from base64 import b16decode
 from subprocess import PIPE, Popen
+import ssl
+import sys
+import os
+
+try:
+    IsADirectoryError  # pylint: disable=used-before-assignment
+except NameError:
+    IsADirectoryError = OSError
+
+try:
+    from urllib.parse import urlencode
+    from urllib.request import Request, urlopen
+except ImportError:
+    from urllib import urlencode
+
+    from urllib2 import urlopen
+
 
 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"
@@ -65,3 +82,96 @@ def getVO(proxy_data):
             if match:
                 return match.groups()[0].decode()
     raise NotImplementedError("Something went very wrong")
+
+
+class BaseConnectedRequest(object):
+    """This class helps supporting multiple kinds of requests that requires connections"""
+
+    def __init__(self, url, caPath, name="unknown"):
+        self.name = name
+        self.url = url
+        self.caPath = caPath
+        # We assume we have only one context, so this variable could be shared to avoid opening n times a cert.
+        # On the contrary, to avoid race conditions, we do avoid using "self.data" and "self.headers"
+        self._context = None
+
+        self._prepareRequest()
+
+    def generateUserAgent(self, pilotUUID):
+        """To analyse the traffic, we can send a taylor-made User-Agent
+
+        Args:
+            pilotUUID (str): Unique ID of the Pilot
+
+        Returns:
+            str: The generated user agent
+        """
+        return "Dirac Pilot [%s]" % pilotUUID
+
+    def _prepareRequest(self):
+        """As previously, loads the SSL certificates of the server (to avoid "unknown issuer")"""
+        # Load the SSL context
+        self._context = ssl.create_default_context()
+        self._context.load_verify_locations(capath=self.caPath)
+
+    def executeRequest(self, raw_data, headers={"User-Agent": "Dirac Pilot [Unknown ID]"}):
+        """Execute a HTTP request with the data, headers, and the pre-defined data (SSL + auth)
+
+        Args:
+            raw_data (dict): Data to send
+            headers (dict, optional): Headers to send, helps to track requests. Defaults to {"User-Agent": "Dirac Pilot [Unknown ID]"}.
+
+        Returns:
+            str: Response of the HTTP request
+        """
+        if sys.version_info[0] == 3:
+            data = urlencode(raw_data).encode("utf-8")  # encode to bytes ! for python3
+        else:
+            # Python2
+            data = urlencode(raw_data)
+
+        request = Request(self.url, data=data, headers=headers)
+
+        res = urlopen(request, context=self._context)
+        res.close()
+
+        return res.read()
+
+
+class TokenBasedRequest(BaseConnectedRequest):
+    """Connected Request with JWT support"""
+
+    def __init__(self, url, caPath, jwtData):
+        super(TokenBasedRequest, self).__init__(url, caPath, "TokenBasedConnection")
+
+        self.jwtData = jwtData
+
+    def executeRequest(self, raw_data, headers={"User-Agent": "Dirac Pilot [Unknown ID]"}):
+        # Adds the JWT in the HTTP request (in the Bearer field)
+        headers["Bearer"] = self.jwtData
+        return super(TokenBasedRequest, self).executeRequest(raw_data, headers)
+
+
+class X509BasedRequest(BaseConnectedRequest):
+    """Connected Request with X509 support"""
+
+    def __init__(self, url, caPath, certEnv):
+        super(X509BasedRequest, self).__init__(url, caPath, "X509BasedConnection")
+
+        self.certEnv = certEnv
+        self._hasExtraCredentials = False
+
+        # Load X509 once
+        try:
+            self._context.load_cert_chain(self.certEnv)
+        except IsADirectoryError:  # assuming it'a dir containing cert and key
+            self._context.load_cert_chain(
+                os.path.join(self.certEnv, "hostcert.pem"), os.path.join(self.certEnv, "hostkey.pem")
+            )
+            self._hasExtraCredentials = True
+
+    def executeRequest(self, raw_data, headers={"User-Agent": "Dirac Pilot [Unknown ID]"}):
+        # Adds a flag if the passed cert is a Directory
+        if self._hasExtraCredentials:
+            raw_data["extraCredentials"] = '"hosts"'
+        return super(X509BasedRequest, self).executeRequest(raw_data, headers)