100 lines
3.1 KiB
Python
100 lines
3.1 KiB
Python
"""Guest command transport over vsock-compatible JSON protocol."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import socket
|
|
from dataclasses import dataclass
|
|
from typing import Any, Callable, Protocol
|
|
|
|
|
|
class SocketLike(Protocol):
|
|
def settimeout(self, timeout: int) -> None: ...
|
|
|
|
def connect(self, address: Any) -> None: ...
|
|
|
|
def sendall(self, data: bytes) -> None: ...
|
|
|
|
def recv(self, size: int) -> bytes: ...
|
|
|
|
def close(self) -> None: ...
|
|
|
|
|
|
SocketFactory = Callable[[int, int], SocketLike]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class GuestExecResponse:
|
|
stdout: str
|
|
stderr: str
|
|
exit_code: int
|
|
duration_ms: int
|
|
|
|
|
|
class VsockExecClient:
|
|
"""Minimal JSON-over-stream client for a guest exec agent."""
|
|
|
|
def __init__(self, socket_factory: SocketFactory | None = None) -> None:
|
|
self._socket_factory = socket_factory or socket.socket
|
|
|
|
def exec(
|
|
self,
|
|
guest_cid: int,
|
|
port: int,
|
|
command: str,
|
|
timeout_seconds: int,
|
|
*,
|
|
uds_path: str | None = None,
|
|
) -> GuestExecResponse:
|
|
request = {
|
|
"command": command,
|
|
"timeout_seconds": timeout_seconds,
|
|
}
|
|
family = getattr(socket, "AF_VSOCK", None)
|
|
if family is not None:
|
|
sock = self._socket_factory(family, socket.SOCK_STREAM)
|
|
connect_address: Any = (guest_cid, port)
|
|
elif uds_path is not None:
|
|
sock = self._socket_factory(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
connect_address = uds_path
|
|
else:
|
|
raise RuntimeError("vsock sockets are not supported on this host Python runtime")
|
|
try:
|
|
sock.settimeout(timeout_seconds)
|
|
sock.connect(connect_address)
|
|
if family is None:
|
|
sock.sendall(f"CONNECT {port}\n".encode("utf-8"))
|
|
status = self._recv_line(sock)
|
|
if not status.startswith("OK "):
|
|
raise RuntimeError(f"vsock unix bridge rejected port {port}: {status.strip()}")
|
|
sock.sendall((json.dumps(request) + "\n").encode("utf-8"))
|
|
chunks: list[bytes] = []
|
|
while True:
|
|
data = sock.recv(65536)
|
|
if data == b"":
|
|
break
|
|
chunks.append(data)
|
|
finally:
|
|
sock.close()
|
|
|
|
payload = json.loads(b"".join(chunks).decode("utf-8"))
|
|
if not isinstance(payload, dict):
|
|
raise RuntimeError("guest exec response must be a JSON object")
|
|
return GuestExecResponse(
|
|
stdout=str(payload.get("stdout", "")),
|
|
stderr=str(payload.get("stderr", "")),
|
|
exit_code=int(payload.get("exit_code", -1)),
|
|
duration_ms=int(payload.get("duration_ms", 0)),
|
|
)
|
|
|
|
@staticmethod
|
|
def _recv_line(sock: SocketLike) -> str:
|
|
chunks: list[bytes] = []
|
|
while True:
|
|
data = sock.recv(1)
|
|
if data == b"":
|
|
break
|
|
chunks.append(data)
|
|
if data == b"\n":
|
|
break
|
|
return b"".join(chunks).decode("utf-8", errors="replace")
|