74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Minimal guest-side exec agent for pyro runtime bundles."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import socket
|
|
import subprocess
|
|
import time
|
|
from typing import Any
|
|
|
|
PORT = 5005
|
|
BUFFER_SIZE = 65536
|
|
|
|
|
|
def _read_request(conn: socket.socket) -> dict[str, Any]:
|
|
chunks: list[bytes] = []
|
|
while True:
|
|
data = conn.recv(BUFFER_SIZE)
|
|
if data == b"":
|
|
break
|
|
chunks.append(data)
|
|
if b"\n" in data:
|
|
break
|
|
payload = json.loads(b"".join(chunks).decode("utf-8").strip())
|
|
if not isinstance(payload, dict):
|
|
raise RuntimeError("request must be a JSON object")
|
|
return payload
|
|
|
|
|
|
def _run_command(command: str, timeout_seconds: int) -> dict[str, Any]:
|
|
started = time.monotonic()
|
|
try:
|
|
proc = subprocess.run(
|
|
["/bin/sh", "-lc", command],
|
|
text=True,
|
|
capture_output=True,
|
|
timeout=timeout_seconds,
|
|
check=False,
|
|
)
|
|
return {
|
|
"stdout": proc.stdout,
|
|
"stderr": proc.stderr,
|
|
"exit_code": proc.returncode,
|
|
"duration_ms": int((time.monotonic() - started) * 1000),
|
|
}
|
|
except subprocess.TimeoutExpired:
|
|
return {
|
|
"stdout": "",
|
|
"stderr": f"command timed out after {timeout_seconds}s",
|
|
"exit_code": 124,
|
|
"duration_ms": int((time.monotonic() - started) * 1000),
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
family = getattr(socket, "AF_VSOCK", None)
|
|
if family is None:
|
|
raise SystemExit("AF_VSOCK is unavailable")
|
|
with socket.socket(family, socket.SOCK_STREAM) as server:
|
|
server.bind((socket.VMADDR_CID_ANY, PORT))
|
|
server.listen(1)
|
|
while True:
|
|
conn, _ = server.accept()
|
|
with conn:
|
|
request = _read_request(conn)
|
|
command = str(request.get("command", ""))
|
|
timeout_seconds = int(request.get("timeout_seconds", 30))
|
|
response = _run_command(command, timeout_seconds)
|
|
conn.sendall((json.dumps(response) + "\n").encode("utf-8"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|