Merge pull request #8250 from ic-hep/cherry-pick-2-6c2c0c52b-integration

[sweep:integration] Token expiry in RSS

Author: fstagni

src/DIRAC/Interfaces/API/DiracAdmin.py

@@ -4,7 +4,9 @@
 site banning and unbanning, WMS proxy uploading etc.
 
 """
+
 import os
+from datetime import datetime, timedelta
 
 from DIRAC import S_ERROR, S_OK, gConfig, gLogger
 from DIRAC.ConfigurationSystem.Client.CSAPI import CSAPI
@@ -150,7 +152,7 @@ def getSiteSection(self, site, printOutput=False):
         return result
 
     #############################################################################
-    def allowSite(self, site, comment, printOutput=False):
+    def allowSite(self, site, comment, printOutput=False, days=1):
         """Adds the site to the site mask. The site must be a valid DIRAC site name
 
         Example usage:
@@ -174,7 +176,13 @@ def allowSite(self, site, comment, printOutput=False):
                 gLogger.notice(f"Site {site} is already Active")
             return S_OK(f"Site {site} is already Active")
 
-        if not (result := self.sitestatus.setSiteStatus(site, "Active", comment))["OK"]:
+        tokenLifetime = int(days)
+        if tokenLifetime <= 0:
+            tokenExpiration = datetime.max
+        else:
+            tokenExpiration = datetime.utcnow().replace(microsecond=0) + timedelta(days=tokenLifetime)
+
+        if not (result := self.sitestatus.setSiteStatus(site, "Active", comment, expiry=tokenExpiration))["OK"]:
             return result
 
         if printOutput:
@@ -223,7 +231,7 @@ def getSiteMaskLogging(self, site=None, printOutput=False):
         return S_OK()
 
     #############################################################################
-    def banSite(self, site, comment, printOutput=False):
+    def banSite(self, site, comment, printOutput=False, days=1):
         """Removes the site from the site mask.
 
         Example usage:
@@ -236,7 +244,6 @@ def banSite(self, site, comment, printOutput=False):
         """
         if not (result := self._checkSiteIsValid(site))["OK"]:
             return result
-
         mask = self.getSiteMask(status="Banned")
         if not mask["OK"]:
             return mask
@@ -246,7 +253,12 @@ def banSite(self, site, comment, printOutput=False):
                 gLogger.notice(f"Site {site} is already Banned")
             return S_OK(f"Site {site} is already Banned")
 
-        if not (result := self.sitestatus.setSiteStatus(site, "Banned", comment))["OK"]:
+        tokenLifetime = int(days)
+        if tokenLifetime <= 0:
+            tokenExpiration = datetime.max
+        else:
+            tokenExpiration = datetime.utcnow().replace(microsecond=0) + timedelta(days=tokenLifetime)
+        if not (result := self.sitestatus.setSiteStatus(site, "Banned", comment, expiry=tokenExpiration))["OK"]:
             return result
 
         if printOutput:

src/DIRAC/Interfaces/scripts/dirac_admin_allow_site.py

@@ -13,6 +13,9 @@
 @Script()
 def main():
     Script.registerSwitch("E:", "email=", "Boolean True/False (True by default)")
+    Script.registerSwitch(
+        "", "days=", "Number of days the token is valid for. Default is 1 day. 0 or less days denotes forever."
+    )
     # Registering arguments will automatically add their description to the help menu
     Script.registerArgument("Site:     Name of the Site")
     Script.registerArgument("Comment:  Reason of the action")
@@ -32,17 +35,20 @@ def getBoolean(value):
             Script.showHelp()
 
     email = True
+    days = 1
     for switch in Script.getUnprocessedSwitches():
         if switch[0] == "email":
             email = getBoolean(switch[1])
+        if switch[0] == "days":
+            days = int(switch[1])
 
     diracAdmin = DiracAdmin()
     exitCode = 0
     errorList = []
 
     # parseCommandLine show help when mandatory arguments are not specified or incorrect argument
     site, comment = Script.getPositionalArgs(group=True)
-    result = diracAdmin.allowSite(site, comment, printOutput=True)
+    result = diracAdmin.allowSite(site, comment, printOutput=True, days=days)
     if not result["OK"]:
         errorList.append((site, result["Message"]))
         exitCode = 2

src/DIRAC/Interfaces/scripts/dirac_admin_ban_site.py

@@ -13,6 +13,9 @@
 @Script()
 def main():
     Script.registerSwitch("E:", "email=", "Boolean True/False (True by default)")
