90 lines
4.1 KiB
Python
90 lines
4.1 KiB
Python
"""Interactive, process-pinned guest terminal transport over the existing vsock agent."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import struct
|
|
import time
|
|
|
|
from .errors import GuestAgentUnavailableError, GuestExecutionError, ValidationError
|
|
from .guest_agent import (
|
|
CONNECT_TIMEOUT_SECONDS, MAX_TERMINAL_FRAME, PROTOCOL_VERSION, ProtocolError,
|
|
receive_frame, send_frame, terminal_data, validate_terminal_size,
|
|
)
|
|
from .guest_exec import GuestSession
|
|
|
|
|
|
def validate_size(rows: object, cols: object) -> None:
|
|
try:
|
|
validate_terminal_size(rows, cols)
|
|
except ProtocolError as error:
|
|
raise ValidationError(str(error)) from error
|
|
|
|
|
|
class GuestTerminal(GuestSession):
|
|
def start(self, rows: int, cols: int) -> None:
|
|
validate_size(rows, cols)
|
|
try:
|
|
deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS
|
|
send_frame(self._connection, {"version": PROTOCOL_VERSION, "terminal": {"rows": rows, "cols": cols}},
|
|
MAX_TERMINAL_FRAME, deadline)
|
|
response = receive_frame(self._connection, MAX_TERMINAL_FRAME, deadline)
|
|
if response != {"type": "ready"}:
|
|
raise ProtocolError("terminal support is unavailable")
|
|
self._connection.setblocking(False)
|
|
self._write_lock = asyncio.Lock()
|
|
except (OSError, EOFError, ProtocolError) as error:
|
|
self.close()
|
|
raise GuestAgentUnavailableError(
|
|
"Interactive terminal support is unavailable. Stop this VM, use Install / update terminal support "
|
|
"in its Sessions view, then start it and launch a new session."
|
|
) from error
|
|
|
|
async def _receive_exact(self, size: int) -> bytes:
|
|
data = bytearray()
|
|
while len(data) < size:
|
|
chunk = await asyncio.get_running_loop().sock_recv(self._connection, size - len(data))
|
|
if not chunk:
|
|
raise EOFError("guest terminal connection closed")
|
|
data.extend(chunk)
|
|
return bytes(data)
|
|
|
|
async def read(self) -> dict:
|
|
try:
|
|
# An idle shell has no output deadline; partial messages are bounded.
|
|
first = await self._receive_exact(1)
|
|
async with asyncio.timeout(CONNECT_TIMEOUT_SECONDS):
|
|
size, = struct.unpack("!I", first + await self._receive_exact(3))
|
|
if not 0 < size <= MAX_TERMINAL_FRAME:
|
|
raise ProtocolError("invalid terminal frame length")
|
|
message = json.loads(await self._receive_exact(size))
|
|
if isinstance(message, dict):
|
|
if set(message) == {"type", "data"} and message["type"] == "output":
|
|
return {"type": "output", "data": terminal_data(message["data"])}
|
|
if (set(message) == {"type", "exit_code"} and message["type"] == "exit"
|
|
and type(message["exit_code"]) is int and -64 <= message["exit_code"] <= 255):
|
|
return message
|
|
raise ProtocolError("invalid guest terminal response")
|
|
except (OSError, EOFError, ValueError, UnicodeError, RecursionError) as error:
|
|
raise GuestExecutionError("The guest terminal connection ended unexpectedly.", outcome_unknown=True) from error
|
|
|
|
async def _send(self, message: dict) -> None:
|
|
payload = json.dumps(message, separators=(",", ":")).encode("ascii")
|
|
if len(payload) > MAX_TERMINAL_FRAME:
|
|
raise ValidationError("terminal input exceeds its size limit")
|
|
async with self._write_lock:
|
|
async with asyncio.timeout(CONNECT_TIMEOUT_SECONDS):
|
|
await asyncio.get_running_loop().sock_sendall(self._connection, struct.pack("!I", len(payload)) + payload)
|
|
|
|
async def write(self, data: bytes) -> None:
|
|
await self._send({"type": "input", "data": base64.b64encode(data).decode("ascii")})
|
|
|
|
async def resize(self, rows: int, cols: int) -> None:
|
|
validate_size(rows, cols)
|
|
await self._send({"type": "resize", "rows": rows, "cols": cols})
|
|
|
|
async def end(self) -> None:
|
|
await self._send({"type": "close"})
|