Merge pull request #5899 from TaykYoku/8.0_fixOAuth-2

[8.0][OAuth 2] dirac-login (AS) added VOMS extension

Author: fstagni

src/DIRAC/FrameworkSystem/private/authorization/AuthServer.py

@@ -21,6 +21,7 @@
     wrapIDAsDN,
     getDNForUsername,
     getIdPForGroup,
+    getGroupOption,
 )
 from DIRAC.FrameworkSystem.Client.ProxyManagerClient import ProxyManagerClient
 from DIRAC.FrameworkSystem.Client.TokenManagerClient import TokenManagerClient
@@ -156,10 +157,15 @@ def generateProxyOrToken(
             # Try every DN to generate a proxy
             for dn in userDNs:
                 sLog.debug("Try to get proxy for %s" % dn)
+                params = {}
                 if lifetime:
-                    result = self.proxyCli.downloadProxy(dn, group, requiredTimeLeft=int(lifetime))
+                    params["requiredTimeLeft"] = int(lifetime)
+                # if the configuration describes adding a VOMS extension, we will do so
+                if getGroupOption(group, "AutoAddVOMS", False):
+                    result = self.proxyCli.downloadVOMSProxy(dn, group, **params)
                 else:
-                    result = self.proxyCli.downloadProxy(dn, group)
+                    # otherwise we will return the usual proxy
+                    result = self.proxyCli.downloadProxy(dn, group, **params)
                 if not result["OK"]:
                     err.append(result["Message"])
                 else:

src/DIRAC/FrameworkSystem/scripts/dirac_login.py

@@ -21,6 +21,7 @@
 import DIRAC
 from DIRAC import gConfig, gLogger, S_OK, S_ERROR
 from DIRAC.Core.Security.Locations import getDefaultProxyLocation, getCertificateAndKeyLocation
+from DIRAC.Core.Security.VOMS import VOMS
 from DIRAC.Core.Security.ProxyFile import writeToProxyFile
 from DIRAC.Core.Security.ProxyInfo import getProxyInfo, formatProxyInfoAsString
 from DIRAC.Core.Security.X509Chain import X509Chain  # pylint: disable=import-error
@@ -33,6 +34,12 @@
     readTokenFromFile,
     getTokenFileLocation,
 )
+from DIRAC.ConfigurationSystem.Client.Helpers.Registry import (
+    getGroupOption,
+    getVOMSAttributeForGroup,
+    getVOMSVOForGroup,
+    findDefaultGroupForDN,
+)
 
 # This value shows what authorization way will be by default
 DEFAULT_AUTH_WAY = "certificate"  # possible values are "certificate", "diracas"
@@ -240,8 +247,7 @@ def loginWithCertificate(self):
 
         chain = X509Chain()
         # Load user cert and key
-        result = chain.loadChainFromFile(self.certLoc)
-        if result["OK"]:
+        if (result := chain.loadChainFromFile(self.certLoc))["OK"]:
             result = chain.loadKeyFromFile(
                 self.keyLoc, password=prompt("Enter Certificate password: ", is_password=True)
             )
@@ -259,16 +265,29 @@ def loginWithCertificate(self):
 
         # Create local proxy with group
         self.outputFile = self.outputFile or getDefaultProxyLocation()
-        result = chain.generateProxyToFile(self.outputFile, int(self.lifetime or 12) * 3600, self.group)
+        parameters = (self.outputFile, int(self.lifetime or 12) * 3600, self.group)
+
+        # Add a VOMS extension if the group requires it
+        if (result := chain.generateProxyToFile(*parameters))["OK"] and (result := self.__enableCS())["OK"]:
+            if not self.group and (result := findDefaultGroupForDN(credentials["DN"]))["OK"]:
+                self.group = result["Value"]  # Use default group if user don't set it
+            # based on the configuration we decide whether to add VOMS extensions
+            if getGroupOption(self.group, "AutoAddVOMS", False):
+                if not (vomsAttr := getVOMSAttributeForGroup(self.group)):
+                    print(HTML(f"<yellow>No VOMS attribute foud for {self.group}</yellow>"))
+                else:
+                    vo = getVOMSVOForGroup(self.group)
+                    if not (result := VOMS().setVOMSAttributes(chain, attribute=vomsAttr, vo=vo))["OK"]:
+                        return S_ERROR(f"Failed adding VOMS attribute: {result['Message']}")
+                    chain = result["Value"]
+                    result = chain.generateProxyToFile(*parameters)
         if not result["OK"]:
             return S_ERROR(f"Couldn't generate proxy: {result['Message']}")
 
         if self.enableCS:
             # After creating the proxy, we can try to connect to the server
-            result = Script.enableCS()
-            if not result["OK"]:
-                return S_ERROR(f"Cannot contact CS: {result['Message']}")
-            gConfig.forceRefresh()
+            if not (result := self.__enableCS())["OK"]:
+                return result
 
             # Step 2: Upload proxy to DIRAC server
             result = gProxyManager.getUploadedProxyLifeTime(credentials["subject"])
@@ -282,6 +301,11 @@ def loginWithCertificate(self):
                 return gProxyManager.uploadProxy(proxy)
         return S_OK()
 
+    def __enableCS(self):
+        if not (result := Script.enableCS())["OK"] or not (result := gConfig.forceRefresh())["OK"]:
+            return S_ERROR(f"Cannot contact CS: {result['Message']}")
+        return result
+
     def howToSwitch(self) -> bool:
         """Helper message, how to switch access type(proxy or access token)"""
         if "DIRAC_USE_ACCESS_TOKEN" in self.ENV:
@@ -303,13 +327,10 @@ def howToSwitch(self) -> bool:
 
     def getAuthStatus(self):
         """Try to get user authorization status.
-
         :return: S_OK()/S_ERROR()
         """
-        result = Script.enableCS()
-        if not result["OK"]:
-            return S_ERROR("Cannot contact CS.")
-        gConfig.forceRefresh()
+        if not (result := self.__enableCS())["OK"]:
+            return result
 
         if self.response == "proxy":
             result = getProxyInfo(self.outputFile)