83 lines
2.8 KiB
Python
83 lines
2.8 KiB
Python
import base64
|
|
import time
|
|
import types
|
|
from dataclasses import dataclass
|
|
|
|
from clan_cli.errors import ClanError
|
|
from clan_cli.qemu.qmp import QEMUMonitorProtocol
|
|
|
|
|
|
@dataclass
|
|
class VmCommandResult:
|
|
returncode: int
|
|
stdout: str | None
|
|
stderr: str | None
|
|
|
|
|
|
# qga is almost like qmp, but not quite, because:
|
|
# - server doesn't send initial message
|
|
# - no need to initialize by asking for capabilities
|
|
# - results need to be base64 decoded
|
|
class QgaSession:
|
|
def __init__(self, address: str) -> None:
|
|
self.client = QEMUMonitorProtocol(address)
|
|
self.client.connect(negotiate=False)
|
|
|
|
def __enter__(self) -> "QgaSession":
|
|
# Implement context manager enter function.
|
|
return self
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc_value: BaseException | None,
|
|
traceback: types.TracebackType | None,
|
|
) -> None:
|
|
# Implement context manager exit function.
|
|
self.client.close()
|
|
|
|
def run_nonblocking(self, cmd: list[str]) -> int:
|
|
result_pid = self.client.cmd(
|
|
"guest-exec", {"path": cmd[0], "arg": cmd[1:], "capture-output": True}
|
|
)
|
|
if result_pid is None:
|
|
msg = "Could not get PID from QGA"
|
|
raise ClanError(msg)
|
|
try:
|
|
return result_pid["return"]["pid"]
|
|
except KeyError as e:
|
|
if "error" in result_pid:
|
|
msg = f"Could not run command: {result_pid['error']['desc']}"
|
|
raise ClanError(msg) from e
|
|
msg = f"PID could not be found: {result_pid}"
|
|
raise ClanError(msg) from e
|
|
|
|
# run, wait for result, return exitcode and output
|
|
def run(self, cmd: list[str], check: bool = True) -> VmCommandResult:
|
|
pid = self.run_nonblocking(cmd)
|
|
# loop until exited=true
|
|
while True:
|
|
result = self.client.cmd("guest-exec-status", {"pid": pid})
|
|
if result is None:
|
|
msg = "Could not get status from QGA"
|
|
raise ClanError(msg)
|
|
if "error" in result and result["error"]["desc"].startswith("PID"):
|
|
msg = "PID could not be found"
|
|
raise ClanError(msg)
|
|
if result["return"]["exited"]:
|
|
break
|
|
time.sleep(0.1)
|
|
|
|
exitcode = result["return"]["exitcode"]
|
|
err_data = result["return"].get("err-data")
|
|
stdout = None
|
|
stderr = None
|
|
if out_data := result["return"].get("out-data"):
|
|
stdout = base64.b64decode(out_data).decode("utf-8")
|
|
if err_data is not None:
|
|
stderr = base64.b64decode(err_data).decode("utf-8")
|
|
if check and exitcode != 0:
|
|
msg = f"Command on guest failed\nCommand: {cmd}\nExitcode {exitcode}\nStdout: {stdout}\nStderr: {stderr}"
|
|
raise ClanError(msg)
|
|
return VmCommandResult(exitcode, stdout, stderr)
|