+    Script.registerSwitch(
+        "", "days=", "Number of days the token is valid for. Default is 1 day. 0 or less days denotes forever."
+    )
     # Registering arguments will automatically add their description to the help menu
     Script.registerArgument("Site:     Name of the Site")
     Script.registerArgument("Comment:  Reason of the action")
@@ -32,9 +35,12 @@ def getBoolean(value):
             Script.showHelp()
 
     email = True
+    days = 1
     for switch in Script.getUnprocessedSwitches():
         if switch[0] == "email":
             email = getBoolean(switch[1])
+        if switch[0] == "days":
+            days = int(switch[1])
 
     diracAdmin = DiracAdmin()
     exitCode = 0
@@ -50,7 +56,7 @@ def getBoolean(value):
 
     # parseCommandLine show help when mandatory arguments are not specified or incorrect argument
     site, comment = Script.getPositionalArgs(group=True)
-    result = diracAdmin.banSite(site, comment, printOutput=True)
+    result = diracAdmin.banSite(site, comment, printOutput=True, days=days)
     if not result["OK"]:
         errorList.append((site, result["Message"]))
         exitCode = 2

src/DIRAC/ResourceStatusSystem/Client/SiteStatus.py

@@ -1,4 +1,4 @@
-""" SiteStatus helper
+"""SiteStatus helper
 
 Module that acts as a helper for knowing the status of a site.
 It takes care of switching between the CS and the RSS.
@@ -195,7 +195,7 @@ def getSites(self, siteState="Active"):
 
         return S_OK(siteList)
 
-    def setSiteStatus(self, site, status, comment="No comment"):
+    def setSiteStatus(self, site, status, comment="No comment", expiry=None):
         """
         Set the status of a site in the 'SiteStatus' table of RSS
 
@@ -231,6 +231,8 @@ def setSiteStatus(self, site, status, comment="No comment"):
             return S_ERROR(f"Unable to get user proxy info {result['Message']} ")
 
         tokenExpiration = datetime.utcnow() + timedelta(days=1)
+        if expiry:
+            tokenExpiration = expiry
 
         self.rssCache.acquireLock()
         try:

src/DIRAC/ResourceStatusSystem/scripts/dirac_rss_set_status.py

@@ -29,6 +29,7 @@ def registerSwitches():
         ("reason=", "Reason to set the Status"),
         ("VO=", "VO to change a status for. When omitted, status will be changed for all VOs"),
         ("tokenOwner=", "Owner of the token"),
+        ("days=", "Number of days the token is valid for. Default is 1 day. 0 or less days denotes forever."),
     )
 
     for switch in switches:
@@ -50,6 +51,7 @@ def parseSwitches():
     switches = dict(Script.getUnprocessedSwitches())
     switches.setdefault("statusType", None)
     switches.setdefault("VO", None)
+    switches.setdefault("days", 1)
 
     for key in ("element", "name", "status", "reason"):
         if key not in switches:
@@ -183,7 +185,11 @@ def setStatus(switchDict, tokenOwner):
         )
         return S_OK()
 
-    tomorrow = datetime.utcnow().replace(microsecond=0) + timedelta(days=1)
+    tokenLifetime = int(switchDict["days"])
+    if tokenLifetime <= 0:
+        tokenExpiration = datetime.max
+    else:
+        tokenExpiration = datetime.utcnow().replace(microsecond=0) + timedelta(days=tokenLifetime)
 
     for status, statusType in elements:
         gLogger.debug(f"{status} {statusType}")
@@ -193,8 +199,16 @@ def setStatus(switchDict, tokenOwner):
             continue
 
         gLogger.debug(
-            "About to set status %s -> %s for %s, statusType: %s, VO: %s, reason: %s"
-            % (status, switchDict["status"], switchDict["name"], statusType, switchDict["VO"], switchDict["reason"])
+            "About to set status %s -> %s for %s, statusType: %s, VO: %s, reason: %s, days: %s"
+            % (
+                status,
+                switchDict["status"],
+                switchDict["name"],
+                statusType,
+                switchDict["VO"],
+                switchDict["reason"],
+                switchDict["days"],
+            )
         )
         result = rssClient.modifyStatusElement(
             switchDict["element"],
@@ -205,7 +219,7 @@ def setStatus(switchDict, tokenOwner):
             reason=switchDict["reason"],
             vO=switchDict["VO"],
             tokenOwner=tokenOwner,
-            tokenExpiration=tomorrow,
+            tokenExpiration=tokenExpiration,
         )
         if not result["OK"]:
             return result