mirror of
https://github.com/LGUG2Z/komorebi.git
synced 2026-11-15 05:12:28 +01:00
LGUG2Z
80f221b8e7
The host side of the agent spoke a protocol that is the same for every platform, and lived in the Windows port. Move it to komorebi_vm/agent.py, which the Windows module re-exports so that its imports stay as they are, and give it the calls the protocol gained: spawn with an environment and an output file, `status` of a spawned program and `send` to its standard input.
120 lines
4.3 KiB
Python
120 lines
4.3 KiB
Python
"""The host side of `komorebi-e2e-agent`: JSON lines over TCP, the same for every platform."""
|
|
|
|
import json
|
|
import socket
|
|
import time
|
|
from itertools import count
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
class AgentError(Exception):
|
|
pass
|
|
|
|
|
|
class AgentLost(AgentError):
|
|
"""The connection dropped during a call: the call is not repeated, as it
|
|
may have run, and the caller names it, so that a dead agent is not taken
|
|
for the command having broken what it was testing."""
|
|
|
|
|
|
class Agent:
|
|
def __init__(self, host: str = "127.0.0.1", port: int = 7878) -> None:
|
|
self._addr = (host, port)
|
|
self._ids = count(1)
|
|
self._sock: socket.socket | None = None
|
|
self._file: Any = None
|
|
|
|
def connect(self, timeout: float = 120) -> None:
|
|
"""Waits for the guest to boot and the agent to listen."""
|
|
deadline = time.monotonic() + timeout
|
|
while True:
|
|
try:
|
|
self._sock = socket.create_connection(self._addr, timeout=5)
|
|
self._sock.settimeout(60)
|
|
self._file = self._sock.makefile("rwb")
|
|
if self.call("ping") == "pong":
|
|
return
|
|
except (OSError, ValueError, AgentLost):
|
|
# QEMU accepts the forwarded connection before the guest
|
|
# listens, then closes it: an empty answer
|
|
if self._sock:
|
|
self._sock.close()
|
|
if time.monotonic() > deadline:
|
|
raise AgentError(f"no agent at {self._addr} after {timeout}s")
|
|
time.sleep(1)
|
|
|
|
def call(self, verb: str, **fields: Any) -> Any:
|
|
request_id = next(self._ids)
|
|
line = json.dumps({"id": request_id, "verb": verb, **fields}) + "\n"
|
|
try:
|
|
self._file.write(line.encode())
|
|
self._file.flush()
|
|
return self._read_reply(request_id)
|
|
except (OSError, ValueError) as e:
|
|
what = " ".join([verb, *map(str, fields.get("args", []))])
|
|
raise AgentLost(f"agent connection lost during {what}: {e}") from e
|
|
|
|
def _read_reply(self, request_id: int) -> Any:
|
|
reply = json.loads(self._file.readline())
|
|
if reply.get("id") != request_id:
|
|
got = reply.get("id")
|
|
raise AgentError(f"answer to {got}, expected {request_id}")
|
|
if "error" in reply:
|
|
raise AgentError(reply["error"])
|
|
return reply["ok"]
|
|
|
|
def exec(self, *args: str) -> dict[str, Any]:
|
|
return self.call("exec", args=list(args))
|
|
|
|
def spawn(
|
|
self, *args: str, env: dict[str, str] | None = None, out: str | None = None
|
|
) -> int:
|
|
"""Starts a program that keeps running and returns its process id. `env` is added to the
|
|
environment it starts with, and `out` is a file in the guest that gets what it prints."""
|
|
fields: dict[str, Any] = {"args": list(args)}
|
|
if env:
|
|
fields["env"] = env
|
|
if out:
|
|
fields["out"] = out
|
|
return self.call("spawn", **fields)["pid"]
|
|
|
|
def status(self, pid: int) -> dict[str, Any]:
|
|
"""`{"running": bool, "code": int | None}` of a program `spawn` started."""
|
|
return self.call("status", pid=pid)
|
|
|
|
def send(self, pid: int, line: str) -> None:
|
|
"""Writes a line to the standard input of a program `spawn` started."""
|
|
self.call("send", pid=pid, line=line)
|
|
|
|
def os_windows(self) -> list[dict[str, Any]]:
|
|
return self.call("os_windows")
|
|
|
|
def displays(self) -> list[dict[str, Any]]:
|
|
return self.call("displays")
|
|
|
|
def put_file(self, local: Path, remote: str) -> None:
|
|
data = local.read_bytes()
|
|
request_id = next(self._ids)
|
|
header = {
|
|
"id": request_id,
|
|
"verb": "put_file",
|
|
"path": remote,
|
|
"size": len(data),
|
|
}
|
|
self._file.write((json.dumps(header) + "\n").encode() + data)
|
|
self._file.flush()
|
|
self._read_reply(request_id)
|
|
|
|
def get_file(self, remote: str) -> bytes:
|
|
request_id = next(self._ids)
|
|
header = {"id": request_id, "verb": "get_file", "path": remote}
|
|
self._file.write((json.dumps(header) + "\n").encode())
|
|
self._file.flush()
|
|
size = self._read_reply(request_id)["size"]
|
|
return self._file.read(size)
|
|
|
|
def close(self) -> None:
|
|
if self._sock:
|
|
self._sock.close()
|