Merge branch 'main' into driver-cli-integration

Author: kirkbrauer

.github/workflows/pytest.yaml

@@ -62,7 +62,10 @@ jobs:
           done
 
       - name: Run pytest
-        run: make test
+        run: |
+            export UV_PYTHON=${{ matrix.python-version }}
+            make test
+
 
   # https://github.com/orgs/community/discussions/26822
   pytest:

packages/jumpstarter-cli/jumpstarter_cli/shell.py

@@ -12,16 +12,25 @@
 
 @click.command("shell")
 @opt_config()
+@click.argument("command", nargs=-1)
 # client specific
 # TODO: warn if these are specified with exporter config
 @click.option("--lease", "lease_name")
 @opt_selector
 @opt_duration_partial(default=timedelta(minutes=30), show_default="00:30:00")
 # end client specific
 @handle_exceptions
-def shell(config, lease_name, selector, duration):
+def shell(config, command: tuple[str, ...], lease_name, selector, duration):
     """
-    Spawns a shell connecting to a local or remote exporter
+    Spawns a shell (or custom command) connecting to a local or remote exporter
+
+    COMMAND is the custom command to run instead of shell.
+
+    Example:
+
+    .. code-block:: bash
+
+        $ jmp shell --exporter foo -- python bar.py
     """
 
     match config:
@@ -31,11 +40,23 @@ def shell(config, lease_name, selector, duration):
             with config.lease(selector=selector, lease_name=lease_name, duration=duration) as lease:
                 with lease.serve_unix() as path:
                     with lease.monitor():
-                        exit_code = launch_shell(path, "remote", config.drivers.allow, config.drivers.unsafe)
+                        exit_code = launch_shell(
+                            path,
+                            "remote",
+                            config.drivers.allow,
+                            config.drivers.unsafe,
+                            command=command,
+                        )
 
             sys.exit(exit_code)
 
         case ExporterConfigV1Alpha1():
             with config.serve_unix() as path:
                 # SAFETY: the exporter config is local thus considered trusted
-                launch_shell(path, "local", allow=[], unsafe=True)
+                launch_shell(
+                    path,
+                    "local",
+                    allow=[],
+                    unsafe=True,
+                    command=command,
+                )

packages/jumpstarter-driver-qemu/jumpstarter_driver_qemu/driver.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import logging
 import os
 import platform
 from collections.abc import AsyncGenerator
@@ -25,6 +26,11 @@
 from jumpstarter.driver import Driver, export
 
 
+class QmpLogFilter(logging.Filter):
+    def filter(self, record):
+        return False
+
+
 @dataclass(kw_only=True)
 class QemuFlasher(FlasherInterface, Driver):
     parent: Qemu
@@ -50,6 +56,10 @@ class QemuPower(PowerInterface, Driver):
 
     @export
     async def on(self) -> None:  # noqa: C901
+        if hasattr(self, "_process"):
+            self.logger.warning("already powered on, ignoring request")
+            return
+
         root = self.parent.validate_partition("root")
         bios = self.parent.validate_partition("bios")
         ovmf_code = self.parent.validate_partition("OVMF_CODE.fd")
@@ -165,6 +175,10 @@ async def on(self) -> None:  # noqa: C901
 
         qmp = QMPClient(self.parent.hostname)
 
+        logging.getLogger(
+            "qemu.qmp.protocol.{}".format(self.parent.hostname),
+        ).addFilter(QmpLogFilter())
+
         with fail_after(10):
             while qmp.runstate != Runstate.RUNNING:
                 try:
@@ -174,6 +188,7 @@ async def on(self) -> None:  # noqa: C901
 
         chardevs = await qmp.execute("query-chardev")
         pty = next(c for c in chardevs if c["label"] == "serial0")["filename"].lstrip("pty:")
+        Path(self.parent._pty).unlink(missing_ok=True)
         Path(self.parent._pty).symlink_to(pty)
 
         await qmp.execute("system_reset")
