okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

nekomimi.petcolab: refresh tokens in the daemon, share the websocket and google oauth code, drop the ssh config editsc9cb3a0

main
branchesmainall refs →
10.9 KiB311 linesraw
1"""A small asyncio WebSocket client, and a stdio bridge for ssh's ProxyCommand.
2
3    python -m albedo_plugins.ws    # ALBEDO_WS_URL, and ALBEDO_WS_HEADERS as a JSON object
4
5The bridge copies stdin to binary messages and messages to stdout until either
6side ends. The url and headers travel in the environment because they carry
7tokens, which a command line would show to every process on the host.
8"""
9
10from __future__ import annotations
11
12import asyncio
13import base64
14from dataclasses import dataclass
15import hashlib
16import json
17import os
18import struct
19import sys
20import urllib.parse
21
22GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
23TEXT, BINARY, CLOSE, PING, PONG, CONTINUATION = 0x1, 0x2, 0x8, 0x9, 0xA, 0x0
24
25
26@dataclass
27class CloseFrame:
28    code: int
29    reason: str
30
31
32class ConnectionClosed(Exception):
33    def __init__(self, code: int = 1006, reason: str = "") -> None:
34        super().__init__(f"WebSocket closed: {code} {reason}")
35        self.code = code
36        self.reason = reason
37        self.rcvd = CloseFrame(code, reason)
38        self.sent: CloseFrame | None = None
39
40
41def mask(payload: bytes, key: bytes) -> bytes:
42    """`payload` XORed with the repeating 4-byte `key`; applying it twice undoes it."""
43    size = len(payload)
44    stream = (key * (size // 4 + 1))[:size]
45    return (int.from_bytes(payload, "big") ^ int.from_bytes(stream, "big")).to_bytes(
46        size, "big"
47    )
48
49
50def frame(opcode: int, payload: bytes) -> bytes:
51    """One final, masked client frame."""
52    key = os.urandom(4)
53    size = len(payload)
54    if size <= 125:
55        head = struct.pack("!BB", 0x80 | opcode, 0x80 | size)
56    elif size <= 0xFFFF:
57        head = struct.pack("!BBH", 0x80 | opcode, 0x80 | 126, size)
58    else:
59        head = struct.pack("!BBQ", 0x80 | opcode, 0x80 | 127, size)
60    return head + key + mask(payload, key)
61
62
63class SimpleWebSocketClient:
64    def __init__(
65        self,
66        reader: asyncio.StreamReader,
67        writer: asyncio.StreamWriter,
68        max_message_bytes: int = 16 * 1024 * 1024,
69    ) -> None:
70        self.reader = reader
71        self.writer = writer
72        self.max_message_bytes = max_message_bytes
73        self._closed = False
74        self._close_sent = False
75        self.close_code: int | None = None
76        self.close_reason: str = ""
77
78    @classmethod
79    async def connect(
80        cls,
81        url: str,
82        *,
83        headers: dict[str, str] | None = None,
84        timeout: float = 15,
85        max_message_bytes: int = 16 * 1024 * 1024,
86    ) -> SimpleWebSocketClient:
87        parsed = urllib.parse.urlsplit(url)
88        if parsed.scheme not in ("ws", "wss"):
89            raise ValueError(f"Expected ws:// or wss:// URL, got: {url!r}")
90        hostname = parsed.hostname or "127.0.0.1"
91        port = parsed.port or (443 if parsed.scheme == "wss" else 80)
92        ssl_ctx = True if parsed.scheme == "wss" else None
93
94        async with asyncio.timeout(timeout):
95            reader, writer = await asyncio.open_connection(
96                hostname, port, ssl=ssl_ctx, limit=max_message_bytes
97            )
98            key = base64.b64encode(os.urandom(16)).decode("ascii")
99            path = parsed.path or "/"
100            if parsed.query:
101                path = f"{path}?{parsed.query}"
102
103            host_header = f"{hostname}:{port}" if parsed.port else hostname
104            lines = [
105                f"GET {path} HTTP/1.1",
106                f"Host: {host_header}",
107                "Upgrade: websocket",
108                "Connection: Upgrade",
109                f"Sec-WebSocket-Key: {key}",
110                "Sec-WebSocket-Version: 13",
111            ]
112            if headers:
113                for k, v in headers.items():
114                    lines.append(f"{k}: {v}")
115            raw_req = "\r\n".join(lines) + "\r\n\r\n"
116            writer.write(raw_req.encode("ascii"))
117            await writer.drain()
118
119            status_line = await reader.readline()
120            if not status_line:
121                writer.close()
122                raise ConnectionClosed(1006, "EOF during WebSocket handshake")
123            parts = status_line.decode("latin1").split(" ", 2)
124            if len(parts) < 2 or parts[1] != "101":
125                body = await reader.read(4096)
126                writer.close()
127                raise ConnectionClosed(
128                    1002,
129                    f"Handshake failed: {status_line.decode('latin1').strip()} {body.decode('latin1', errors='replace')}",
130                )
131
132            resp_headers: dict[str, str] = {}
133            while True:
134                line = await reader.readline()
135                if line in (b"\r\n", b"\n", b""):
136                    break
137                hparts = line.decode("latin1").split(":", 1)
138                if len(hparts) == 2:
139                    resp_headers[hparts[0].strip().lower()] = hparts[1].strip()
140
141            expected_accept = base64.b64encode(
142                hashlib.sha1((key + GUID).encode("ascii")).digest()
143            ).decode("ascii")
144            accept = resp_headers.get("sec-websocket-accept")
145            if accept != expected_accept:
146                writer.close()
147                raise ConnectionClosed(
148                    1002,
149                    f"Sec-WebSocket-Accept mismatch: expected {expected_accept}, got {accept}",
150                )
151
152            return cls(reader, writer, max_message_bytes=max_message_bytes)
153
154    async def send(self, data: str | bytes) -> None:
155        """Send one text message for a str, one binary message for bytes."""
156        if self._closed:
157            raise ConnectionClosed(self.close_code or 1000, self.close_reason)
158        if isinstance(data, str):
159            self.writer.write(frame(TEXT, data.encode("utf-8")))
160        else:
161            self.writer.write(frame(BINARY, data))
162        await self.writer.drain()
163
164    async def recv_bytes(self) -> bytes:
165        """The next message's bytes; control frames are answered on the way."""
166        fragments: list[bytes] = []
167        total_len = 0
168
169        while True:
170            if self._closed:
171                raise ConnectionClosed(self.close_code or 1000, self.close_reason)
172
173            try:
174                head = await self.reader.readexactly(2)
175            except (
176                asyncio.IncompleteReadError,
177                ConnectionResetError,
178                BrokenPipeError,
179            ) as exc:
180                self._closed = True
181                raise ConnectionClosed(1006, "Connection closed abruptly") from exc
182
183            b1, b2 = head[0], head[1]
184            fin = bool(b1 & 0x80)
185            opcode = b1 & 0x0F
186            is_masked = bool(b2 & 0x80)
187            payload_len = b2 & 0x7F
188
189            if payload_len == 126:
190                ext = await self.reader.readexactly(2)
191                payload_len = struct.unpack("!H", ext)[0]
192            elif payload_len == 127:
193                ext = await self.reader.readexactly(8)
194                payload_len = struct.unpack("!Q", ext)[0]
195
196            total_len += payload_len
197            if total_len > self.max_message_bytes:
198                await self.close(1009, "Message exceeds size limit")
199                exc = ConnectionClosed(
200                    1009, f"Message exceeded max_message_bytes {self.max_message_bytes}"
201                )
202                exc.sent = CloseFrame(1009, "Message exceeds size limit")
203                raise exc
204
205            key = await self.reader.readexactly(4) if is_masked else b""
206            payload = await self.reader.readexactly(payload_len)
207            if is_masked:
208                payload = mask(payload, key)
209
210            if opcode == CLOSE:
211                code = 1000
212                reason = ""
213                if len(payload) >= 2:
214                    code = struct.unpack("!H", payload[:2])[0]
215                    reason = payload[2:].decode("utf-8", errors="replace")
216                self.close_code = code
217                self.close_reason = reason
218                self._closed = True
219                if not self._close_sent:
220                    try:
221                        self._close_sent = True
222                        self.writer.write(frame(CLOSE, b""))
223                        await self.writer.drain()
224                        self.writer.close()
225                    except Exception:
226                        pass
227                raise ConnectionClosed(code, reason)
228            elif opcode == PING:
229                self.writer.write(frame(PONG, payload))
230                await self.writer.drain()
231            elif opcode == PONG:
232                continue
233            elif opcode in (TEXT, BINARY, CONTINUATION):
234                fragments.append(payload)
235                if fin:
236                    return b"".join(fragments)
237            else:
238                raise ConnectionClosed(1002, f"Unknown WebSocket opcode: {opcode}")
239
240    async def recv_message(self) -> str:
241        return (await self.recv_bytes()).decode("utf-8", errors="replace")
242
243    def __aiter__(self) -> SimpleWebSocketClient:
244        return self
245
246    async def __anext__(self) -> str:
247        try:
248            return await self.recv_message()
249        except ConnectionClosed:
250            raise StopAsyncIteration
251
252    async def close(self, code: int = 1000, reason: str = "") -> None:
253        if self._closed and self._close_sent:
254            return
255        self._closed = True
256        if not self._close_sent:
257            self._close_sent = True
258            try:
259                payload = struct.pack("!H", code) + reason.encode("utf-8")[:123]
260                self.writer.write(frame(CLOSE, payload))
261                await self.writer.drain()
262                self.writer.close()
263                await self.writer.wait_closed()
264            except Exception:
265                pass
266
267
268async def bridge(url: str, headers: dict[str, str]) -> None:
269    """Copy stdin to the socket and the socket to stdout until either ends."""
270    socket = await SimpleWebSocketClient.connect(url, headers=headers)
271    loop = asyncio.get_running_loop()
272    stdin = asyncio.StreamReader()
273    await loop.connect_read_pipe(
274        lambda: asyncio.StreamReaderProtocol(stdin), sys.stdin.buffer
275    )
276
277    async def upstream() -> None:
278        while chunk := await stdin.read(65536):
279            await socket.send(chunk)
280
281    async def downstream() -> None:
282        try:
283            while True:
284                sys.stdout.buffer.write(await socket.recv_bytes())
285                sys.stdout.buffer.flush()
286        except ConnectionClosed:
287            pass
288
289    tasks = [asyncio.create_task(upstream()), asyncio.create_task(downstream())]
290    try:
291        await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
292    finally:
293        for task in tasks:
294            task.cancel()
295        await asyncio.gather(*tasks, return_exceptions=True)
296        await socket.close()
297
298
299def main() -> int:
300    try:
301        url = os.environ["ALBEDO_WS_URL"]
302        headers = json.loads(os.environ.get("ALBEDO_WS_HEADERS", "{}"))
303        asyncio.run(bridge(url, headers))
304    except (KeyError, ValueError, OSError, ConnectionClosed) as error:
305        sys.stderr.write(f"ws bridge: {error!r}\n")
306        return 1
307    return 0
308
309
310if __name__ == "__main__":
311    sys.exit(main())