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

1import os 

2import sys 

3from contextlib import ExitStack, asynccontextmanager, contextmanager 

4from subprocess import Popen 

5 

6from anyio.from_thread import BlockingPortal, start_blocking_portal 

7 

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 

13 

14__all__ = ["env"] 

15 

16 

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() 

28 

29 

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() 

40 

41 

42ANSI_GRAY = "\\[\\e[90m\\]" 

43ANSI_YELLOW = "\\[\\e[93m\\]" 

44ANSI_WHITE = "\\[\\e[97m\\]" 

45ANSI_RESET = "\\[\\e[0m\\]" 

46PROMPT_CWD = "\\W" 

47 

48 

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. 

58 

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 """ 

65 

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 } 

71 

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") 

79 

80 process = Popen(cmd, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr, env=env) 

81 

82 return process.wait()