Merge pull request #8213 from ic-hep/rss_time

[8.0] 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 gLogger, gConfig, S_OK, S_ERROR
 from DIRAC.Core.Utilities.PromptUser import promptUser
@@ -175,7 +177,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.
 
         Example usage:
@@ -200,7 +202,12 @@ def allowSite(self, site, comment, printOutput=False):
             return S_OK(f"Site {site} is already Active")
 
         if self.rssFlag:
-            result = self.sitestatus.setSiteStatus(site, "Active", comment)
+            tokenLifetime = int(days)
+            if tokenLifetime <= 0:
+                tokenExpiration = datetime.max
+            else:
+                tokenExpiration = datetime.utcnow().replace(microsecond=0) + timedelta(days=tokenLifetime)
+            result = self.sitestatus.setSiteStatus(site, "Active", comment, expiry=tokenExpiration)
         else:
             result = WMSAdministratorClient().allowSite(site, comment)
         if not result["OK"]:
@@ -258,7 +265,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:
@@ -272,7 +279,6 @@ def banSite(self, site, comment, printOutput=False):
         result = self._checkSiteIsValid(site)
         if not result["OK"]:
             return result
-
         mask = self.getSiteMask(status="Banned")
         if not mask["OK"]:
             return mask
@@ -281,9 +287,13 @@ def banSite(self, site, comment, printOutput=False):
             if printOutput:
                 gLogger.notice(f"Site {site} is already Banned")
             return S_OK(f"Site {site} is already Banned")
-
         if self.rssFlag:
-            result = self.sitestatus.setSiteStatus(site, "Banned", comment)
+            tokenLifetime = int(days)
+            if tokenLifetime <= 0:
+                tokenExpiration = datetime.max
+            else:
+                tokenExpiration = datetime.utcnow().replace(microsecond=0) + timedelta(days=tokenLifetime)
+            result = self.sitestatus.setSiteStatus(site, "Banned", comment, expiry=tokenExpiration)
         else:
             result = WMSAdministratorClient().banSite(site, comment)
         if not result["OK"]:

src/DIRAC/Interfaces/scripts/dirac_admin_allow_site.py

@@ -17,6 +17,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")
@@ -35,9 +38,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
@@ -58,7 +64,7 @@ def getBoolean(value):
 
     # 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

@@ -17,6 +17,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")
@@ -36,9 +39,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
@@ -59,7 +65,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.
@@ -221,7 +221,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
 
@@ -258,6 +258,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():
         ("status=", "Status to be changed"),
         ("reason=", "Reason to set the Status"),
         ("VO=", "VO to change a status for. When omitted, status will be changed for all VOs"),
+        ("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:
@@ -180,7 +182,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}")
@@ -190,8 +196,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"],
@@ -202,7 +216,7 @@ def setStatus(switchDict, tokenOwner):
             reason=switchDict["reason"],
             vO=switchDict["VO"],
             tokenOwner=tokenOwner,
-            tokenExpiration=tomorrow,
+            tokenExpiration=tokenExpiration,
         )
         if not result["OK"]:
             return result