diff --git a/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/exceptions.py b/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/exceptions.py index 32d8f9e51..e66859460 100644 --- a/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/exceptions.py +++ b/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/exceptions.py @@ -211,13 +211,26 @@ def wrapped(*args, **kwargs): return wrapped +class _ReauthSucceeded(Exception): + """Internal signal that re-authentication succeeded and the caller should retry.""" + + def _handle_connection_error_with_reauth(exc, login_func): - """Handle ConnectionError with reauthentication logic.""" + """Handle ConnectionError with reauthentication logic. + + If the error indicates an expired token, triggers re-authentication via + ``login_func`` and raises ``_ReauthSucceeded`` so the calling decorator + can retry the original operation. Non-expired connection errors are + surfaced as ``ClickExceptionRed`` directly. + """ if "expired" in str(exc).lower(): - click.echo(click.style("Token is expired, triggering re-authentication", fg="red")) + click.echo(click.style("Token is expired, triggering re-authentication", fg="yellow")) config = exc.get_config() - login_func(config) - raise ClickExceptionRed("Please try again now") from None + try: + login_func(config) + except Exception as reauth_exc: + raise ClickExceptionRed(f"Re-authentication failed: {reauth_exc}") from None + raise _ReauthSucceeded() from None else: raise ClickExceptionRed(str(exc)) from None @@ -239,26 +252,78 @@ def _handle_exception_group_with_reauth(eg, login_func) -> NoReturn: raise eg +def _raise_if_mappable(exc: BaseException) -> None: + """Raise a user-friendly ClickException if *exc* maps to one, otherwise return.""" + if cli_exc := _map_cli_exception(exc): + raise cli_exc from None + + +def _try_reauth_or_handle(handler, exc, login_func, *, allow_reauth: bool) -> bool: + """Attempt re-auth via *handler*; return True if re-auth succeeded. + + When *allow_reauth* is False the handler is called without catching + ``_ReauthSucceeded`` — any re-auth signal propagates as a normal error. + """ + if allow_reauth: + try: + handler(exc, login_func) + except _ReauthSucceeded: + return True + else: + handler(exc, login_func) + return False + + +def _call_with_exception_handling(func, args, kwargs, login_func, *, allow_reauth: bool): + """Call *func* and handle exceptions, optionally allowing re-authentication. + + Returns ``(result, needs_retry)`` where *needs_retry* is ``True`` when + re-authentication succeeded and the caller should retry the call. + """ + try: + return func(*args, **kwargs), False + except _ReauthSucceeded: + raise ClickExceptionRed("Unexpected re-auth signal") from None + except BaseExceptionGroup as eg: + if _try_reauth_or_handle(_handle_exception_group_with_reauth, eg, login_func, allow_reauth=allow_reauth): + return None, True + except (ConnectionError, JumpstarterException, click.ClickException) as e: + if isinstance(e, ConnectionError) and not allow_reauth: + raise ClickExceptionRed(str(e)) from None + if _try_reauth_or_handle(_handle_single_exception_with_reauth, e, login_func, allow_reauth=allow_reauth): + return None, True + except Exception as e: + _raise_if_mappable(e) + raise + except KeyboardInterrupt as e: + _raise_if_mappable(e) + raise + return None, False # pragma: no cover — handlers above always raise or return + + def handle_exceptions_with_reauthentication(login_func): - """Decorator to handle exceptions in blocking functions, including those wrapped in BaseExceptionGroup.""" + """Decorator to handle exceptions in blocking functions, including those wrapped in BaseExceptionGroup. + + When a ``ConnectionError`` with an expired-token message is caught, the + decorator triggers re-authentication via *login_func* and **retries the + original call exactly once**. If the retry also fails, the error is + surfaced normally — no infinite loop. + """ def decorator(func): @wraps(func) def wrapped(*args, **kwargs): - try: - return func(*args, **kwargs) - except BaseExceptionGroup as eg: - _handle_exception_group_with_reauth(eg, login_func) - except (ConnectionError, JumpstarterException, click.ClickException) as e: - _handle_single_exception_with_reauth(e, login_func) - except Exception as e: - if cli_exc := _map_cli_exception(e): - raise cli_exc from None - raise - except KeyboardInterrupt as e: - if cli_exc := _map_cli_exception(e): - raise cli_exc from None - raise + result, needs_retry = _call_with_exception_handling( + func, args, kwargs, login_func, allow_reauth=True + ) + if not needs_retry: + return result + + click.echo(click.style("Re-authenticated, retrying...", fg="yellow")) + result, _ = _call_with_exception_handling( + func, args, kwargs, login_func, allow_reauth=False + ) + return result return wrapped diff --git a/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/exceptions_test.py b/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/exceptions_test.py index d6d761508..2821c1b1d 100644 --- a/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/exceptions_test.py +++ b/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/exceptions_test.py @@ -1,5 +1,6 @@ import ssl from json import JSONDecodeError +from unittest.mock import MagicMock import click import pytest @@ -10,6 +11,8 @@ handle_exceptions_with_reauthentication, ) +from jumpstarter.common.exceptions import ConnectionError as JmpConnectionError + @pytest.fixture def anyio_backend(): @@ -183,3 +186,79 @@ def grpc_precondition_no_details_fn(): with pytest.raises(click.ClickException, match="precondition"): grpc_precondition_no_details_fn() + + +# --------------------------------------------------------------------------- +# Tests for automatic retry after successful re-authentication (NS-REQ-1/3) +# --------------------------------------------------------------------------- + + +def _make_expired_connection_error(): + """Create a ConnectionError that looks like an expired-token error.""" + exc = JmpConnectionError("token expired") + config = MagicMock(name="client_config") + exc.set_config(config) + return exc, config + + +def test_reauth_retries_on_success() -> None: + """After successful re-auth the decorator retries and returns the result (TS-NS-1).""" + call_count = 0 + + def login_func(_config): + pass # success + + @handle_exceptions_with_reauthentication(login_func) + def fn(): + nonlocal call_count + call_count += 1 + if call_count == 1: + exc, _ = _make_expired_connection_error() + raise exc + return "sentinel" + + result = fn() + assert result == "sentinel" + assert call_count == 2 + + +def test_reauth_failure_raises_click_exception() -> None: + """If login_func raises, the decorator surfaces a ClickException (TS-NS-2).""" + call_count = 0 + + def login_func(_config): + raise RuntimeError("IdP unreachable") + + @handle_exceptions_with_reauthentication(login_func) + def fn(): + nonlocal call_count + call_count += 1 + exc, _ = _make_expired_connection_error() + raise exc + + with pytest.raises(click.ClickException, match="Re-authentication failed"): + fn() + + # The wrapped function should have been called only once (no retry). + assert call_count == 1 + + +def test_reauth_retry_bounded_to_one_attempt() -> None: + """The decorator retries at most once; a second failure surfaces normally (TS-NS-3).""" + call_count = 0 + + def login_func(_config): + pass # always succeeds + + @handle_exceptions_with_reauthentication(login_func) + def fn(): + nonlocal call_count + call_count += 1 + exc, _ = _make_expired_connection_error() + raise exc + + with pytest.raises(click.ClickException): + fn() + + # Original call + exactly one retry = 2 total. + assert call_count == 2 diff --git a/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/oidc.py b/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/oidc.py index 08661c853..7af009f86 100644 --- a/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/oidc.py +++ b/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/oidc.py @@ -1,7 +1,9 @@ +import asyncio import json import os import ssl import time +import warnings from dataclasses import dataclass from functools import wraps from typing import ClassVar @@ -12,11 +14,18 @@ from aiohttp import web from anyio import create_memory_object_stream from anyio.to_thread import run_sync -from authlib.integrations.requests_client import OAuth2Session -from joserfc.jws import extract_compact -from yarl import URL -from jumpstarter.config.env import JMP_OIDC_CALLBACK_PORT +# Suppress AuthlibDeprecationWarning emitted unconditionally during the authlib +# import chain (authlib._joserfc_helpers -> authlib.jose). The project already +# uses joserfc directly for JWT operations; authlib is only needed for +# OAuth2Session. The warning provides no actionable information to end users. +warnings.filterwarnings("ignore", category=DeprecationWarning, module=r"authlib\.") + +from authlib.integrations.requests_client import OAuth2Session # noqa: E402 +from joserfc.jws import extract_compact # noqa: E402 +from yarl import URL # noqa: E402 + +from jumpstarter.config.env import JMP_OIDC_CALLBACK_PORT, JMP_OIDC_DEVICE_FLOW # noqa: E402 def _get_ssl_context() -> ssl.SSLContext: @@ -46,6 +55,14 @@ def opt_oidc(f): default=True, help="Request offline_access scope (refresh token)", ) + @click.option( + "--device-flow", + "device_flow", + is_flag=True, + default=False, + help="Use OAuth 2.0 Device Authorization Grant (RFC 8628) instead of authorization code flow. " + "Useful in headless or containerized environments where localhost callbacks are not available.", + ) @wraps(f) def wrapper(*args, **kwds): return f(*args, **kwds) @@ -53,6 +70,18 @@ def wrapper(*args, **kwds): return wrapper +def should_use_device_flow(device_flow_flag: bool) -> bool: + """Determine whether to use the device authorization grant flow. + + Returns True if: + - The --device-flow CLI flag was explicitly passed, OR + - The JMP_OIDC_DEVICE_FLOW environment variable is set to "1" + """ + if device_flow_flag: + return True + return os.environ.get(JMP_OIDC_DEVICE_FLOW) == "1" + + @dataclass(kw_only=True) class Config: issuer: str @@ -78,7 +107,15 @@ def _scopes(self) -> list[str]: def client(self, **kwargs): session = OAuth2Session(client_id=self.client_id, scope=self._scopes(), **kwargs) - session.verify = False if self.insecure_tls else (os.environ.get("SSL_CERT_FILE") or certifi.where()) + if self.insecure_tls: + session.verify = False + # The user has already opted into insecure TLS (via --insecure flag + # or config), so urllib3's InsecureRequestWarning is redundant noise. + import urllib3 + + urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + else: + session.verify = os.environ.get("SSL_CERT_FILE") or certifi.where() return session async def token_exchange_grant(self, token: str, **kwargs): @@ -177,6 +214,109 @@ async def callback(request): lambda: client.fetch_token(config["token_endpoint"], authorization_response=authorization_response) ) + async def device_authorization_grant(self): + """Perform OAuth 2.0 Device Authorization Grant (RFC 8628). + + This flow is suitable for headless or containerized environments where + a localhost callback server is not accessible from the user's browser. + + The flow: + 1. Request a device code from the authorization server. + 2. Display a verification URL and user code to the user. + 3. Poll the token endpoint until the user completes authorization. + """ + config = await self.configuration() + + device_endpoint = config.get("device_authorization_endpoint") + if not device_endpoint: + raise click.ClickException( + "The identity provider does not support Device Authorization Grant (RFC 8628). " + "The OIDC discovery document does not include a 'device_authorization_endpoint'. " + "Contact your IdP administrator to enable device flow, or use a different login method." + ) + + token_endpoint = config["token_endpoint"] + + ssl_context: ssl.SSLContext | bool = False if self.insecure_tls else _get_ssl_context() + connector = aiohttp.TCPConnector(ssl=ssl_context) + + async with aiohttp.ClientSession(connector=connector) as session: + # Step 1: Request device authorization + async with session.post( + device_endpoint, + data={ + "client_id": self.client_id, + "scope": " ".join(self._scopes()), + }, + ) as response: + if response.status != 200: + text = await response.text() + raise click.ClickException( + f"Device authorization request failed (HTTP {response.status}): {text}" + ) + device_data = await response.json() + + device_code = device_data["device_code"] + interval = device_data.get("interval", 5) + expires_in = device_data.get("expires_in", 600) + + # Step 2: Display verification URI to user + verification_uri_complete = device_data.get("verification_uri_complete") + if verification_uri_complete: + click.echo(f"To sign in, open the following URL in your browser:\n\n {verification_uri_complete}\n") + else: + verification_uri = device_data.get("verification_uri") + user_code = device_data.get("user_code") + click.echo( + f"To sign in, open the following URL in your browser:\n\n {verification_uri}\n\n" + f"Then enter the code: {user_code}\n" + ) + + click.echo("Waiting for authentication...") + + # Step 3: Poll the token endpoint + deadline = time.monotonic() + expires_in + while time.monotonic() < deadline: + await asyncio.sleep(interval) + + async with session.post( + token_endpoint, + data={ + "grant_type": "urn:ietf:params:oauth:grant-type:device_code", + "device_code": device_code, + "client_id": self.client_id, + }, + ) as token_response: + token_data = await token_response.json() + + if token_response.status == 200: + return token_data + + error = token_data.get("error", "") + if error == "authorization_pending": + continue + elif error == "slow_down": + interval += 5 + continue + elif error == "expired_token": + raise click.ClickException( + "Device authorization has expired. Please try again." + ) + elif error == "access_denied": + raise click.ClickException( + "Authorization request was denied by the user." + ) + else: + error_description = token_data.get("error_description", "") + raise click.ClickException( + f"Device authorization failed: {error}" + + (f" - {error_description}" if error_description else "") + ) + + raise click.ClickException( + "Device authorization timed out waiting for user approval. Please try again." + ) + def decode_jwt(token: str): try: diff --git a/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/oidc_test.py b/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/oidc_test.py index 3b67db02a..b290c3fd4 100644 --- a/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/oidc_test.py +++ b/python/packages/jumpstarter-cli-common/jumpstarter_cli_common/oidc_test.py @@ -1,6 +1,11 @@ import ssl +import warnings +from unittest.mock import AsyncMock, MagicMock, patch -from jumpstarter_cli_common.oidc import Config, _get_ssl_context +import click +import pytest + +from jumpstarter_cli_common.oidc import Config, _get_ssl_context, should_use_device_flow class TestConfigInsecureTls: @@ -23,3 +28,300 @@ class TestGetSslContext: def test_returns_ssl_context(self) -> None: ctx = _get_ssl_context() assert isinstance(ctx, ssl.SSLContext) + + +class TestShouldUseDeviceFlow: + def test_returns_true_when_flag_is_set(self) -> None: + assert should_use_device_flow(device_flow_flag=True) is True + + def test_returns_true_when_env_var_is_1(self, monkeypatch) -> None: + monkeypatch.setenv("JMP_OIDC_DEVICE_FLOW", "1") + assert should_use_device_flow(device_flow_flag=False) is True + + def test_returns_false_when_env_var_is_not_1(self, monkeypatch) -> None: + monkeypatch.setenv("JMP_OIDC_DEVICE_FLOW", "0") + assert should_use_device_flow(device_flow_flag=False) is False + + def test_returns_false_when_env_var_unset(self, monkeypatch) -> None: + monkeypatch.delenv("JMP_OIDC_DEVICE_FLOW", raising=False) + assert should_use_device_flow(device_flow_flag=False) is False + + def test_flag_takes_priority_over_env(self, monkeypatch) -> None: + monkeypatch.delenv("JMP_OIDC_DEVICE_FLOW", raising=False) + assert should_use_device_flow(device_flow_flag=True) is True + + +def _make_async_cm(response): + """Create an async context manager wrapper for a MagicMock response.""" + cm = MagicMock() + cm.__aenter__ = AsyncMock(return_value=response) + cm.__aexit__ = AsyncMock(return_value=False) + return cm + + +class TestDeviceAuthorizationGrant: + @pytest.mark.asyncio + async def test_raises_when_no_device_endpoint_in_discovery(self) -> None: + config = Config(issuer="https://auth.example.com", client_id="test") + with patch.object(config, "configuration", new_callable=AsyncMock) as mock_config: + mock_config.return_value = { + "token_endpoint": "https://auth.example.com/token", + # No device_authorization_endpoint + } + with pytest.raises(click.ClickException, match="does not support Device Authorization Grant"): + await config.device_authorization_grant() + + @pytest.mark.asyncio + async def test_error_message_mentions_device_authorization_endpoint(self) -> None: + config = Config(issuer="https://auth.example.com", client_id="test") + with patch.object(config, "configuration", new_callable=AsyncMock) as mock_config: + mock_config.return_value = { + "token_endpoint": "https://auth.example.com/token", + } + with pytest.raises(click.ClickException, match="device_authorization_endpoint"): + await config.device_authorization_grant() + + @pytest.mark.asyncio + async def test_successful_device_flow_with_verification_uri_complete(self) -> None: + config = Config(issuer="https://auth.example.com", client_id="test") + + discovery = { + "token_endpoint": "https://auth.example.com/token", + "device_authorization_endpoint": "https://auth.example.com/device", + } + + device_response_data = { + "device_code": "test-device-code", + "user_code": "ABCD-EFGH", + "verification_uri": "https://auth.example.com/device", + "verification_uri_complete": "https://auth.example.com/device?user_code=ABCD-EFGH", + "interval": 0.01, # Speed up test + "expires_in": 300, + } + + token_data = { + "access_token": "test-access-token", + "refresh_token": "test-refresh-token", + "token_type": "Bearer", + } + + # Track poll count: first returns authorization_pending, second returns success + poll_count = 0 + + def mock_post(url, data=None, **kwargs): + nonlocal poll_count + response = MagicMock() + + if "device" in str(url) and "grant_type" not in (data or {}): + # Device authorization endpoint + response.status = 200 + response.json = AsyncMock(return_value=device_response_data) + response.text = AsyncMock(return_value="") + else: + # Token endpoint + poll_count += 1 # ty: ignore[unresolved-reference] + if poll_count == 1: + response.status = 400 + response.json = AsyncMock(return_value={"error": "authorization_pending"}) + else: + response.status = 200 + response.json = AsyncMock(return_value=token_data) + + return _make_async_cm(response) + + mock_session = MagicMock() + mock_session.post = mock_post + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=False) + + with ( + patch.object(config, "configuration", new_callable=AsyncMock, return_value=discovery), + patch("jumpstarter_cli_common.oidc.aiohttp.ClientSession", return_value=mock_session), + ): + result = await config.device_authorization_grant() + + assert result["access_token"] == "test-access-token" + assert result["refresh_token"] == "test-refresh-token" + assert poll_count == 2 + + @pytest.mark.asyncio + async def test_handles_slow_down_response(self) -> None: + config = Config(issuer="https://auth.example.com", client_id="test") + + discovery = { + "token_endpoint": "https://auth.example.com/token", + "device_authorization_endpoint": "https://auth.example.com/device", + } + + device_response_data = { + "device_code": "test-device-code", + "user_code": "ABCD-EFGH", + "verification_uri": "https://auth.example.com/device", + "interval": 0.01, + "expires_in": 300, + } + + token_data = {"access_token": "test-access-token", "token_type": "Bearer"} + + poll_count = 0 + + def mock_post(url, data=None, **kwargs): + nonlocal poll_count + response = MagicMock() + + if "device" in str(url) and "grant_type" not in (data or {}): + response.status = 200 + response.json = AsyncMock(return_value=device_response_data) + response.text = AsyncMock(return_value="") + else: + poll_count += 1 # ty: ignore[unresolved-reference] + if poll_count == 1: + response.status = 400 + response.json = AsyncMock(return_value={"error": "slow_down"}) + else: + response.status = 200 + response.json = AsyncMock(return_value=token_data) + + return _make_async_cm(response) + + mock_session = MagicMock() + mock_session.post = mock_post + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=False) + + with ( + patch.object(config, "configuration", new_callable=AsyncMock, return_value=discovery), + patch("jumpstarter_cli_common.oidc.aiohttp.ClientSession", return_value=mock_session), + ): + result = await config.device_authorization_grant() + + assert result["access_token"] == "test-access-token" + + @pytest.mark.asyncio + async def test_raises_on_access_denied(self) -> None: + config = Config(issuer="https://auth.example.com", client_id="test") + + discovery = { + "token_endpoint": "https://auth.example.com/token", + "device_authorization_endpoint": "https://auth.example.com/device", + } + + device_response_data = { + "device_code": "test-device-code", + "user_code": "ABCD-EFGH", + "verification_uri": "https://auth.example.com/device", + "interval": 0.01, + "expires_in": 300, + } + + def mock_post(url, data=None, **kwargs): + response = MagicMock() + + if "device" in str(url) and "grant_type" not in (data or {}): + response.status = 200 + response.json = AsyncMock(return_value=device_response_data) + response.text = AsyncMock(return_value="") + else: + response.status = 400 + response.json = AsyncMock(return_value={"error": "access_denied"}) + + return _make_async_cm(response) + + mock_session = MagicMock() + mock_session.post = mock_post + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=False) + + with ( + patch.object(config, "configuration", new_callable=AsyncMock, return_value=discovery), + patch("jumpstarter_cli_common.oidc.aiohttp.ClientSession", return_value=mock_session), + ): + with pytest.raises(click.ClickException, match="denied by the user"): + await config.device_authorization_grant() + + @pytest.mark.asyncio + async def test_raises_on_expired_token(self) -> None: + config = Config(issuer="https://auth.example.com", client_id="test") + + discovery = { + "token_endpoint": "https://auth.example.com/token", + "device_authorization_endpoint": "https://auth.example.com/device", + } + + device_response_data = { + "device_code": "test-device-code", + "user_code": "ABCD-EFGH", + "verification_uri": "https://auth.example.com/device", + "interval": 0.01, + "expires_in": 300, + } + + def mock_post(url, data=None, **kwargs): + response = MagicMock() + + if "device" in str(url) and "grant_type" not in (data or {}): + response.status = 200 + response.json = AsyncMock(return_value=device_response_data) + response.text = AsyncMock(return_value="") + else: + response.status = 400 + response.json = AsyncMock(return_value={"error": "expired_token"}) + + return _make_async_cm(response) + + mock_session = MagicMock() + mock_session.post = mock_post + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=False) + + with ( + patch.object(config, "configuration", new_callable=AsyncMock, return_value=discovery), + patch("jumpstarter_cli_common.oidc.aiohttp.ClientSession", return_value=mock_session), + ): + with pytest.raises(click.ClickException, match="expired"): + await config.device_authorization_grant() + + +# --------------------------------------------------------------------------- +# Warning suppression tests (NS-REQ-4, NS-REQ-5) +# --------------------------------------------------------------------------- + + +class TestAuthlibDeprecationWarningSuppressed: + """TS-NS-4: importing oidc must not emit AuthlibDeprecationWarning.""" + + def test_no_authlib_deprecation_warning_on_import(self) -> None: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + # Force re-evaluation of the filter + import chain + import importlib + + import jumpstarter_cli_common.oidc + + importlib.reload(jumpstarter_cli_common.oidc) + + authlib_warnings = [ + w + for w in caught + if issubclass(w.category, DeprecationWarning) and "authlib" in str(w.message).lower() + ] + assert authlib_warnings == [], f"Unexpected authlib deprecation warnings: {authlib_warnings}" + + +class TestInsecureRequestWarningSuppressed: + """TS-NS-5: Config.client() with insecure_tls=True suppresses InsecureRequestWarning.""" + + def test_urllib3_insecure_request_warning_suppressed(self) -> None: + import urllib3.exceptions + + config = Config(issuer="https://auth.example.com", client_id="test", insecure_tls=True) + config.client() + + # After calling client() with insecure_tls=True, the urllib3 + # InsecureRequestWarning should be in the warning filters. + matching_filters = [ + f + for f in warnings.filters + if len(f) >= 3 and f[0] == "ignore" and f[2] is urllib3.exceptions.InsecureRequestWarning + ] + assert len(matching_filters) > 0, "InsecureRequestWarning was not suppressed" diff --git a/python/packages/jumpstarter-cli/jumpstarter_cli/login.py b/python/packages/jumpstarter-cli/jumpstarter_cli/login.py index b1b0f240d..111d89f7c 100644 --- a/python/packages/jumpstarter-cli/jumpstarter_cli/login.py +++ b/python/packages/jumpstarter-cli/jumpstarter_cli/login.py @@ -8,7 +8,7 @@ from jumpstarter_cli_common.blocking import blocking from jumpstarter_cli_common.config import opt_config from jumpstarter_cli_common.exceptions import handle_exceptions -from jumpstarter_cli_common.oidc import Config, decode_jwt_issuer, opt_oidc +from jumpstarter_cli_common.oidc import Config, decode_jwt_issuer, opt_oidc, should_use_device_flow from jumpstarter_cli_common.opt import confirm_insecure_tls, opt_insecure_tls, opt_nointeractive from jumpstarter.common.exceptions import ReauthenticationFailed @@ -160,6 +160,7 @@ async def login( # noqa: C901 connector_id: str, callback_port: int | None, offline_access: bool, + device_flow: bool, unsafe, insecure_tls: bool, nointeractive: bool, @@ -341,6 +342,8 @@ def save_config() -> None: tokens = await oidc.token_exchange_grant(token, **kwargs) elif username is not None and password is not None: tokens = await oidc.password_grant(username, password) + elif should_use_device_flow(device_flow): + tokens = await oidc.device_authorization_grant() else: tokens = await oidc.authorization_code_grant(callback_port=callback_port) @@ -389,7 +392,10 @@ async def relogin_client(config: ClientConfigV1Alpha1): except Exception: pass - tokens = await oidc.authorization_code_grant() + if should_use_device_flow(device_flow_flag=False): + tokens = await oidc.device_authorization_grant() + else: + tokens = await oidc.authorization_code_grant() config.token = tokens["access_token"] refresh_token = tokens.get("refresh_token") if refresh_token is not None: diff --git a/python/packages/jumpstarter-cli/jumpstarter_cli/login_test.py b/python/packages/jumpstarter-cli/jumpstarter_cli/login_test.py index 090c27380..4e08ce5ae 100644 --- a/python/packages/jumpstarter-cli/jumpstarter_cli/login_test.py +++ b/python/packages/jumpstarter-cli/jumpstarter_cli/login_test.py @@ -279,3 +279,157 @@ async def authorization_code_grant(self, **kwargs): assert result.exit_code != 0 assert "TLS certificate validation failed" in result.output assert "Traceback" not in result.output + + +def test_login_uses_device_flow_when_flag_is_passed(monkeypatch) -> None: + """When --device-flow is passed, device_authorization_grant is called instead of authorization_code_grant.""" + auth_config = { + "grpcEndpoint": "grpc.example.com:443", + "namespace": "default", + "oidc": [{"issuer": "https://auth.example.com", "clientId": "test-client"}], + } + + async def fake_fetch_auth_config(*args, **kwargs): + return auth_config + + device_flow_called = False + auth_code_called = False + + class FakeOidcConfig: + def __init__(self, *args, **kwargs): + pass + + async def device_authorization_grant(self): + nonlocal device_flow_called + device_flow_called = True + return {"access_token": "test-token"} + + async def authorization_code_grant(self, **kwargs): + nonlocal auth_code_called + auth_code_called = True + return {"access_token": "test-token"} + + monkeypatch.setattr("jumpstarter_cli.login.fetch_auth_config", fake_fetch_auth_config) + monkeypatch.setattr("jumpstarter_cli.login.Config", FakeOidcConfig) + + runner = CliRunner() + runner.invoke( + jmp, + [ + "login", + "test-client@login.example.com", + "--client-config", + "/tmp/nonexistent-client.yaml", + "--nointeractive", + "--unsafe", + "--device-flow", + ], + ) + + assert device_flow_called is True + assert auth_code_called is False + + +def test_login_uses_device_flow_when_env_var_is_set(monkeypatch) -> None: + """When JMP_OIDC_DEVICE_FLOW=1, device_authorization_grant is called automatically.""" + auth_config = { + "grpcEndpoint": "grpc.example.com:443", + "namespace": "default", + "oidc": [{"issuer": "https://auth.example.com", "clientId": "test-client"}], + } + + async def fake_fetch_auth_config(*args, **kwargs): + return auth_config + + device_flow_called = False + auth_code_called = False + + class FakeOidcConfig: + def __init__(self, *args, **kwargs): + pass + + async def device_authorization_grant(self): + nonlocal device_flow_called + device_flow_called = True + return {"access_token": "test-token"} + + async def authorization_code_grant(self, **kwargs): + nonlocal auth_code_called + auth_code_called = True + return {"access_token": "test-token"} + + monkeypatch.setattr("jumpstarter_cli.login.fetch_auth_config", fake_fetch_auth_config) + monkeypatch.setattr("jumpstarter_cli.login.Config", FakeOidcConfig) + monkeypatch.setenv("JMP_OIDC_DEVICE_FLOW", "1") + + runner = CliRunner() + runner.invoke( + jmp, + [ + "login", + "test-client@login.example.com", + "--client-config", + "/tmp/nonexistent-client.yaml", + "--nointeractive", + "--unsafe", + ], + ) + + assert device_flow_called is True + assert auth_code_called is False + + +def test_login_uses_auth_code_flow_without_device_flow_signals(monkeypatch) -> None: + """Without --device-flow or JMP_OIDC_DEVICE_FLOW, authorization_code_grant is used (no regression).""" + auth_config = { + "grpcEndpoint": "grpc.example.com:443", + "namespace": "default", + "oidc": [{"issuer": "https://auth.example.com", "clientId": "test-client"}], + } + + async def fake_fetch_auth_config(*args, **kwargs): + return auth_config + + device_flow_called = False + auth_code_called = False + + class FakeOidcConfig: + def __init__(self, *args, **kwargs): + pass + + async def device_authorization_grant(self): + nonlocal device_flow_called + device_flow_called = True + return {"access_token": "test-token"} + + async def authorization_code_grant(self, **kwargs): + nonlocal auth_code_called + auth_code_called = True + return {"access_token": "test-token"} + + monkeypatch.setattr("jumpstarter_cli.login.fetch_auth_config", fake_fetch_auth_config) + monkeypatch.setattr("jumpstarter_cli.login.Config", FakeOidcConfig) + monkeypatch.delenv("JMP_OIDC_DEVICE_FLOW", raising=False) + + runner = CliRunner() + runner.invoke( + jmp, + [ + "login", + "test-client@login.example.com", + "--client-config", + "/tmp/nonexistent-client.yaml", + "--nointeractive", + "--unsafe", + ], + ) + + assert auth_code_called is True + assert device_flow_called is False + + +def test_env_py_contains_jmp_oidc_device_flow_constant() -> None: + """The JMP_OIDC_DEVICE_FLOW constant must exist in env.py.""" + from jumpstarter.config.env import JMP_OIDC_DEVICE_FLOW + + assert JMP_OIDC_DEVICE_FLOW == "JMP_OIDC_DEVICE_FLOW" diff --git a/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py b/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py index 396d7f66c..c31e7bfc6 100644 --- a/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py +++ b/python/packages/jumpstarter-cli/jumpstarter_cli/shell_test.py @@ -365,11 +365,27 @@ def b64url(data: bytes) -> str: return f"{header}.{payload}.{sig}" +def _make_valid_jwt() -> str: + """Create a JWT with an exp claim in the future (no signature verification needed).""" + + def b64url(data: bytes) -> str: + return base64.urlsafe_b64encode(data).rstrip(b"=").decode() + + header = b64url(json.dumps({"alg": "HS256", "typ": "JWT"}).encode()) + payload = b64url(json.dumps({"exp": int(time.time()) + 3600, "iss": "https://example.com"}).encode()) + sig = b64url(b"fakesig") + return f"{header}.{payload}.{sig}" + + def test_expired_token_triggers_reauth(): config = _DummyConfig() config.token = _make_expired_jwt() - login_mock = Mock() + def _login_side_effect(cfg): + # Simulate successful re-authentication by updating the token + cfg.token = _make_valid_jwt() + + login_mock = Mock(side_effect=_login_side_effect) @handle_exceptions_with_reauthentication(login_mock) def run_shell(): @@ -385,8 +401,11 @@ def run_shell(): None, ) - with pytest.raises(click.ClickException, match="Please try again now"): - run_shell() + with patch( + "jumpstarter_cli.shell._run_shell_with_lease_async", + new=AsyncMock(return_value=0), + ): + run_shell() # Should succeed after retry — no exception raised login_mock.assert_called_once_with(config) diff --git a/python/packages/jumpstarter/jumpstarter/config/env.py b/python/packages/jumpstarter/jumpstarter/config/env.py index 34fba06f0..a99233d82 100644 --- a/python/packages/jumpstarter/jumpstarter/config/env.py +++ b/python/packages/jumpstarter/jumpstarter/config/env.py @@ -14,6 +14,7 @@ JMP_DISABLE_COMPRESSION = "JMP_DISABLE_COMPRESSION" JMP_OIDC_CALLBACK_PORT = "JMP_OIDC_CALLBACK_PORT" +JMP_OIDC_DEVICE_FLOW = "JMP_OIDC_DEVICE_FLOW" JMP_GRPC_INSECURE = "JMP_GRPC_INSECURE" JUMPSTARTER_GRPC_INSECURE = "JUMPSTARTER_GRPC_INSECURE" JMP_GRPC_PASSPHRASE = "JMP_GRPC_PASSPHRASE"