"""Small RFC 6455 client for Firefox's loopback RDP endpoint."""
from __future__ import annotations
import asyncio
import base64
from dataclasses import dataclass
import hashlib
import os
import struct
from typing import Final
from .errors import ConnectionLostError, NotConnectedError, RdpProtocolError
WEBSOCKET_GUID: Final = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
MAX_HEADER_BYTES: Final = 65_536
DEFAULT_MAX_MESSAGE_BYTES: Final = 16 * 1024 * 1024
@dataclass(frozen=True)
class Frame:
"""Represent one decoded WebSocket frame."""
final: bool
opcode: int
payload: bytes
class FrameDecoder:
"""Incrementally decode unmasked server WebSocket frames."""
def __init__(
self,
max_message_bytes: int = DEFAULT_MAX_MESSAGE_BYTES,
) -> None:
self._buffer = bytearray()
self._fragment_opcode: int | None = None
self._fragment_payload = bytearray()
self.max_message_bytes = max_message_bytes
def feed(self, data: bytes) -> list[Frame]:
"""Consume bytes and return all complete validated frames."""
self._buffer.extend(data)
frames: list[Frame] = []
while True:
frame = self._takeFrame()
if frame is None:
break
frames.append(frame)
return frames
def _takeFrame(self) -> Frame | None:
if len(self._buffer) < 2:
return None
first = self._buffer[0]
second = self._buffer[1]
final = bool(first & 0x80)
if first & 0x70:
raise RdpProtocolError("WebSocket reserved bits are nonzero.")
opcode = first & 0x0f
if second & 0x80:
raise RdpProtocolError("Server WebSocket frames must be unmasked.")
length = second & 0x7f
offset = 2
if length == 126:
if len(self._buffer) < 4:
return None
length = struct.unpack("!H", self._buffer[2:4])[0]
offset = 4
if length < 126:
raise RdpProtocolError("Non-minimal WebSocket frame length.")
elif length == 127:
if len(self._buffer) < 10:
return None
length = struct.unpack("!Q", self._buffer[2:10])[0]
offset = 10
if length & (1 << 63):
raise RdpProtocolError("Invalid 64-bit WebSocket length.")
if length <= 65_535:
raise RdpProtocolError("Non-minimal WebSocket frame length.")
if opcode >= 0x8:
if not final or length > 125:
raise RdpProtocolError("Invalid WebSocket control frame.")
if length > self.max_message_bytes:
raise RdpProtocolError("WebSocket message exceeds size limit.")
if len(self._buffer) < offset + length:
return None
payload = bytes(self._buffer[offset:offset + length])
del self._buffer[:offset + length]
if opcode >= 0x8:
if opcode not in {0x8, 0x9, 0xA}:
raise RdpProtocolError("Unsupported WebSocket control opcode.")
return Frame(final, opcode, payload)
return self._handleData(final, opcode, payload)
def _handleData(
self,
final: bool,
opcode: int,
payload: bytes,
) -> Frame:
if opcode == 0:
if self._fragment_opcode is None:
raise RdpProtocolError("Unexpected continuation frame.")
self._fragment_payload.extend(payload)
if len(self._fragment_payload) > self.max_message_bytes:
raise RdpProtocolError(
"WebSocket message exceeds size limit."
)
if not final:
return Frame(False, 0, b"")
result = bytes(self._fragment_payload)
result_opcode = self._fragment_opcode
self._fragment_opcode = None
self._fragment_payload.clear()
self._validateData(result_opcode, result)
return Frame(True, result_opcode, result)
if opcode not in {1, 2}:
raise RdpProtocolError("Unsupported WebSocket data opcode.")
if self._fragment_opcode is not None:
raise RdpProtocolError("Data frame interrupted fragmented message.")
if final:
self._validateData(opcode, payload)
return Frame(True, opcode, payload)
self._fragment_opcode = opcode
self._fragment_payload.extend(payload)
return Frame(False, opcode, b"")
@staticmethod
def _validateData(opcode: int, payload: bytes) -> None:
if opcode == 2:
raise RdpProtocolError("Binary WebSocket messages are unsupported.")
try:
payload.decode("utf-8", errors="strict")
except UnicodeDecodeError as error:
raise RdpProtocolError(
"WebSocket text message is not valid UTF-8."
) from error
def encodeFrame(
payload: bytes,
opcode: int = 1,
*,
final: bool = True,
mask_key: bytes | None = None,
) -> bytes:
"""Encode one mandatory-masked client WebSocket frame."""
key = os.urandom(4) if mask_key is None else mask_key
if len(key) != 4:
raise ValueError("A WebSocket mask key must contain four bytes.")
if opcode >= 0x8 and (not final or len(payload) > 125):
raise ValueError("Control frames must be final and at most 125 bytes.")
first = (0x80 if final else 0) | opcode
length = len(payload)
if length < 126:
header = struct.pack("!BB", first, 0x80 | length)
elif length <= 65_535:
header = struct.pack("!BBH", first, 0x80 | 126, length)
else:
if length >= (1 << 63):
raise ValueError("WebSocket payload is too large.")
header = struct.pack("!BBQ", first, 0x80 | 127, length)
masked = bytes(value ^ key[index % 4]
for index, value in enumerate(payload))
return header + key + masked
class WebSocketClient:
"""Connect and exchange text messages with a local WebSocket server."""
def __init__(
self,
host: str,
port: int,
*,
max_message_bytes: int = DEFAULT_MAX_MESSAGE_BYTES,
) -> None:
self.host = host
self.port = port
self._reader: asyncio.StreamReader | None = None
self._writer: asyncio.StreamWriter | None = None
self._decoder = FrameDecoder(max_message_bytes)
self._frames: list[Frame] = []
self._write_lock = asyncio.Lock()
self._close_sent = False
self.last_pong: bytes | None = None
@property
def connected(self) -> bool:
"""Return whether an open stream is currently attached."""
return self._writer is not None and not self._writer.is_closing()
async def connect(self, timeout: float) -> None:
"""Open TCP and validate the WebSocket HTTP Upgrade response."""
await self.close()
try:
reader, writer = await asyncio.wait_for(
asyncio.open_connection(self.host, self.port),
timeout,
)
except (OSError, asyncio.TimeoutError) as error:
raise NotConnectedError(
f"Cannot connect to Firefox RDP at "
f"{self.host}:{self.port}."
) from error
self._reader = reader
self._writer = writer
self._decoder = FrameDecoder(self._decoder.max_message_bytes)
self._frames.clear()
self._close_sent = False
key = base64.b64encode(os.urandom(16)).decode("ascii")
request = (
"GET / HTTP/1.1\r\n"
f"Host: {self.host}:{self.port}\r\n"
"Upgrade: websocket\r\n"
"Connection: Upgrade\r\n"
f"Sec-WebSocket-Key: {key}\r\n"
"Sec-WebSocket-Version: 13\r\n\r\n"
).encode("ascii")
writer.write(request)
await writer.drain()
response = await self._readHandshake(timeout)
self._validateHandshake(response, key)
async def _readHandshake(self, timeout: float) -> bytes:
assert self._reader is not None
data = bytearray()
try:
while b"\r\n\r\n" not in data:
chunk = await asyncio.wait_for(
self._reader.read(4096),
timeout,
)
if not chunk:
raise ConnectionLostError(
"Firefox closed during the WebSocket handshake."
)
data.extend(chunk)
if len(data) > MAX_HEADER_BYTES:
raise RdpProtocolError(
"WebSocket Upgrade headers exceed 64 KiB."
)
except asyncio.TimeoutError as error:
raise ConnectionLostError(
"Timed out during the WebSocket handshake."
) from error
delimiter = data.index(b"\r\n\r\n") + 4
remainder = bytes(data[delimiter:])
if remainder:
self._frames.extend(self._decoder.feed(remainder))
return bytes(data[:delimiter])
@staticmethod
def _validateHandshake(response: bytes, key: str) -> None:
try:
lines = response.decode("iso-8859-1").split("\r\n")
status = lines[0].split(" ", 2)
if len(status) < 2 or status[1] != "101":
raise RdpProtocolError(
"WebSocket server did not return status 101."
)
headers: dict[str, list[str]] = {}
for line in lines[1:]:
if not line:
continue
name, separator, value = line.partition(":")
if not separator:
raise RdpProtocolError(
"Malformed WebSocket response header."
)
headers.setdefault(name.lower(), []).append(value.strip())
upgrades = ",".join(headers.get("upgrade", [])).lower()
connections = {
token.strip().lower()
for token in ",".join(
headers.get("connection", [])
).split(",")
}
expected = base64.b64encode(
hashlib.sha1(
(key + WEBSOCKET_GUID).encode("ascii")
).digest()
).decode("ascii")
actual = ",".join(headers.get(
"sec-websocket-accept",
[],
))
if upgrades != "websocket" or "upgrade" not in connections:
raise RdpProtocolError(
"WebSocket Upgrade response is missing required tokens."
)
if actual != expected:
raise RdpProtocolError(
"WebSocket Sec-WebSocket-Accept is incorrect."
)
except UnicodeDecodeError as error:
raise RdpProtocolError(
"WebSocket Upgrade response headers are invalid."
) from error
async def sendText(self, text: str) -> None:
"""Send one UTF-8 text message as a final masked frame."""
await self._sendFrame(text.encode("utf-8"), 1)
async def _sendFrame(self, payload: bytes, opcode: int) -> None:
if not self.connected or self._writer is None:
raise ConnectionLostError("WebSocket is not connected.")
async with self._write_lock:
try:
self._writer.write(encodeFrame(payload, opcode))
await self._writer.drain()
except (OSError, ConnectionError) as error:
raise ConnectionLostError(
"Firefox RDP WebSocket write failed."
) from error
async def receiveText(self) -> str:
"""Receive the next data message while servicing control frames."""
if self._reader is None:
raise ConnectionLostError("WebSocket is not connected.")
while True:
if not self._frames:
try:
data = await self._reader.read(65_536)
except (OSError, ConnectionError) as error:
raise ConnectionLostError(
"Firefox RDP WebSocket read failed."
) from error
if not data:
raise ConnectionLostError(
"Firefox RDP WebSocket closed."
)
self._frames.extend(self._decoder.feed(data))
if not self._frames:
continue
frame = self._frames.pop(0)
if frame.opcode == 0 or not frame.final:
continue
if frame.opcode == 0x9:
await self._sendFrame(frame.payload, 0xA)
continue
if frame.opcode == 0xA:
self.last_pong = frame.payload
continue
if frame.opcode == 0x8:
if not self._close_sent:
self._close_sent = True
await self._sendFrame(frame.payload, 0x8)
await self.close()
raise ConnectionLostError(
"Firefox RDP WebSocket sent a close frame."
)
return frame.payload.decode("utf-8")
async def close(self) -> None:
"""Close streams and discard all decoder state."""
writer = self._writer
self._reader = None
self._writer = None
self._frames.clear()
if writer is not None:
writer.close()
try:
await writer.wait_closed()
except (OSError, ConnectionError):
pass