Nexus/server/infrastructure/guacamole/protocol.py

"""Guacamole protocol helpers (instruction encode/decode + guacd handshake)."""

from __future__ import annotations

import asyncio
import logging
import uuid
from typing import Optional

logger = logging.getLogger("nexus.guacamole.protocol")


def encode_element(value: str) -> str:
    s = value if value is not None else ""
    return f"{len(s)}.{s}"


def encode_instruction(opcode: str, *args: str) -> str:
    parts = [encode_element(opcode)] + [encode_element(a) for a in args]
    return ",".join(parts) + ";"


def parse_instruction(message: str) -> tuple[str, list[str]]:
    """Parse one complete instruction (without trailing semicolon)."""
    elements: list[str] = []
    i = 0
    length = len(message)
    while i < length:
        dot = message.find(".", i)
        if dot < 0:
            raise ValueError(f"invalid guacamole instruction at {i}")
        el_len = int(message[i:dot])
        start = dot + 1
        end = start + el_len
        elements.append(message[start:end])
        i = end
        if i < length and message[i] == ",":
            i += 1
        else:
            break
    if not elements:
        raise ValueError("empty guacamole instruction")
    return elements[0], elements[1:]


class GuacdStream:
    """Buffered guacd TCP reader — one instruction per read()."""

    def __init__(self, reader: asyncio.StreamReader):
        self._reader = reader
        self._buf = b""

    async def read_instruction(self) -> tuple[str, list[str]]:
        while True:
            semi = self._buf.find(b";")
            if semi >= 0:
                raw = self._buf[:semi]
                self._buf = self._buf[semi + 1 :]
                msg = raw.decode("utf-8", errors="replace")
                return parse_instruction(msg)
            chunk = await self._reader.read(8192)
            if not chunk:
                raise ConnectionError("guacd closed during read")
            self._buf += chunk

    async def read_raw_instruction(self) -> bytes:
        """Read one complete guacd instruction (including trailing ``;``)."""
        while True:
            semi = self._buf.find(b";")
            if semi >= 0:
                raw = self._buf[: semi + 1]
                self._buf = self._buf[semi + 1 :]
                return raw
            chunk = await self._reader.read(8192)
            if not chunk:
                raise ConnectionError("guacd closed during read")
            self._buf += chunk

    async def read_raw_batch(self) -> bytes:
        """Read all currently buffered complete instructions (may be multiple)."""
        parts: list[bytes] = []
        while True:
            semi = self._buf.find(b";")
            if semi < 0:
                if not parts:
                    chunk = await self._reader.read(8192)
                    if not chunk:
                        if parts:
                            break
                        raise ConnectionError("guacd closed during read")
                    self._buf += chunk
                    continue
                break
            parts.append(await self.read_raw_instruction())
        return b"".join(parts)


def _param_for_arg(arg_name: str, params: dict[str, str]) -> str:
    """Map guacd arg name to connection parameter value."""
    normalized = arg_name.replace("_", "-").lower()
    for key, value in params.items():
        if key.replace("_", "-").lower() == normalized:
            return value or ""
    return ""


async def guacd_rdp_handshake(
    stream: GuacdStream,
    writer: asyncio.StreamWriter,
    params: dict[str, str],
    *,
    width: int = 1024,
    height: int = 768,
    dpi: int = 96,
) -> str:
    """Run guacd handshake for RDP; return connection id from ``ready``."""
    writer.write(encode_instruction("select", "rdp").encode("utf-8"))
    await writer.drain()

    arg_names: list[str] = []
    while True:
        opcode, args = await stream.read_instruction()
        if opcode == "args":
            arg_names = list(args)
            break
        if opcode == "error":
            raise RuntimeError(args[0] if args else "guacd error during args")
        logger.warning("guacd unexpected during select: %s %s", opcode, args)

    writer.write(encode_instruction("size", str(width), str(height), str(dpi)).encode("utf-8"))
    await writer.drain()
    writer.write(encode_instruction("audio").encode("utf-8"))
    await writer.drain()
    writer.write(encode_instruction("video").encode("utf-8"))
    await writer.drain()

    connect_values = [_param_for_arg(name, params) for name in arg_names]
    writer.write(encode_instruction("connect", *connect_values).encode("utf-8"))
    await writer.drain()

    while True:
        opcode, args = await stream.read_instruction()
        if opcode == "ready":
            if not args:
                raise RuntimeError("guacd ready without connection id")
            return args[0]
        if opcode == "error":
            raise RuntimeError(args[0] if args else "guacd connection failed")
        logger.debug("guacd during handshake: %s %s", opcode, args)


def tunnel_open_instruction(connection_id: Optional[str] = None) -> str:
    """Empty opcode + UUID — tells guacamole-common-js the tunnel is OPEN."""
    cid = connection_id or str(uuid.uuid4())
    return encode_instruction("", cid)


def is_internal_instruction(message: str) -> bool:
    """True if message uses internal (empty) opcode — not for guacd."""
    try:
        opcode, _ = parse_instruction(message.rstrip(";"))
        return opcode == ""
    except ValueError:
        return False