@@ -188,6 +203,8 @@ def off(self) -> None:
             except TimeoutExpired:
                 self._process.kill()
             del self._process
+        else:
+            self.logger.warning("already powered off, ignoring request")
 
         if hasattr(self, "_cidata"):
             del self._cidata

packages/jumpstarter-driver-shell/jumpstarter_driver_shell/driver.py

@@ -38,9 +38,11 @@ def call_method(self, method: str, env, *args):
             if result.returncode != 0:
                 self.logger.info(f"{method} return code: {result.returncode}")
             if result.stderr != "":
-                self.logger.debug(f"{method} stderr:\n{result.stderr.rstrip('\n')}")
+                stderr = result.stderr.rstrip("\n")
+                self.logger.debug(f"{method} stderr:\n{stderr}")
             if result.stdout != "":
-                self.logger.debug(f"{method} stdout:\n{result.stdout.rstrip('\n')}")
+                stdout = result.stdout.rstrip("\n")
+                self.logger.debug(f"{method} stdout:\n{stdout}")
             return result.stdout, result.stderr, result.returncode
         except subprocess.TimeoutExpired as e:
             self.logger.error(f"Timeout expired while running {method}: {e}")

packages/jumpstarter/jumpstarter/client/lease.py

@@ -163,13 +163,20 @@ async def monitor_async(self, threshold: timedelta = timedelta(minutes=5)):
         async def _monitor():
             while True:
                 lease = await self.get()
+                # TODO: use effective_end_time as the authoritative source for lease end time
                 if lease.effective_begin_time:
                     end_time = lease.effective_begin_time + lease.duration
                     remain = end_time - datetime.now(tz=datetime.now().astimezone().tzinfo)
-                    if remain < threshold:
+                    if remain < timedelta(0):
+                        # lease already expired, stopping monitor
+                        logger.info("Lease {} ended at {}".format(self.name, end_time))
+                        break
+                    elif remain < threshold:
+                        # lease expiring soon, check again on expected expiration time in case it's extended
                         logger.info("Lease {} ending soon in {} at {}".format(self.name, remain, end_time))
                         await sleep(threshold.total_seconds())
                     else:
+                        # lease still active, check again in 5 seconds
                         await sleep(5)
                 else:
                     await sleep(1)

packages/jumpstarter/jumpstarter/common/grpc.py

@@ -1,3 +1,4 @@
+import asyncio
 import base64
 import os
 import socket
@@ -7,11 +8,12 @@
 from urllib.parse import urlparse
 
 import grpc
+from anyio import fail_after
 
 from jumpstarter.common.exceptions import ConfigurationError, ConnectionError
 
 
-def ssl_channel_credentials(target: str, tls_config):
+async def ssl_channel_credentials(target: str, tls_config, timeout=5):
     configure_grpc_env()
     if tls_config.insecure or os.getenv("JUMPSTARTER_GRPC_INSECURE") == "1":
         try:
@@ -21,12 +23,21 @@ def ssl_channel_credentials(target: str, tls_config):
             raise ConfigurationError(f"Failed parsing {target}") from e
 
         try:
-            root_certificates = ssl.get_server_certificate((parsed.hostname, port))
+            with fail_after(timeout):
+                ssl_context = ssl.create_default_context()
+                ssl_context.check_hostname = False
+                ssl_context.verify_mode = ssl.CERT_NONE
+                _, writer = await asyncio.open_connection(parsed.hostname, port, ssl=ssl_context)
+                root_certificates = ""
+                for cert in writer.get_extra_info("ssl_object")._sslobj.get_unverified_chain():
+                    root_certificates += cert.public_bytes()
             return grpc.ssl_channel_credentials(root_certificates=root_certificates.encode())
         except socket.gaierror as e:
             raise ConnectionError(f"Failed resolving {parsed.hostname}") from e
         except ConnectionRefusedError as e:
             raise ConnectionError(f"Failed connecting to {parsed.hostname}:{port}") from e
