1#!/usr/bin/env python3 2# Copyright 2021 The ChromiumOS Authors 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6import shutil 7from typing import Iterable, Optional 8 9from impl.common import run_commands, argh, console, strip_ansi_escape_sequences 10from impl import testvm 11from impl.testvm import Arch, VmState 12 13USAGE = """Manages VMs for testing crosvm. 14 15Can run an x86_64 and an aarch64 vm via `./tools/x86vm` and `./tools/aarch64vm`. 16Both are a wrapper around `./tools/testvm --arch=x86_64/aarch64`. 17 18The easiest way to use the VM is: 19 20 $ ./tools/aarch64vm ssh 21 22Which will initialize and boot the VM, then wait for SSH to be available and 23opens an SSH session. The VM will stay alive between calls. 24 25Available commands are: 26 - up: Start the VM if it is not already running. 27 - stop: Gracefully stop the VM 28 - kill: Send SIGKILL to the VM 29 - clean: Stop the VM and delete all images 30 - logs: Print logs of the VM console 31 32All of these can be called on `./tools/x86vm` or `./tools/aarch64vm`, but also on 33`tools/testvm` to apply to both VMs. 34""" 35 36 37def cli_shorthand(arch: Arch): 38 if arch == "x86_64": 39 return "tools/x86vm" 40 elif arch == "aarch64": 41 return "tools/aarch64vm" 42 else: 43 raise Exception(f"Unknown architecture: {arch}") 44 45 46def arch_or_all(arch: Optional[Arch]): 47 return (arch,) if arch else testvm.ARCH_OPTIONS 48 49 50ARCHS = testvm.ARCH_OPTIONS 51 52 53@argh.arg("--arch-list", "--arch", nargs="*", type=str, default=ARCHS, choices=ARCHS) 54def up(arch_list: Iterable[Arch] = [], reset: bool = False, wait: bool = False, timeout: int = 120): 55 "Start the VM if it's not already running." 56 for arch in arch_list: 57 testvm.up(arch, reset, wait, timeout) 58 59 60@argh.arg("--arch", required=True, choices=testvm.ARCH_OPTIONS) 61def run(arch: Arch = "x86_64", reset: bool = False): 62 "Run the VM in foreground for debugging purposes." 63 if testvm.is_running(arch): 64 raise Exception("VM is already running") 65 testvm.build_if_needed(arch, reset) 66 testvm.run_qemu( 67 arch, 68 testvm.rootfs_img_path(arch), 69 background=False, 70 ) 71 72 73@argh.arg("--arch", required=True, choices=testvm.ARCH_OPTIONS) 74def shell(arch: Arch = "x86_64", timeout: int = 120): 75 "Starts an interactive shell via SSH, will ensure the VM is running." 76 testvm.up(arch, wait=True, timeout=timeout) 77 testvm.ssh_exec(arch) 78 79 80@argh.arg("--arch-list", "--arch", nargs="*", type=str, default=ARCHS, choices=ARCHS) 81def stop(arch_list: Iterable[Arch] = []): 82 "Gracefully stops the running VM." 83 for arch in arch_list: 84 if not testvm.is_running(arch): 85 print(f"{arch} VM is not running") 86 break 87 console.print(f"Stopping {arch} VM") 88 testvm.ssh_exec(arch, "sudo poweroff") 89 90 91@argh.arg("--arch-list", "--arch", nargs="*", type=str, default=ARCHS, choices=ARCHS) 92def kill(arch_list: Iterable[Arch] = []): 93 "Kills the running VM with SIGKILL." 94 for arch in arch_list: 95 if not testvm.is_running(arch): 96 console.print(f"{arch} VM is not running") 97 break 98 console.print(f"Killing {arch} VM process") 99 testvm.kill_vm(arch) 100 print() 101 102 103@argh.arg("--arch-list", "--arch", nargs="*", type=str, default=ARCHS, choices=ARCHS) 104def clean(arch_list: Iterable[Arch] = []): 105 "Stops the VM or VMs and deletes all data." 106 for arch in arch_list: 107 if testvm.is_running(arch): 108 kill(arch) 109 if testvm.data_dir(arch).exists(): 110 console.print("Cleaning data directory", testvm.data_dir(arch)) 111 shutil.rmtree(testvm.data_dir(arch)) 112 print() 113 114 115def vm_status(arch: Arch): 116 def cli_tip(*args: str): 117 return f"[green][bold]{cli_shorthand(arch)} {' '.join(args)}[/bold][/green]" 118 119 vm = f"{arch} VM" 120 port = f"[blue]{testvm.SSH_PORTS[arch]}[/blue]" 121 122 state = testvm.state(arch) 123 if state == VmState.REACHABLE: 124 console.print(f"{vm} is [green]reachable[/green] on port {port}") 125 console.print(f"Start a shell with {cli_tip('shell')}") 126 elif state == VmState.STOPPED: 127 console.print(f"{vm} is [red]stopped[/red]") 128 console.print(f"Start the VM with {cli_tip('up')}") 129 else: 130 console.print(f"{vm} is running but [red]not reachable[/red] on port {port}") 131 console.print(f"Recent logs:") 132 logs(arch, 10, style="light_slate_grey") 133 console.print(f"See all logs with {cli_tip('logs')}") 134 135 136@argh.arg("--arch-list", "--arch", nargs="*", type=str, default=ARCHS, choices=ARCHS) 137def status(arch_list: Iterable[Arch] = []): 138 for arch in arch_list: 139 vm_status(arch) 140 print() 141 142 143@argh.arg("--arch", required=True, choices=testvm.ARCH_OPTIONS) 144def logs(arch: Arch = "x86_64", n: int = 0, style: Optional[str] = None): 145 log_lines = testvm.log_path(arch).read_text().splitlines() 146 if n > 0 and len(log_lines) > n: 147 log_lines = log_lines[-n:] 148 for line in log_lines: 149 if style: 150 console.print( 151 strip_ansi_escape_sequences(line), style=style, markup=False, highlight=False 152 ) 153 else: 154 print(line) 155 156 157if __name__ == "__main__": 158 run_commands(up, run, shell, stop, kill, clean, status, logs, usage=USAGE, default_fn=status) 159