82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
import base64
|
|
import json
|
|
import socket
|
|
from pathlib import Path
|
|
from time import sleep
|
|
|
|
|
|
# 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, socket_file: Path | str) -> None:
|
|
self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
# try to reconnect a couple of times if connection refused
|
|
for _ in range(100):
|
|
try:
|
|
self.sock.connect(str(socket_file))
|
|
return
|
|
except ConnectionRefusedError:
|
|
sleep(0.1)
|
|
self.sock.connect(str(socket_file))
|
|
|
|
def get_response(self) -> dict:
|
|
result = self.sock.recv(9999999)
|
|
return json.loads(result)
|
|
|
|
# only execute, don't wait for response
|
|
def exec_cmd(self, cmd: str) -> None:
|
|
self.sock.send(
|
|
json.dumps(
|
|
{
|
|
"execute": "guest-exec",
|
|
"arguments": {
|
|
"path": "/bin/sh",
|
|
"arg": ["-l", "-c", cmd],
|
|
"capture-output": True,
|
|
},
|
|
}
|
|
).encode("utf-8")
|
|
)
|
|
|
|
# run, wait for result, return exitcode and output
|
|
def run(self, cmd: str, check: bool = False) -> tuple[int, str, str]:
|
|
self.exec_cmd(cmd)
|
|
result_pid = self.get_response()
|
|
pid = result_pid["return"]["pid"]
|
|
# loop until exited=true
|
|
status_payload = json.dumps(
|
|
{
|
|
"execute": "guest-exec-status",
|
|
"arguments": {
|
|
"pid": pid,
|
|
},
|
|
}
|
|
).encode("utf-8")
|
|
while True:
|
|
self.sock.send(status_payload)
|
|
result = self.get_response()
|
|
if "error" in result and result["error"]["desc"].startswith("PID"):
|
|
raise Exception("PID could not be found")
|
|
if result["return"]["exited"]:
|
|
break
|
|
sleep(0.1)
|
|
|
|
exitcode = result["return"]["exitcode"]
|
|
stdout = (
|
|
""
|
|
if "out-data" not in result["return"]
|
|
else base64.b64decode(result["return"]["out-data"]).decode("utf-8")
|
|
)
|
|
stderr = (
|
|
""
|
|
if "err-data" not in result["return"]
|
|
else base64.b64decode(result["return"]["err-data"]).decode("utf-8")
|
|
)
|
|
if check and exitcode != 0:
|
|
raise Exception(
|
|
f"Command on guest failed\nCommand: {cmd}\nExitcode {exitcode}\nStdout: {stdout}\nStderr: {stderr}"
|
|
)
|
|
return exitcode, stdout, stderr
|