Coverage for /Users/ajo/work/jumpstarter/jumpstarter/packages/jumpstarter/jumpstarter/common/utils.py: 80%
44 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:21 +0200
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:21 +0200
1import os
2import sys
3from contextlib import ExitStack, asynccontextmanager, contextmanager
4from subprocess import Popen
6from anyio.from_thread import BlockingPortal, start_blocking_portal
8from jumpstarter.client import client_from_path
9from jumpstarter.config.env import JMP_DRIVERS_ALLOW, JUMPSTARTER_HOST
10from jumpstarter.driver import Driver
11from jumpstarter.exporter import Session
12from jumpstarter.utils.env import env
14__all__ = ["env"]
17@asynccontextmanager
18async def serve_async(root_device: Driver, portal: BlockingPortal, stack: ExitStack):
19 with Session(root_device=root_device) as session:
20 async with session.serve_unix_async() as path:
21 # SAFETY: the root_device instance is constructed locally thus considered trusted
22 async with client_from_path(path, portal, stack, allow=[], unsafe=True) as client:
23 try:
24 yield client
25 finally:
26 if hasattr(client, "close"):
27 client.close()
30@contextmanager
31def serve(root_device: Driver):
32 with start_blocking_portal() as portal:
33 with ExitStack() as stack:
34 with portal.wrap_async_context_manager(serve_async(root_device, portal, stack)) as client:
35 try:
36 yield client
37 finally:
38 if hasattr(client, "close"):
39 client.close()
42ANSI_GRAY = "\\[\\e[90m\\]"
43ANSI_YELLOW = "\\[\\e[93m\\]"
44ANSI_WHITE = "\\[\\e[97m\\]"
45ANSI_RESET = "\\[\\e[0m\\]"
46PROMPT_CWD = "\\W"
49def launch_shell(
50 host: str,
51 context: str,
52 allow: [str],
53 unsafe: bool,
54 *,
55 command: tuple[str, ...] | None = None,
56) -> int:
57 """Launch a shell with a custom prompt indicating the exporter type.
59 Args:
60 host: The jumpstarter host path
61 context: The context of the shell ("local" or "remote")
62 allow: List of allowed drivers
63 unsafe: Whether to allow drivers outside of the allow list
64 """
66 env = os.environ | {
67 JUMPSTARTER_HOST: host,
68 JMP_DRIVERS_ALLOW: "UNSAFE" if unsafe else ",".join(allow),
69 "PS1": f"{ANSI_GRAY}{PROMPT_CWD} {ANSI_YELLOW}⚡{ANSI_WHITE}{context} {ANSI_YELLOW}➤{ANSI_RESET} ",
70 }
72 if command:
73 process = Popen(command, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr, env=env)
74 else:
75 cmd = [os.environ.get("SHELL", "bash")]
76 if cmd[0].endswith("bash"):
77 cmd.append("--norc")
78 cmd.append("--noprofile")
80 process = Popen(cmd, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr, env=env)
82 return process.wait()