114 lines
4.5 KiB
Python
114 lines
4.5 KiB
Python
"""Bounded JSON QMP over private Unix sockets, with reply ID matching."""
|
|
|
|
import json
|
|
import socket
|
|
import time
|
|
from pathlib import Path
|
|
|
|
|
|
class QMPError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class QMPClient:
|
|
def __init__(self, path: Path, timeout: float = 2.0):
|
|
if timeout <= 0:
|
|
raise ValueError('QMP timeout must be positive')
|
|
self.timeout = timeout
|
|
self.socket = socket.socket(socket.AF_UNIX)
|
|
self.buffer = bytearray()
|
|
self.counter = 0
|
|
deadline = time.monotonic() + timeout
|
|
try:
|
|
self.socket.settimeout(timeout)
|
|
self.socket.connect(str(path))
|
|
greeting = self._receive(deadline)
|
|
# Real QEMU can queue a SPICE event across monitor reconnects.
|
|
# Consume only protocol events, retaining the original deadline.
|
|
while 'event' in greeting:
|
|
if not self._async_event(greeting):
|
|
raise QMPError('Invalid QMP event before greeting')
|
|
greeting = self._receive(deadline)
|
|
if not isinstance(greeting.get('QMP'), dict):
|
|
raise QMPError('Invalid QMP greeting')
|
|
self._command('qmp_capabilities', None, deadline)
|
|
except (OSError, QMPError) as exc:
|
|
self.close()
|
|
if isinstance(exc, QMPError):
|
|
raise
|
|
raise QMPError(f'QMP connection failed: {exc}') from exc
|
|
|
|
@staticmethod
|
|
def _async_event(message: dict) -> bool:
|
|
if not {'event', 'timestamp'} <= message.keys() <= {'event', 'timestamp', 'data'}:
|
|
return False
|
|
if not isinstance(message['event'], str) or not message['event']:
|
|
return False
|
|
if 'data' in message and not isinstance(message['data'], dict):
|
|
return False
|
|
stamp = message['timestamp']
|
|
if (not isinstance(stamp, dict) or set(stamp) != {'seconds', 'microseconds'}
|
|
or any(type(value) is not int for value in stamp.values())):
|
|
return False
|
|
seconds, micros = stamp['seconds'], stamp['microseconds']
|
|
# QMP specifies (-1, -1) when the server cannot retrieve host time.
|
|
return (seconds == -1 and micros == -1) or (seconds >= 0 and 0 <= micros < 1000000)
|
|
|
|
def _receive(self, deadline: float) -> dict:
|
|
while True:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
raise QMPError('QMP deadline exceeded')
|
|
if b'\n' in self.buffer:
|
|
line, _, rest = self.buffer.partition(b'\n')
|
|
self.buffer = bytearray(rest)
|
|
try:
|
|
message = json.loads(line)
|
|
except (ValueError, UnicodeDecodeError) as exc:
|
|
raise QMPError('Malformed QMP JSON') from exc
|
|
if not isinstance(message, dict):
|
|
raise QMPError('QMP response must be an object')
|
|
return message
|
|
self.socket.settimeout(remaining)
|
|
try:
|
|
chunk = self.socket.recv(65536)
|
|
except OSError as exc:
|
|
raise QMPError(f'QMP receive failed: {exc}') from exc
|
|
if not chunk:
|
|
raise QMPError('QMP connection closed')
|
|
self.buffer.extend(chunk)
|
|
if len(self.buffer) > 1048576:
|
|
raise QMPError('QMP message exceeds 1 MiB')
|
|
|
|
def _command(self, name: str, arguments: dict | None, deadline: float) -> dict:
|
|
self.counter += 1
|
|
request = {'execute': name, 'id': self.counter}
|
|
if arguments is not None:
|
|
request['arguments'] = arguments
|
|
try:
|
|
self.socket.settimeout(max(0.001, deadline - time.monotonic()))
|
|
self.socket.sendall(json.dumps(request).encode() + b'\r\n')
|
|
except OSError as exc:
|
|
raise QMPError(f'QMP send failed: {exc}') from exc
|
|
while True:
|
|
response = self._receive(deadline)
|
|
if response.get('id') != self.counter:
|
|
continue
|
|
if 'error' in response:
|
|
raise QMPError(f'QMP {name}: {response["error"]}')
|
|
if 'return' not in response:
|
|
raise QMPError('QMP reply has no return value')
|
|
return response['return']
|
|
|
|
def command(self, name: str, arguments: dict | None = None) -> dict:
|
|
return self._command(name, arguments, time.monotonic() + self.timeout)
|
|
|
|
def close(self):
|
|
self.socket.close()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
self.close()
|