Fix OIDC re-login

packages/jumpstarter-cli-common/jumpstarter_cli_common/exceptions.py

@@ -4,7 +4,7 @@
 
 import click
 
-from jumpstarter.common.exceptions import JumpstarterException
+from jumpstarter.common.exceptions import ConnectionError, JumpstarterException
 
 
 class ClickExceptionRed(click.ClickException):
@@ -46,6 +46,32 @@ def wrapped(*args, **kwargs):
     return wrapped
 
 
+def handle_exceptions_with_reauthentication(login_func):
+    """Decorator to handle exceptions in blocking functions."""
+    def decorator(func):
+        @wraps(func)
+        def wrapped(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except ConnectionError as e:
+                if "expired" in str(e).lower():
+                    click.echo(click.style("Token is expired, triggering re-authentication", fg="red"))
+                    config = e.get_config()
+                    login_func(config)
+                    raise ClickExceptionRed("Please try again now") from None
+                else:
+                    raise ClickExceptionRed(str(e)) from None
+            except JumpstarterException as e:
+                raise ClickExceptionRed(str(e)) from None
+            except click.ClickException:
+                raise  # if it was already a click exception from the cli commands, just re-raise it
+            except Exception:
+                raise
+        return wrapped
+
+    return decorator
+
+
 # https://peps.python.org/pep-0785/#reference-implementation
 def leaf_exceptions(self: BaseExceptionGroup, *, fix_tracebacks: bool = True) -> list[BaseException]:
     """

packages/jumpstarter-cli/jumpstarter_cli/create.py

@@ -2,11 +2,12 @@
 
 import click
 from jumpstarter_cli_common.config import opt_config
-from jumpstarter_cli_common.exceptions import handle_exceptions
+from jumpstarter_cli_common.exceptions import handle_exceptions_with_reauthentication
 from jumpstarter_cli_common.opt import OutputType, opt_output_all
 from jumpstarter_cli_common.print import model_print
 
 from .common import opt_duration_partial, opt_selector
+from .login import relogin_client
 
 
 @click.group()
@@ -21,7 +22,7 @@ def create():
 @opt_selector
 @opt_duration_partial(required=True)
 @opt_output_all
-@handle_exceptions
+@handle_exceptions_with_reauthentication(relogin_client)
 def create_lease(config, selector: str, duration: timedelta, output: OutputType):
     """
     Create a lease

packages/jumpstarter-cli/jumpstarter_cli/delete.py

@@ -1,9 +1,10 @@
 import click
 from jumpstarter_cli_common.config import opt_config
-from jumpstarter_cli_common.exceptions import handle_exceptions
+from jumpstarter_cli_common.exceptions import handle_exceptions_with_reauthentication
 from jumpstarter_cli_common.opt import OutputMode, OutputType, opt_output_name_only
 
 from .common import opt_selector
+from .login import relogin_client
 
 
 @click.group()
@@ -19,7 +20,7 @@ def delete():
 @opt_selector
 @click.option("--all", "all", is_flag=True)
 @opt_output_name_only
-@handle_exceptions
+@handle_exceptions_with_reauthentication(relogin_client)
 def delete_leases(config, name: str, selector: str | None, all: bool, output: OutputType):
     """
     Delete leases

packages/jumpstarter-cli/jumpstarter_cli/get.py

@@ -1,10 +1,11 @@
 import click
 from jumpstarter_cli_common.config import opt_config
-from jumpstarter_cli_common.exceptions import handle_exceptions
+from jumpstarter_cli_common.exceptions import handle_exceptions_with_reauthentication
 from jumpstarter_cli_common.opt import OutputType, opt_output_all
 from jumpstarter_cli_common.print import model_print
 
 from .common import opt_selector
+from .login import relogin_client
 
 
 @click.group()
@@ -19,7 +20,7 @@ def get():
 @opt_selector
 @opt_output_all
 @click.option("--with", "with_options", multiple=True, help="Include additional information (e.g., 'leases')")
-@handle_exceptions
+@handle_exceptions_with_reauthentication(relogin_client)
 def get_exporters(config, selector: str | None, output: OutputType, with_options: tuple[str, ...]):
     """
     Display one or many exporters
@@ -35,7 +36,7 @@ def get_exporters(config, selector: str | None, output: OutputType, with_options
 @opt_config(exporter=False)
 @opt_selector
 @opt_output_all
-@handle_exceptions
+@handle_exceptions_with_reauthentication(relogin_client)
 def get_leases(config, selector: str | None, output: OutputType):
     """
     Display one or many leases

packages/jumpstarter-cli/jumpstarter_cli/login.py

@@ -2,11 +2,17 @@
 from jumpstarter_cli_common.blocking import blocking
 from jumpstarter_cli_common.config import opt_config
 from jumpstarter_cli_common.oidc import Config, decode_jwt_issuer, opt_oidc
-from jumpstarter_cli_common.opt import confirm_insecure_tls, opt_insecure_tls_config, opt_nointeractive
+from jumpstarter_cli_common.opt import (
+    confirm_insecure_tls,
+    opt_insecure_tls_config,
+    opt_nointeractive,
+)
 
+from jumpstarter.common.exceptions import ReauthenticationFailed
 from jumpstarter.config.client import ClientConfigV1Alpha1, ClientConfigV1Alpha1Drivers
 from jumpstarter.config.common import ObjectMeta
 from jumpstarter.config.exporter import ExporterConfigV1Alpha1
+from jumpstarter.config.tls import TLSConfigV1Alpha1
 
 
 @click.command("login", short_help="Login")
@@ -50,12 +56,18 @@ async def login(  # noqa: C901
 
     confirm_insecure_tls(insecure_tls_config, nointeractive)
 
+    config_kind = None
     match config:
+        # we are updating an existing config
         case ClientConfigV1Alpha1():
             issuer = decode_jwt_issuer(config.token)
+            config_kind = "client"
         case ExporterConfigV1Alpha1():
             issuer = decode_jwt_issuer(config.token)
+            config_kind = "exporter"
+        # we are creating a new config
         case (kind, value):
+            config_kind = kind
             if namespace is None:
                 if nointeractive:
                     raise click.UsageError("Namespace is required in non-interactive mode.")
@@ -83,6 +95,7 @@ async def login(  # noqa: C901
                 config = ClientConfigV1Alpha1(
                     alias=value if kind == "client" else "default",
                     metadata=ObjectMeta(namespace=namespace, name=name),
+                    tls=TLSConfigV1Alpha1(insecure=insecure_tls_config),
                     endpoint=endpoint,
                     token="",
                     drivers=ClientConfigV1Alpha1Drivers(allow=allow.split(","), unsafe=unsafe),
@@ -91,6 +104,7 @@ async def login(  # noqa: C901
             if kind.startswith("exporter"):
                 config = ExporterConfigV1Alpha1(
                     alias=value if kind == "exporter" else "default",
+                    tls=TLSConfigV1Alpha1(insecure=insecure_tls_config),
                     metadata=ObjectMeta(namespace=namespace, name=name),
                     endpoint=endpoint,
                     token="",
@@ -112,9 +126,8 @@ async def login(  # noqa: C901
         tokens = await oidc.authorization_code_grant()
 
     config.token = tokens["access_token"]
-    config.tls.insecure = insecure_tls_config
 
-    match kind:
+    match config_kind:
         case "client":
             ClientConfigV1Alpha1.save(config)  # ty: ignore[invalid-argument-type]
         case "client_config":
@@ -123,3 +136,21 @@ async def login(  # noqa: C901
             ExporterConfigV1Alpha1.save(config)  # ty: ignore[invalid-argument-type]
         case "exporter_config":
             ExporterConfigV1Alpha1.save(config, value)  # ty: ignore[invalid-argument-type]
+
+@blocking
+async def relogin_client(config: ClientConfigV1Alpha1):
+    """Relogin into a jumpstarter instance"""
+    client_id = "jumpstarter-cli" # TODO: store this metadata in the config
+    try:
+        issuer = decode_jwt_issuer(config.token)
+    except Exception as e:
+        raise ReauthenticationFailed(f"Failed to decode JWT issuer: {e}") from e
+
+    try:
+        oidc = Config(issuer=issuer, client_id=client_id)
+        tokens = await oidc.authorization_code_grant()
+        config.token = tokens["access_token"]
+        ClientConfigV1Alpha1.save(config) # ty: ignore[invalid-argument-type]
+    except Exception as e:
+        raise ReauthenticationFailed(f"Failed to re-authenticate: {e}") from e
+

packages/jumpstarter-cli/jumpstarter_cli/shell.py

@@ -3,9 +3,10 @@
 
 import click
 from jumpstarter_cli_common.config import opt_config
-from jumpstarter_cli_common.exceptions import handle_exceptions
+from jumpstarter_cli_common.exceptions import handle_exceptions_with_reauthentication
 
 from .common import opt_duration_partial, opt_selector
+from .login import relogin_client
 from jumpstarter.common.utils import launch_shell
 from jumpstarter.config.client import ClientConfigV1Alpha1
 from jumpstarter.config.exporter import ExporterConfigV1Alpha1
@@ -20,7 +21,7 @@
 @opt_selector
 @opt_duration_partial(default=timedelta(minutes=30), show_default="00:30:00")
 # end client specific
-@handle_exceptions
+@handle_exceptions_with_reauthentication(relogin_client)
 def shell(config, command: tuple[str, ...], lease_name, selector, duration):
     """
     Spawns a shell (or custom command) connecting to a local or remote exporter

packages/jumpstarter-cli/jumpstarter_cli/update.py

@@ -2,11 +2,12 @@
 
 import click
 from jumpstarter_cli_common.config import opt_config
-from jumpstarter_cli_common.exceptions import handle_exceptions
+from jumpstarter_cli_common.exceptions import handle_exceptions_with_reauthentication
 from jumpstarter_cli_common.opt import OutputType, opt_output_all
 from jumpstarter_cli_common.print import model_print
 
 from .common import opt_duration_partial
+from .login import relogin_client
 
 
 @click.group()
@@ -21,7 +22,7 @@ def update():
 @click.argument("name")
 @opt_duration_partial(required=True)
 @opt_output_all
-@handle_exceptions
+@handle_exceptions_with_reauthentication(relogin_client)
 def update_lease(config, name: str, duration: timedelta, output: OutputType):
     """
     Update a lease

packages/jumpstarter/jumpstarter/common/exceptions.py

@@ -16,12 +16,23 @@ class for all jumpstarter-specific errors.
     def __init__(self, message: str):
         super().__init__(message)
         self.message = message
+        self._config = None
 
     def __str__(self):
         if self.__cause__:
             return f"{self.message} (Caused by: {self.__cause__})"
         return f"{self.message}"
 
+
+    # some exceptions need to able to set the config that caused the error
+    # to attempt recovery, or re-authentication if the token is expired
+    def set_config(self, config):
+        self._config = config
+
+    def get_config(self):
+        return self._config
+
+
     def print(self, message: str | None = None):
         ANSI_RED = "\033[91m"
         ANSI_CLEAR = "\033[0m"
@@ -56,3 +67,9 @@ class FileNotFoundError(JumpstarterException, FileNotFoundError):
     """Raised when a file is not found."""
 
     pass
+
+
+class ReauthenticationFailed(JumpstarterException):
+    """Raised when a re-authentication fails."""
+
+    pass