"""Interactive pseudo-terminal lifecycle and WebSocket streaming.""" import base64 import hashlib import json import secrets import socket import struct import time import unittest import urllib.parse import urllib.request from harness import Albedo class TerminalWebSocket: def __init__(self, app, terminal_id, **params): address = urllib.parse.urlsplit(app.base) self.socket = socket.create_connection((address.hostname, address.port), 20) self.reader = self.socket.makefile("rb") key = base64.b64encode(secrets.token_bytes(16)).decode() query = urllib.parse.urlencode(params) path = f"/terminals/{terminal_id}" + (f"?{query}" if query else "") self.socket.sendall( ( f"GET {path} HTTP/1.1\r\nHost: {address.netloc}\r\n" "Connection: Upgrade\r\nUpgrade: websocket\r\n" f"Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n" f"Authorization: Bearer {app.connection['token']}\r\n\r\n" ).encode() ) try: status = self.reader.readline() if b" 101 " not in status: raise AssertionError(f"WebSocket upgrade failed: {status!r}") headers = {} while (line := self.reader.readline()) not in (b"\r\n", b""): name, value = line.decode().split(":", 1) headers[name.lower()] = value.strip() expected = base64.b64encode( hashlib.sha1( (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode() ).digest() ).decode() if headers.get("sec-websocket-accept") != expected: raise AssertionError(f"bad WebSocket accept: {headers}") except BaseException: self.close() raise def close(self): try: self.reader.close() self.socket.close() except OSError: pass def frame(self, opcode, payload): mask = secrets.token_bytes(4) size = len(payload) header = bytes([0x80 | opcode]) if size < 126: header += bytes([0x80 | size]) elif size < 65536: header += bytes([0x80 | 126]) + struct.pack("!H", size) else: header += bytes([0x80 | 127]) + struct.pack("!Q", size) self.socket.sendall( header + mask + bytes(value ^ mask[i % 4] for i, value in enumerate(payload)) ) def send_text(self, **message): self.frame(1, json.dumps(message).encode("utf-8")) def send_bytes(self, data): self.frame(2, data) def read_exact(self, size): value = self.reader.read(size) if len(value) != size: raise AssertionError("WebSocket closed before a complete frame") return value def receive_frame(self, timeout=10): self.socket.settimeout(timeout) flags, length = self.read_exact(2) opcode = flags & 15 if length & 128: raise AssertionError("server sent a masked frame") if length == 126: length = struct.unpack("!H", self.read_exact(2))[0] elif length == 127: length = struct.unpack("!Q", self.read_exact(8))[0] payload = self.read_exact(length) if opcode == 9: self.frame(10, payload) return self.receive_frame(timeout) if opcode == 10: return self.receive_frame(timeout) return opcode, payload class TerminalsTest(unittest.TestCase): def setUp(self): self.app = Albedo(providers={}).__enter__() self.addCleanup(self.app.__exit__, None, None, None) def test_terminal_session_io_and_lifecycle(self): tid = "term-test-1" ws = TerminalWebSocket(self.app, tid, cols=80, rows=24, shell="/bin/sh") self.addCleanup(ws.close) # 1. Expect initial ready frame (text) opcode, payload = ws.receive_frame() self.assertEqual(opcode, 1) ready = json.loads(payload.decode("utf-8")) self.assertEqual(ready["type"], "ready", f"terminal returned error: {ready}") self.assertEqual(ready["id"], tid) self.assertEqual(ready["cols"], 80) self.assertEqual(ready["rows"], 24) self.assertEqual(ready["state"], "running") # 2. Send command to PTY ws.send_bytes(b"echo ALBEDO_TERM_OK\r\n") # 3. Read output frames until our string appears output = bytearray() deadline = time.monotonic() + 10.0 while time.monotonic() < deadline: opcode, chunk = ws.receive_frame(timeout=5) if opcode == 2: output.extend(chunk) if b"ALBEDO_TERM_OK" in output: break self.assertIn(b"ALBEDO_TERM_OK", output) # 4. Test resize command ws.send_text(type="resize", cols=100, rows=30) ws.send_bytes(b"stty size\r\n") deadline = time.monotonic() + 10.0 size_output = bytearray() while time.monotonic() < deadline: opcode, chunk = ws.receive_frame(timeout=5) if opcode == 2: size_output.extend(chunk) if b"30 100" in size_output: break self.assertIn(b"30 100", size_output) # 5. Test signal (SIGINT stops running cat) ws.send_bytes(b"cat\r\n") time.sleep(0.05) ws.send_text(type="signal", signal="SIGINT") ws.send_bytes(b"echo AFTER_SIGINT\r\n") deadline = time.monotonic() + 10.0 sig_output = bytearray() while time.monotonic() < deadline: opcode, chunk = ws.receive_frame(timeout=5) if opcode == 2: sig_output.extend(chunk) if b"AFTER_SIGINT" in sig_output: break self.assertIn(b"AFTER_SIGINT", sig_output) ws.close() # 6. Reconnect to the same terminal: scrollback history should be preserved ws2 = TerminalWebSocket(self.app, tid, shell="/bin/sh") self.addCleanup(ws2.close) opcode, payload = ws2.receive_frame() self.assertEqual(opcode, 1) ready2 = json.loads(payload.decode("utf-8")) self.assertEqual(ready2["id"], tid) # Following frame is binary history opcode, history = ws2.receive_frame() self.assertEqual(opcode, 2) self.assertIn(b"ALBEDO_TERM_OK", history) # 7. Test in-stream restart ws2.send_text(type="restart", cols=80, rows=24) while True: opcode, payload = ws2.receive_frame(timeout=5) if opcode == 1: break restarted = json.loads(payload.decode("utf-8")) self.assertEqual(restarted["type"], "ready") self.assertEqual(restarted["id"], tid) ws2.close() # 8. Explicit delete endpoint req = urllib.request.Request( f"{self.app.base}/terminals/{tid}", headers={"Authorization": f"Bearer {self.app.connection['token']}"}, method="DELETE", ) with urllib.request.urlopen(req, timeout=5) as resp: self.assertEqual(resp.status, 200) res = json.load(resp) self.assertEqual(res["closed"], tid) # 9. Connecting after deletion gives fresh terminal with no old history ws3 = TerminalWebSocket(self.app, tid, shell="/bin/sh") self.addCleanup(ws3.close) opcode, payload = ws3.receive_frame() self.assertEqual(opcode, 1) ready3 = json.loads(payload.decode("utf-8")) self.assertEqual(ready3["type"], "ready") ws3.send_bytes(b"echo FRESH_AFTER_DELETE\r\n") fresh_out = bytearray() deadline = time.monotonic() + 10.0 while time.monotonic() < deadline: opcode, chunk = ws3.receive_frame(timeout=5) if opcode == 2: fresh_out.extend(chunk) if b"FRESH_AFTER_DELETE" in fresh_out: break self.assertIn(b"FRESH_AFTER_DELETE", fresh_out) self.assertNotIn(b"AFTER_SIGINT", fresh_out) ws3.close() if __name__ == "__main__": unittest.main()