+        except TimeoutError as e:
+            raise ConnectionError(f"Timeout connecting to {parsed.hostname}:{port}") from e
 
     elif tls_config.ca != "":
         ca_certificate = base64.b64decode(tls_config.ca)

packages/jumpstarter/jumpstarter/common/streams.py

@@ -34,7 +34,7 @@ class StreamRequestMetadata(BaseModel):
 @asynccontextmanager
 async def connect_router_stream(endpoint, token, stream, tls_config, grpc_options):
     credentials = grpc.composite_channel_credentials(
-        ssl_channel_credentials(endpoint, tls_config),
+        await ssl_channel_credentials(endpoint, tls_config),
         grpc.access_token_call_credentials(token),
     )
 

packages/jumpstarter/jumpstarter/common/utils.py

@@ -80,7 +80,14 @@ def env():
 PROMPT_CWD = "\\W"
 
 
-def launch_shell(host: str, context: str, allow: list[str], unsafe: bool) -> int:
+def launch_shell(
+    host: str,
+    context: str,
+    allow: [str],
+    unsafe: bool,
+    *,
+    command: tuple[str, ...] | None = None,
+) -> int:
     """Launch a shell with a custom prompt indicating the exporter type.
 
     Args:
@@ -89,21 +96,21 @@ def launch_shell(host: str, context: str, allow: list[str], unsafe: bool) -> int
         allow: List of allowed drivers
         unsafe: Whether to allow drivers outside of the allow list
     """
-    cmd = [os.environ.get("SHELL", "bash")]
-    if cmd[0].endswith("bash"):
-        cmd.append("--norc")
-        cmd.append("--noprofile")
-
-    process = Popen(
-        cmd,
-        stdin=sys.stdin,
-        stdout=sys.stdout,
-        stderr=sys.stderr,
-        env=os.environ
-        | {
-            JUMPSTARTER_HOST: host,
-            JMP_DRIVERS_ALLOW: "UNSAFE" if unsafe else ",".join(allow),
-            "PS1": f"{ANSI_GRAY}{PROMPT_CWD} {ANSI_YELLOW}⚡{ANSI_WHITE}{context} {ANSI_YELLOW}➤{ANSI_RESET} ",
-        },
-    )
+
+    env = os.environ | {
+        JUMPSTARTER_HOST: host,
+        JMP_DRIVERS_ALLOW: "UNSAFE" if unsafe else ",".join(allow),
+        "PS1": f"{ANSI_GRAY}{PROMPT_CWD} {ANSI_YELLOW}⚡{ANSI_WHITE}{context} {ANSI_YELLOW}➤{ANSI_RESET} ",
+    }
+
+    if command:
+        process = Popen(command, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr, env=env)
+    else:
+        cmd = [os.environ.get("SHELL", "bash")]
+        if cmd[0].endswith("bash"):
+            cmd.append("--norc")
+            cmd.append("--noprofile")
+
+        process = Popen(cmd, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr, env=env)
+
     return process.wait()

packages/jumpstarter/jumpstarter/config/client.py

@@ -54,7 +54,7 @@ class ClientConfigV1Alpha1(BaseModel):
 
     async def channel(self):
         credentials = grpc.composite_channel_credentials(
-            ssl_channel_credentials(self.endpoint, self.tls),
+            await ssl_channel_credentials(self.endpoint, self.tls),
             call_credentials("Client", self.metadata, self.token),
         )
 

packages/jumpstarter/jumpstarter/config/exporter.py

@@ -159,9 +159,9 @@ async def serve(self):
         # dynamic import to avoid circular imports
         from jumpstarter.exporter import Exporter
 
-        def channel_factory():
+        async def channel_factory():
             credentials = grpc.composite_channel_credentials(
-                ssl_channel_credentials(self.endpoint, self.tls),
+                await ssl_channel_credentials(self.endpoint, self.tls),
                 call_credentials("Exporter", self.metadata, self.token),
             )
             return aio_secure_channel(self.endpoint, credentials, self.grpcOptions)