Merge pull request #92 from lsst/tickets/DM-48912

Add refresh function in panda_auth

Author: zhaoyuyoung

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

@@ -33,6 +33,7 @@
 
 __all__ = [
     "panda_auth_clean_driver",
+    "panda_auth_refresh_driver",
     "panda_auth_reset_driver",
     "panda_auth_status_driver",
 ]
@@ -41,7 +42,19 @@
 import logging
 from datetime import datetime
 
-from .panda_auth_utils import panda_auth_clean, panda_auth_status, panda_auth_update
+from lsst.ctrl.bps.panda.panda_exceptions import (
+    PandaAuthError,
+    TokenExpiredError,
+    TokenNotFoundError,
+    TokenTooEarlyError,
+)
+
+from .panda_auth_utils import (
+    panda_auth_clean,
+    panda_auth_refresh,
+    panda_auth_status,
+    panda_auth_update,
+)
 
 _LOG = logging.getLogger(__name__)
 
@@ -56,6 +69,20 @@ def panda_auth_reset_driver():
     panda_auth_update(None, True)
 
 
+def panda_auth_refresh_driver(days, verbose):
+    """Refresh auth token."""
+    try:
+        panda_auth_refresh(days, verbose)
+    except TokenNotFoundError as e:
+        print(f"[ERROR] {e}")
+    except TokenExpiredError as e:
+        print(f"[ERROR] {e}")
+    except TokenTooEarlyError as e:
+        print(f"[INFO] {e}")
+    except PandaAuthError as e:
+        print(f"[FAIL] {e}")
+
+
 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,19 +30,32 @@
 __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 UTC, datetime, timedelta
 
 import idds.common.utils as idds_utils
 import pandaclient.idds_api
 from pandaclient.openidc_utils import OpenIdConnect_Utils
 
+from lsst.ctrl.bps.panda.panda_exceptions import (
+    AuthConfigError,
+    PandaAuthError,
+    TokenExpiredError,
+    TokenNotFoundError,
+    TokenRefreshError,
+    TokenTooEarlyError,
+)
+
 _LOG = logging.getLogger(__name__)
 
 
