From 8e767e3c7c0d120f8ccae9351c06126efcb0e070 Mon Sep 17 00:00:00 2001 From: zhaoyu Date: Tue, 5 Aug 2025 13:06:41 -0400 Subject: [PATCH] Add refresh function in panda_auth --- doc/changes/DM-48912.feature.rst | 1 + .../lsst/ctrl/bps/panda/cli/cmd/__init__.py | 4 +- .../bps/panda/cli/cmd/panda_auth_commands.py | 18 +++- .../lsst/ctrl/bps/panda/panda_auth_drivers.py | 29 +++++- .../lsst/ctrl/bps/panda/panda_auth_utils.py | 97 +++++++++++++++++++ .../lsst/ctrl/bps/panda/panda_exceptions.py | 34 +++++++ tests/test_panda_auth_utils.py | 89 ++++++++++++++++- 7 files changed, 267 insertions(+), 5 deletions(-) create mode 100644 doc/changes/DM-48912.feature.rst create mode 100644 python/lsst/ctrl/bps/panda/panda_exceptions.py diff --git a/doc/changes/DM-48912.feature.rst b/doc/changes/DM-48912.feature.rst new file mode 100644 index 00000000..8de92888 --- /dev/null +++ b/doc/changes/DM-48912.feature.rst @@ -0,0 +1 @@ +Added refresh function in panda_auth diff --git a/python/lsst/ctrl/bps/panda/cli/cmd/__init__.py b/python/lsst/ctrl/bps/panda/cli/cmd/__init__.py index acbd89be..deb15602 100644 --- a/python/lsst/ctrl/bps/panda/cli/cmd/__init__.py +++ b/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 . -__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 diff --git a/python/lsst/ctrl/bps/panda/cli/cmd/panda_auth_commands.py b/python/lsst/ctrl/bps/panda/cli/cmd/panda_auth_commands.py index 5ae78b17..4e581f28 100644 --- a/python/lsst/ctrl/bps/panda/cli/cmd/panda_auth_commands.py +++ b/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) diff --git a/python/lsst/ctrl/bps/panda/panda_auth_drivers.py b/python/lsst/ctrl/bps/panda/panda_auth_drivers.py index e5f408c8..2baebbc1 100644 --- a/python/lsst/ctrl/bps/panda/panda_auth_drivers.py +++ b/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() diff --git a/python/lsst/ctrl/bps/panda/panda_auth_utils.py b/python/lsst/ctrl/bps/panda/panda_auth_utils.py index 4348ffbf..64193454 100644 --- a/python/lsst/ctrl/bps/panda/panda_auth_utils.py +++ b/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 diff --git a/python/lsst/ctrl/bps/panda/panda_exceptions.py b/python/lsst/ctrl/bps/panda/panda_exceptions.py new file mode 100644 index 00000000..827092ad --- /dev/null +++ b/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 diff --git a/tests/test_panda_auth_utils.py b/tests/test_panda_auth_utils.py index 16127f32..bf289fe7 100644 --- a/tests/test_panda_auth_utils.py +++ b/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()