Add refresh function in panda_auth

Author: zhaoyu

doc/changes/DM-48912.feature.rst

@@ -0,0 +1 @@
+Added refresh function in panda_auth

python/lsst/ctrl/bps/panda/cli/cmd/__init__.py

@@ -25,6 +25,6 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-__all__ = ["clean", "reset", "status"]
+__all__ = ["clean", "reset", "refresh", "status"]
 
-from .panda_auth_commands import clean, reset, status
+from .panda_auth_commands import clean, reset, refresh, status

python/lsst/ctrl/bps/panda/cli/cmd/panda_auth_commands.py

@@ -28,6 +28,7 @@
 
 __all__ = [
     "clean",
+    "refresh",
     "reset",
     "status",
 ]
@@ -37,7 +38,12 @@
 
 from lsst.daf.butler.cli.utils import MWCommand
 
-from ...panda_auth_drivers import panda_auth_clean_driver, panda_auth_reset_driver, panda_auth_status_driver
+from ...panda_auth_drivers import (
+    panda_auth_clean_driver,
+    panda_auth_refresh_driver,
+    panda_auth_reset_driver,
+    panda_auth_status_driver,
+)
 
 
 class PandaAuthCommand(MWCommand):
@@ -62,3 +68,13 @@ def reset(*args, **kwargs):
 def clean(*args, **kwargs):
     """Clean up token and token cache files."""
     panda_auth_clean_driver(*args, **kwargs)
+
+
+@click.command(cls=PandaAuthCommand)
+@click.option("--days", default=4, help="The earlist remaining days to refresh the token.")
+@click.option("--verbose", is_flag=True, help="Enable verbose output")
+def refresh(*args, **kwargs):
+    """Refresh auth tocken."""
+    days = kwargs.get("days", 4)
+    verbose = kwargs.get("verbose", False)
+    panda_auth_refresh_driver(days, verbose)

python/lsst/ctrl/bps/panda/panda_auth_drivers.py

@@ -41,7 +41,12 @@
 import logging
 from datetime import datetime
 
-from .panda_auth_utils import panda_auth_clean, panda_auth_status, panda_auth_update
+from .panda_auth_utils import (
+    panda_auth_clean,
+    panda_auth_refresh,
+    panda_auth_status,
+    panda_auth_update,
+)
 
 _LOG = logging.getLogger(__name__)
 
@@ -56,6 +61,11 @@ def panda_auth_reset_driver():
     panda_auth_update(None, True)
 
 
+def panda_auth_refresh_driver(days, verbose):
+    """Get new auth token."""
+    panda_auth_refresh(days, verbose)
+
+
 def panda_auth_status_driver():
     """Gather information about a token if it exists."""
     status = panda_auth_status()

python/lsst/ctrl/bps/panda/panda_auth_utils.py

@@ -30,14 +30,18 @@
 __all__ = [
     "panda_auth_clean",
     "panda_auth_expiration",
+    "panda_auth_refresh",
     "panda_auth_setup",
     "panda_auth_status",
     "panda_auth_update",
 ]
 
 
+import base64
+import json
 import logging
 import os
+from datetime import datetime, timedelta
 
 import idds.common.utils as idds_utils
 import pandaclient.idds_api
@@ -151,3 +155,65 @@ def panda_auth_update(idds_server=None, reset=False):
         # idds server given.  So for now, check result string for keywords.
         if "request_id" not in ret[1][-1] or "status" not in ret[1][-1]:
             raise RuntimeError(f"Error contacting PanDA service: {ret}")
+
+
+def panda_auth_refresh(days=4, verbose=False):
+    """Refresh auth token"""
+    panda_url = os.environ.get("PANDA_URL")
+    panda_auth_vo = os.environ.get("PANDA_AUTH_VO")
+    url_prefix = panda_url.split("/server", 1)[0]
+    auth_url = f"{url_prefix}/auth/{panda_auth_vo}_auth_config.json"
+    open_id = OpenIdConnect_Utils(auth_url, log_stream=_LOG, verbose=verbose)
+
+    token_file = open_id.get_token_path()
+    if os.path.exists(token_file):
+        with open(token_file) as f:
+            data = json.load(f)
+            enc = data["id_token"].split(".")[1]
+            enc += "=" * (-len(enc) % 4)
+            dec = json.loads(base64.urlsafe_b64decode(enc.encode()))
+            exp_time = datetime.utcfromtimestamp(dec["exp"])
+            delta = exp_time - datetime.utcnow()
+            minutes = delta.total_seconds() / 60
+            print(f"Token will expire in {minutes} minutes.")
+            print(f"Token expiration time : {exp_time.strftime('%Y-%m-%d %H:%M:%S')} UTC")
+            if delta < timedelta(minutes=0):
+                print("Token already expired. Cannot refresh.")
+                return
+            elif delta > timedelta(days=days):
+                print("\n" + "=" * 60)
+                print(f" Too early to refresh. More than {days} day(s) until expiration.")
+                print(" To change this threshold, use the --days option:")
+                print("   Example: panda_auth refresh --days 10")
+                print("=" * 60 + "\n")
+                return
+    else:
+        print("Cannot find token file.")
+        return
+    refresh_token_string = data["refresh_token"]
+
+    s, auth_config = open_id.fetch_page(open_id.auth_config_url)
+    if not s:
+        print("Failed to get Auth configuration")
+        return
+
+    s, endpoint_config = open_id.fetch_page(auth_config["oidc_config_url"])
+    if not s:
+        print("Failed to get endpoint configuration")
+        return
+
+    s, o = open_id.refresh_token(
+        endpoint_config["token_endpoint"],
+        auth_config["client_id"],
+        auth_config["client_secret"],
+        refresh_token_string,
+    )
+
+    if not s:
+        print("Failed to refresh token")
+        return
+    else:
+        status = panda_auth_status()
+        if status:
+            print(f"{'New expiration time:':23} {datetime.utcfromtimestamp(status['exp'])} UTC")
+        print("Success to refresh token")