Something went wrong. Try again.
small gleam coding and (not yet) persistent agent daemon with a detachable cli forked from okami.mom/albedo
Something went wrong. Try again.
11 kB · 311 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312"""A small asyncio WebSocket client, and a stdio bridge for ssh's ProxyCommand.
python -m albedo_plugins.ws # ALBEDO_WS_URL, and ALBEDO_WS_HEADERS as a JSON object
The bridge copies stdin to binary messages and messages to stdout until eitherside ends. The url and headers travel in the environment because they carrytokens, which a command line would show to every process on the host."""
from __future__ import annotations
import asyncioimport base64from dataclasses import dataclassimport hashlibimport jsonimport osimport structimport sysimport urllib.parse
GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"TEXT, BINARY, CLOSE, PING, PONG, CONTINUATION = 0x1, 0x2, 0x8, 0x9, 0xA, 0x0
@dataclassclass CloseFrame: code: int reason: str
class ConnectionClosed(Exception): def __init__(self, code: int = 1006, reason: str = "") -> None: super().__init__(f"WebSocket closed: {code} {reason}") self.code = code self.reason = reason self.rcvd = CloseFrame(code, reason) self.sent: CloseFrame | None = None
def mask(payload: bytes, key: bytes) -> bytes: """`payload` XORed with the repeating 4-byte `key`; applying it twice undoes it.""" size = len(payload) stream = (key * (size // 4 + 1))[:size] return (int.from_bytes(payload, "big") ^ int.from_bytes(stream, "big")).to_bytes( size, "big" )
def frame(opcode: int, payload: bytes) -> bytes: """One final, masked client frame.""" key = os.urandom(4) size = len(payload) if size <= 125: head = struct.pack("!BB", 0x80 | opcode, 0x80 | size) elif size <= 0xFFFF: head = struct.pack("!BBH", 0x80 | opcode, 0x80 | 126, size) else: head = struct.pack("!BBQ", 0x80 | opcode, 0x80 | 127, size) return head + key + mask(payload, key)
class SimpleWebSocketClient: def __init__( self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, max_message_bytes: int = 16 * 1024 * 1024, ) -> None: self.reader = reader self.writer = writer self.max_message_bytes = max_message_bytes self._closed = False self._close_sent = False self.close_code: int | None = None self.close_reason: str = ""
@classmethod async def connect( cls, url: str, *, headers: dict[str, str] | None = None, timeout: float = 15, max_message_bytes: int = 16 * 1024 * 1024, ) -> SimpleWebSocketClient: parsed = urllib.parse.urlsplit(url) if parsed.scheme not in ("ws", "wss"): raise ValueError(f"Expected ws:// or wss:// URL, got: {url!r}") hostname = parsed.hostname or "127.0.0.1" port = parsed.port or (443 if parsed.scheme == "wss" else 80) ssl_ctx = True if parsed.scheme == "wss" else None
async with asyncio.timeout(timeout): reader, writer = await asyncio.open_connection( hostname, port, ssl=ssl_ctx, limit=max_message_bytes ) key = base64.b64encode(os.urandom(16)).decode("ascii") path = parsed.path or "/" if parsed.query: path = f"{path}?{parsed.query}"
host_header = f"{hostname}:{port}" if parsed.port else hostname lines = [ f"GET {path} HTTP/1.1", f"Host: {host_header}", "Upgrade: websocket", "Connection: Upgrade", f"Sec-WebSocket-Key: {key}", "Sec-WebSocket-Version: 13", ] if headers: for k, v in headers.items(): lines.append(f"{k}: {v}") raw_req = "\r\n".join(lines) + "\r\n\r\n" writer.write(raw_req.encode("ascii")) await writer.drain()
status_line = await reader.readline() if not status_line: writer.close() raise ConnectionClosed(1006, "EOF during WebSocket handshake") parts = status_line.decode("latin1").split(" ", 2) if len(parts) < 2 or parts[1] != "101": body = await reader.read(4096) writer.close() raise ConnectionClosed( 1002, f"Handshake failed: {status_line.decode('latin1').strip()} {body.decode('latin1', errors='replace')}", )
resp_headers: dict[str, str] = {} while True: line = await reader.readline() if line in (b"\r\n", b"\n", b""): break hparts = line.decode("latin1").split(":", 1) if len(hparts) == 2: resp_headers[hparts[0].strip().lower()] = hparts[1].strip()
expected_accept = base64.b64encode( hashlib.sha1((key + GUID).encode("ascii")).digest() ).decode("ascii") accept = resp_headers.get("sec-websocket-accept") if accept != expected_accept: writer.close() raise ConnectionClosed( 1002, f"Sec-WebSocket-Accept mismatch: expected {expected_accept}, got {accept}", )
return cls(reader, writer, max_message_bytes=max_message_bytes)
async def send(self, data: str | bytes) -> None: """Send one text message for a str, one binary message for bytes.""" if self._closed: raise ConnectionClosed(self.close_code or 1000, self.close_reason) if isinstance(data, str): self.writer.write(frame(TEXT, data.encode("utf-8"))) else: self.writer.write(frame(BINARY, data)) await self.writer.drain()
async def recv_bytes(self) -> bytes: """The next message's bytes; control frames are answered on the way.""" fragments: list[bytes] = [] total_len = 0
while True: if self._closed: raise ConnectionClosed(self.close_code or 1000, self.close_reason)
try: head = await self.reader.readexactly(2) except ( asyncio.IncompleteReadError, ConnectionResetError, BrokenPipeError, ) as exc: self._closed = True raise ConnectionClosed(1006, "Connection closed abruptly") from exc
b1, b2 = head[0], head[1] fin = bool(b1 & 0x80) opcode = b1 & 0x0F is_masked = bool(b2 & 0x80) payload_len = b2 & 0x7F
if payload_len == 126: ext = await self.reader.readexactly(2) payload_len = struct.unpack("!H", ext)[0] elif payload_len == 127: ext = await self.reader.readexactly(8) payload_len = struct.unpack("!Q", ext)[0]
total_len += payload_len if total_len > self.max_message_bytes: await self.close(1009, "Message exceeds size limit") exc = ConnectionClosed( 1009, f"Message exceeded max_message_bytes {self.max_message_bytes}" ) exc.sent = CloseFrame(1009, "Message exceeds size limit") raise exc
key = await self.reader.readexactly(4) if is_masked else b"" payload = await self.reader.readexactly(payload_len) if is_masked: payload = mask(payload, key)
if opcode == CLOSE: code = 1000 reason = "" if len(payload) >= 2: code = struct.unpack("!H", payload[:2])[0] reason = payload[2:].decode("utf-8", errors="replace") self.close_code = code self.close_reason = reason self._closed = True if not self._close_sent: try: self._close_sent = True self.writer.write(frame(CLOSE, b"")) await self.writer.drain() self.writer.close() except Exception: pass raise ConnectionClosed(code, reason) elif opcode == PING: self.writer.write(frame(PONG, payload)) await self.writer.drain() elif opcode == PONG: continue elif opcode in (TEXT, BINARY, CONTINUATION): fragments.append(payload) if fin: return b"".join(fragments) else: raise ConnectionClosed(1002, f"Unknown WebSocket opcode: {opcode}")
async def recv_message(self) -> str: return (await self.recv_bytes()).decode("utf-8", errors="replace")
def __aiter__(self) -> SimpleWebSocketClient: return self
async def __anext__(self) -> str: try: return await self.recv_message() except ConnectionClosed: raise StopAsyncIteration
async def close(self, code: int = 1000, reason: str = "") -> None: if self._closed and self._close_sent: return self._closed = True if not self._close_sent: self._close_sent = True try: payload = struct.pack("!H", code) + reason.encode("utf-8")[:123] self.writer.write(frame(CLOSE, payload)) await self.writer.drain() self.writer.close() await self.writer.wait_closed() except Exception: pass
async def bridge(url: str, headers: dict[str, str]) -> None: """Copy stdin to the socket and the socket to stdout until either ends.""" socket = await SimpleWebSocketClient.connect(url, headers=headers) loop = asyncio.get_running_loop() stdin = asyncio.StreamReader() await loop.connect_read_pipe( lambda: asyncio.StreamReaderProtocol(stdin), sys.stdin.buffer )
async def upstream() -> None: while chunk := await stdin.read(65536): await socket.send(chunk)
async def downstream() -> None: try: while True: sys.stdout.buffer.write(await socket.recv_bytes()) sys.stdout.buffer.flush() except ConnectionClosed: pass
tasks = [asyncio.create_task(upstream()), asyncio.create_task(downstream())] try: await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) finally: for task in tasks: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) await socket.close()
def main() -> int: try: url = os.environ["ALBEDO_WS_URL"] headers = json.loads(os.environ.get("ALBEDO_WS_HEADERS", "{}")) asyncio.run(bridge(url, headers)) except (KeyError, ValueError, OSError, ConnectionClosed) as error: sys.stderr.write(f"ws bridge: {error!r}\n") return 1 return 0
if __name__ == "__main__": sys.exit(main())