@@ -151,3 +164,87 @@ 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 the current valid IAM OpenID authentication token.
+
+    This function checks the expiration time of the existing token stored
+    in the local token file and attempts to refresh it if it is close to
+    expiring (within a specified number of days).
+
+    Parameters
+    ----------
+    days : `int`, optional
+        The minimum number of days before token expiration to trigger a
+        refresh. If the token expires in more than this number of days,
+        the refresh is skipped. Default is 4.
+    verbose : `bool`, optional
+        If True, enables verbose output for debugging or logging.
+        Default is False.
+
+    Returns
+    -------
+    status: `dict`
+        A dictionary containing the refreshed token status
+    """
+    panda_url = os.environ.get("PANDA_URL")
+    panda_auth_vo = os.environ.get("PANDA_AUTH_VO")
+
+    if not panda_url or not panda_auth_vo:
+        raise PandaAuthError("Missing required environment variables: PANDA_URL or 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 not os.path.exists(token_file):
+        raise TokenNotFoundError("Cannot find token file. Use 'panda_auth reset' to obtain a new token.")
+
+    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.fromtimestamp(dec["exp"], tz=UTC)
+    delta = exp_time - datetime.now(UTC)
+    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):
+        raise TokenExpiredError("Token already expired. Cannot refresh.")
+    elif delta > timedelta(days=days):
+        raise TokenTooEarlyError(
+            f"Too early to refresh. More than {days} day(s) until expiration.\n"
+            f"Use '--days' option to adjust threshold, e.g.:\n"
+            f"  panda_auth refresh --days 10"
+        )
+
+    refresh_token_string = data["refresh_token"]
+
+    s, auth_config = open_id.fetch_page(open_id.auth_config_url)
+    if not s:
+        raise AuthConfigError("Failed to get Auth configuration.")
+
+    s, endpoint_config = open_id.fetch_page(auth_config["oidc_config_url"])
+    if not s:
+        raise AuthConfigError("Failed to get endpoint configuration.")
+
+    s, o = open_id.refresh_token(
+        endpoint_config["token_endpoint"],
+        auth_config["client_id"],
+        auth_config["client_secret"],
+        refresh_token_string,
+    )
+
+    if not s:
+        raise TokenRefreshError("Failed to refresh token.")
+
+    status = panda_auth_status()
+    if status:
+        exp_time = datetime.fromtimestamp(status["exp"], tz=UTC)
+        print(f"{'New expiration time:':23} {exp_time.strftime('%Y-%m-%d %H:%M:%S')} UTC")
+    print("Success to refresh token")
+    return status

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

@@ -0,0 +1,34 @@
+class PandaAuthError(Exception):
+    """Base class for authentication errors."""
+
+    pass
+
+
+class TokenNotFoundError(PandaAuthError):
+    """Raised when the token file is missing."""
+
+    pass
+
+
+class TokenExpiredError(PandaAuthError):
+    """Raised when the token has already expired."""
+
+    pass
+
+
+class TokenTooEarlyError(PandaAuthError):
+    """Raised when attempting to refresh too early."""
+
+    pass
+
+
+class AuthConfigError(PandaAuthError):
+    """Raised when fetching the auth or endpoint configuration fails."""
+
+    pass
+
+
+class TokenRefreshError(PandaAuthError):
+    """Raised when token refresh fails."""
+
+    pass

tests/test_panda_auth_utils.py

@@ -27,12 +27,49 @@
 
 """Unit tests for PanDA authentication utilities."""
 
+import base64
+import json
 import os
 import unittest
+from datetime import UTC, datetime, timedelta
 from unittest import mock
 
 from lsst.ctrl.bps.panda import __version__ as version
-from lsst.ctrl.bps.panda.panda_auth_utils import panda_auth_status
+from lsst.ctrl.bps.panda.panda_auth_utils import (
+    TokenExpiredError,
+    panda_auth_refresh,
+    panda_auth_status,
+)
+
+
+def make_fake_jwt(exp_offset_days):
+    """Return a fake id_token that expires in N days."""
+    payload = {"exp": int((datetime.now(UTC) + timedelta(days=exp_offset_days)).timestamp())}
+    b64_payload = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode().rstrip("=")
+    return f"header.{b64_payload}.sig"
+
+
+def fake_token_file(exp_days=3, refresh_token="fake_refresh"):
+    """Generate fake token file data"""
+    token = make_fake_jwt(exp_days)
+    return json.dumps({"id_token": token, "refresh_token": refresh_token})
+
+
+def fetch_page_side_effect(url):
+    """Simulate OpenIdConnect_Utils.fetch_page behavior in tests."""
+    if url.endswith("auth_config.json"):
+        return True, {
+            "client_secret": "secret",
+            "audience": "https://iam.example.com",
+            "client_id": "cid",
+            "oidc_config_url": "https://oidc.example.org/.well-known/openid-configuration",
+            "vo": "fake_vo",
+            "no_verify": "True",
+            "robot_ids": "NONE",
+        }
+    elif url.endswith("openid-configuration"):
+        return True, {"token_endpoint": "https://oidc.example.org/token"}
+    return False, {}
 
 
 class VersionTestCase(unittest.TestCase):
@@ -46,6 +83,21 @@ def test_version(self):
 class TestPandaAuthUtils(unittest.TestCase):
     """Simple test of auth utilities."""
 
+    def setUp(self):
+        self.test_env = {
+            "PANDA_CONFIG_ROOT": "/fake/token",
+            "PANDA_URL_SSL": "https://fake.server.com:8443/server/panda",
+            "PANDA_URL": "https://fake.server.com:8443/server/panda",
+            "PANDACACHE_URL": "https://fake.server.com:8443/server/panda",
+            "PANDAMON_URL": "https://fake.monitor.com:8443/",
+            "PANDA_AUTH": "oidc",
+            "PANDA_VERIFY_HOST": "off",
+            "PANDA_AUTH_VO": "fake_vo",
+            "PANDA_BEHIND_REAL_LB": "true",
+            "PANDA_SYS": "/fake/pandasys",
+            "IDDS_CONFIG": "/fake/pandasys/etc/idds/idds.cfg.client.template",
+        }
+
     def testPandaAuthStatusWrongEnviron(self):
         unwanted = {
             "PANDA_AUTH",
@@ -59,6 +111,41 @@ def testPandaAuthStatusWrongEnviron(self):
             with self.assertRaises(OSError):
                 panda_auth_status()
 
+    @mock.patch("builtins.print")
+    @mock.patch("os.path.exists", return_value=True)
+    @mock.patch("pandaclient.openidc_utils.OpenIdConnect_Utils")
+    def test_expired_token(self, mock_oidc, mock_exists, mock_print):
+        mock_oidc.return_value.get_token_path.return_value = "/fake/token.json"
+
+        with mock.patch.dict("os.environ", self.test_env):
+            with mock.patch("builtins.open", mock.mock_open(read_data=fake_token_file(exp_days=-1))):
+                with self.assertRaises(TokenExpiredError):
+                    panda_auth_refresh(days=4)
+
+    @mock.patch("builtins.print")
+    @mock.patch("lsst.ctrl.bps.panda.panda_auth_utils.panda_auth_status")
+    @mock.patch("os.path.exists", return_value=True)
+    @mock.patch("lsst.ctrl.bps.panda.panda_auth_utils.OpenIdConnect_Utils")
+    def test_successful_refresh(self, mock_oidc, mock_exists, mock_status, mock_print):
+        fake_openid = mock_oidc.return_value
+        fake_openid.get_token_path.return_value = "/fake/token.json"
+        fake_openid.auth_config_url = "https://fake.server/auth_config.json"
+
+        fake_openid.fetch_page.side_effect = fetch_page_side_effect
+
+        fake_openid.refresh_token.return_value = (True, {"access_token": "new_token"})
+
+        mock_status.return_value = {"exp": int((datetime.now(UTC) + timedelta(seconds=3600)).timestamp())}
+
+        with mock.patch.dict("os.environ", self.test_env):
+            token_json = fake_token_file(exp_days=2)
+            with mock.patch("builtins.open", mock.mock_open(read_data=token_json)):
+                panda_auth_refresh(days=4)
+
+        fake_openid.refresh_token.assert_called_once()
+        found = any("Success to refresh token" in str(c[0][0]) for c in mock_print.call_args_list)
+        assert found
+
 
 if __name__ == "__main__":
     unittest.main()