DefEssential 0.1.8__23-py3-none-any.whl

DefEssential/IO/WebSocket/WebSocket.py

@@ -0,0 +1,288 @@
+from logging import getLogger
+from struct import (
+	pack,
+	unpack,
+)
+from threading import Lock
+from typing import (
+	Optional,
+	Union,
+)
+
+from ...Request import Request
+from .WebSocketError import WebSocketError
+from .WebSocketFrame import WebSocketFrame
+from .WebSocketFrameReader import WebSocketFrameReader
+from .WebSocketHandshake import WebSocketHandshake
+
+__all__ = (
+	'WebSocket',
+)
+
+
+class WebSocket:
+	"""Exchange WebSocket messages over an upgraded HTTP connection.
+
+	Call receive() from a single thread. send(), ping(), and close() may also be
+	called from other threads.
+	"""
+
+	logger = getLogger('DefEssential.IO.WebSocket')
+
+	def __init__(
+		self,
+		request: Request,
+		limit: int,
+	):
+		if limit <= 0:
+			raise ValueError('Message size limit must be positive')
+		self.request = request
+		self.limit = limit
+		self.handshake = WebSocketHandshake(request)
+		self.frames = WebSocketFrameReader(
+			request.reader,
+			limit,
+		)
+		self.lock = Lock()
+		self.accepted = False
+		self.close_sent = False
+		self.closed = False
+		self.code: Optional[int] = None
+		self.reason = ''
+		return
+
+	@property
+	def protocols(self) -> list[str]:
+		return self.handshake.protocols
+
+	def accept(
+		self,
+		protocol: Optional[str] = None,
+	) -> None:
+		"""Validate the handshake and send a 101 response. On failure, a regular error response is sent because the connection has not been upgraded."""
+		if self.accepted:
+			raise RuntimeError('WebSocket has already been accepted')
+		self.request.writer.upgrade(self.handshake.response(protocol))
+		self.accepted = True
+		self.logger.debug('WebSocket upgraded: request_id=%s status=101', self.request.id)
+		return
+
+	def valid_code(
+		self,
+		code: int,
+	) -> bool:
+		return code in (
+			1000,
+			1001,
+			1002,
+			1003,
+			1007,
+			1008,
+			1009,
+			1010,
+			1011,
+			1012,
+			1013,
+			1014,
+		) or 3000 <= code <= 4999
+
+	def write(
+		self,
+		frame: WebSocketFrame,
+	) -> None:
+		with self.lock:
+			if not self.accepted:
+				raise RuntimeError('WebSocket must be accepted first')
+			if self.close_sent:
+				raise RuntimeError('Cannot send after the WebSocket close frame')
+			if frame.opcode == WebSocketFrame.CLOSE:
+				self.close_sent = True
+			self.request.writer.write(frame.encode())
+			self.request.writer.flush()
+		return
+
+	def send(
+		self,
+		message: Union[str, bytes],
+	) -> None:
+		"""Send str values as text messages and bytes values as binary messages."""
+		if isinstance(message, str):
+			frame = WebSocketFrame(
+				WebSocketFrame.TEXT,
+				message.encode('utf-8'),
+			)
+		elif isinstance(message, bytes):
+			frame = WebSocketFrame(
+				WebSocketFrame.BINARY,
+				message,
+			)
+		else:
+			raise TypeError('WebSocket messages must be str or bytes')
+		self.write(frame)
+		return
+
+	def ping(
+		self,
+		data: bytes = b'',
+	) -> None:
+		self.write(WebSocketFrame(
+			WebSocketFrame.PING,
+			data,
+		))
+		return
+
+	def close_frame(
+		self,
+		code: int,
+		reason: str,
+	) -> WebSocketFrame:
+		return WebSocketFrame(
+			WebSocketFrame.CLOSE,
+			pack('!H', code) + reason.encode('utf-8'),
+		)
+
+	def receive_close(
+		self,
+		frame: WebSocketFrame,
+	) -> None:
+		payload = frame.payload
+		if len(payload) == 1:
+			raise WebSocketError(
+				1002,
+				'Invalid close frame',
+			)
+		code = unpack('!H', payload[:2])[0] if payload else None
+		if code is not None and not self.valid_code(code):
+			raise WebSocketError(
+				1002,
+				'Invalid close code',
+			)
+		try:
+			reason = payload[2:].decode('utf-8')
+		except UnicodeDecodeError as error:
+			raise WebSocketError(
+				1007,
+				'Invalid close reason',
+			) from error
+		self.code = code
+		self.reason = reason
+		self.closed = True
+		with self.lock:
+			echo = not self.close_sent
+		if echo:
+			self.write(WebSocketFrame(
+				WebSocketFrame.CLOSE,
+				payload[:2],
+			))
+		self.logger.debug('WebSocket closed: request_id=%s code=%s', self.request.id, self.code)
+		return
+
+	def fail(
+		self,
+		error: WebSocketError,
+	) -> None:
+		self.logger.warning(
+			'WebSocket protocol error: request_id=%s code=%s',
+			self.request.id,
+			error.code,
+		)
+		self.code = error.code
+		self.reason = error.reason
+		self.closed = True
+		with self.lock:
+			send = not self.close_sent
+		if send:
+			self.write(self.close_frame(
+				error.code,
+				error.reason,
+			))
+		self.logger.debug('WebSocket failed: request_id=%s code=%s', self.request.id, error.code)
+		return
+
+	def receive(self) -> Union[str, bytes, None]:
+		"""Return the next message, or None once the connection starts closing.
+
+		Automatically reply to ping frames with pong. On a protocol violation,
+		send a close code and return None. Propagate ClientDisconnected if the
+		TCP connection is lost.
+		"""
+		if not self.accepted:
+			raise RuntimeError('WebSocket must be accepted first')
+		if self.closed:
+			return None
+		opcode: Optional[int] = None
+		parts = []
+		size = 0
+		try:
+			while True:
+				frame = self.frames.read()
+				if frame.opcode == WebSocketFrame.CLOSE:
+					self.receive_close(frame)
+					return None
+				if frame.opcode == WebSocketFrame.PING:
+					with self.lock:
+						reply = not self.close_sent
+					if reply:
+						self.write(WebSocketFrame(
+							WebSocketFrame.PONG,
+							frame.payload,
+						))
+					continue
+				if frame.opcode == WebSocketFrame.PONG:
+					continue
+				if frame.opcode == WebSocketFrame.CONTINUATION:
+					if opcode is None:
+						raise WebSocketError(
+							1002,
+							'Unexpected continuation frame',
+						)
+				elif opcode is not None:
+					raise WebSocketError(
+						1002,
+						'Expected a continuation frame',
+					)
+				else:
+					opcode = frame.opcode
+				size += len(frame.payload)
+				if size > self.limit:
+					raise WebSocketError(
+						1009,
+						'Message too big',
+					)
+				parts.append(frame.payload)
+				if frame.final:
+					break
+			data = b''.join(parts)
+			if opcode == WebSocketFrame.BINARY:
+				return data
+			try:
+				return data.decode('utf-8')
+			except UnicodeDecodeError as error:
+				raise WebSocketError(
+					1007,
+					'Invalid UTF-8 text',
+				) from error
+		except WebSocketError as error:
+			self.fail(error)
+			return None
+
+	def close(
+		self,
+		code: int = 1000,
+		reason: str = '',
+	) -> None:
+		"""Send a close frame and discard messages until the client's close frame arrives."""
+		if not self.valid_code(code):
+			raise ValueError('Invalid WebSocket close code')
+		if len(reason.encode('utf-8')) > 123:
+			raise ValueError('WebSocket close reason must be at most 123 bytes')
+		with self.lock:
+			send = not self.close_sent
+		if send:
+			self.write(self.close_frame(
+				code,
+				reason,
+			))
+		while not self.closed:
+			self.receive()
+		return

DefEssential/IO/WebSocket/WebSocketError.py

@@ -0,0 +1,17 @@
+__all__ = (
+	'WebSocketError',
+)
+
+
+class WebSocketError(Exception):
+	"""Represent a client WebSocket protocol violation with a close code."""
+
+	def __init__(
+		self,
+		code: int,
+		reason: str,
+	):
+		super().__init__(reason)
+		self.code = code
+		self.reason = reason
+		return

DefEssential/IO/WebSocket/WebSocketFrame.py

@@ -0,0 +1,68 @@
+from struct import pack
+
+__all__ = (
+	'WebSocketFrame',
+)
+
+
+class WebSocketFrame:
+	"""Represent a WebSocket frame and encode it as unmasked bytes for server transmission."""
+
+	CONTINUATION = 0x0
+	TEXT = 0x1
+	BINARY = 0x2
+	CLOSE = 0x8
+	PING = 0x9
+	PONG = 0xA
+	OPCODES = (
+		CONTINUATION,
+		TEXT,
+		BINARY,
+		CLOSE,
+		PING,
+		PONG,
+	)
+
+	def __init__(
+		self,
+		opcode: int,
+		payload: bytes = b'',
+		final: bool = True,
+	):
+		if opcode not in self.OPCODES:
+			raise ValueError('Unknown WebSocket opcode')
+		if opcode >= self.CLOSE and (not final or len(payload) > 125):
+			raise ValueError('Control frames must be final and at most 125 bytes')
+		self.opcode = opcode
+		self.payload = payload
+		self.final = final
+		return
+
+	@property
+	def control(self) -> bool:
+		return self.opcode >= self.CLOSE
+
+	def encode(self) -> bytes:
+		first = (0x80 if self.final else 0) | self.opcode
+		size = len(self.payload)
+		if size < 126:
+			header = pack(
+				'!BB',
+				first,
+				size,
+			)
+		elif size < 65536:
+			header = pack(
+				'!BBH',
+				first,
+				126,
+				size,
+			)
+		else:
+			header = pack(
+				'!BBQ',
+				first,
+				127,
+				size,
+			)
+		return header + self.payload

DefEssential/IO/WebSocket/WebSocketFrameReader.py

@@ -0,0 +1,98 @@
+from struct import unpack
+
+from ...Error import ClientError
+from ..RequestReader import (
+	ClientDisconnected,
+	RequestReader,
+)
+from .WebSocketError import WebSocketError
+from .WebSocketFrame import WebSocketFrame
+
+__all__ = (
+	'WebSocketFrameReader',
+)
+
+
+class WebSocketFrameReader:
+	"""Read client frames one at a time from an upgraded connection and validate frame rules."""
+
+	def __init__(
+		self,
+		reader: RequestReader,
+		limit: int,
+	):
+		self.reader = reader
+		self.limit = limit
+		return
+
+	def exact(
+		self,
+		size: int,
+	) -> bytes:
+		try:
+			return self.reader.exact(size)
+		except ClientError as error:
+			# RequestReader.exact() treats EOF as a body error, but EOF after an upgrade means disconnection.
+			raise ClientDisconnected(error.reason) from error
+
+	def unmask(
+		self,
+		data: bytes,
+		mask: bytes,
+	) -> bytes:
+		size = len(data)
+		key = (mask * (size // 4 + 1))[:size]
+		return (int.from_bytes(data, 'big') ^ int.from_bytes(key, 'big')).to_bytes(
+			size,
+			'big',
+		)
+
+	def read(self) -> WebSocketFrame:
+		first, second = self.exact(2)
+		if first & 0x70:
+			raise WebSocketError(
+				1002,
+				'Reserved bits must be zero',
+			)
+		final = bool(first & 0x80)
+		opcode = first & 0x0F
+		if opcode not in WebSocketFrame.OPCODES:
+			raise WebSocketError(
+				1002,
+				'Unknown opcode',
+			)
+		if not second & 0x80:
+			raise WebSocketError(
+				1002,
+				'Client frames must be masked',
+			)
+		size = second & 0x7F
+		if size == 126:
+			size = unpack('!H', self.exact(2))[0]
+		elif size == 127:
+			size = unpack('!Q', self.exact(8))[0]
+			if size >> 63:
+				raise WebSocketError(
+					1002,
+					'Invalid payload length',
+				)
+		control = opcode >= WebSocketFrame.CLOSE
+		if control and (not final or size > 125):
+			raise WebSocketError(
+				1002,
+				'Control frames must be final and at most 125 bytes',
+			)
+		if not control and size > self.limit:
+			raise WebSocketError(
+				1009,
+				'Message too big',
+			)
+		mask = self.exact(4)
+		return WebSocketFrame(
+			opcode,
+			self.unmask(
+				self.exact(size),
+				mask,
+			),
+			final,
+		)

DefEssential/IO/WebSocket/WebSocketHandshake.py

@@ -0,0 +1,94 @@
+from base64 import (
+	b64decode,
+	b64encode,
+)
+from hashlib import sha1
+from typing import Optional
+
+from ...Error import ClientError
+from ...Request import Request
+from ...Response import Response
+
+__all__ = (
+	'WebSocketHandshake',
+)
+
+
+class WebSocketHandshake:
+	"""Validate a WebSocket opening handshake request and build a 101 response."""
+
+	GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
+	VERSION = '13'
+
+	def __init__(
+		self,
+		request: Request,
+	):
+		self.request = request
+		return
+
+	def tokens(
+		self,
+		name: str,
+	) -> list[str]:
+		value = self.request.header(name) or ''
+		return [token.strip() for token in value.split(',') if token.strip()]
+
+	@property
+	def protocols(self) -> list[str]:
+		"""Return the subprotocols requested by the client in request order."""
+		return self.tokens('Sec-WebSocket-Protocol')
+
+	def validate(self) -> str:
+		"""Validate the handshake request and return the Sec-WebSocket-Key."""
+		if self.request.method != 'GET':
+			raise ClientError('WebSocket handshake requires GET')
+		if self.request.writer.version != 'HTTP/1.1':
+			raise ClientError('WebSocket handshake requires HTTP/1.1')
+		if 'websocket' not in (token.lower() for token in self.tokens('Upgrade')):
+			raise ClientError(
+				'WebSocket upgrade is required',
+				status=426,
+				headers={'Upgrade': 'websocket'},
+			)
+		if 'upgrade' not in (token.lower() for token in self.tokens('Connection')):
+			raise ClientError('Connection: Upgrade is required')
+		if self.request.header('Sec-WebSocket-Version') != self.VERSION:
+			raise ClientError(
+				'Unsupported WebSocket version',
+				status=426,
+				headers={'Sec-WebSocket-Version': self.VERSION},
+			)
+		key = self.request.header('Sec-WebSocket-Key')
+		try:
+			valid = key is not None and len(b64decode(key, validate=True)) == 16
+		except ValueError:
+			valid = False
+		if key is None or not valid:
+			raise ClientError('Invalid Sec-WebSocket-Key')
+		return key
+
+	def accept_key(
+		self,
+		key: str,
+	) -> str:
+		return b64encode(sha1((key + self.GUID).encode('ascii')).digest()).decode('ascii')
+
+	def response(
+		self,
+		protocol: Optional[str] = None,
+	) -> Response:
+		key = self.validate()
+		if protocol is not None and protocol not in self.protocols:
+			raise ValueError('Selected subprotocol was not requested by the client')
+		headers = {
+			'Upgrade': 'websocket',
+			'Connection': 'Upgrade',
+			'Sec-WebSocket-Accept': self.accept_key(key),
+		}
+		if protocol is not None:
+			headers['Sec-WebSocket-Protocol'] = protocol
+		return Response(
+			status=101,
+			headers=headers,
+		)

DefEssential/IO/WebSocket/__init__.py

@@ -0,0 +1,13 @@
+from .WebSocket import WebSocket
+from .WebSocketError import WebSocketError
+from .WebSocketFrame import WebSocketFrame
+from .WebSocketFrameReader import WebSocketFrameReader
+from .WebSocketHandshake import WebSocketHandshake
+
+__all__ = (
+	'WebSocket',
+	'WebSocketError',
+	'WebSocketFrame',
+	'WebSocketFrameReader',
+	'WebSocketHandshake',
+)

DefEssential/IO/Writers/ChunkedStreamWriter.py

@@ -0,0 +1,19 @@
+from .StreamWriter import StreamWriter
+
+__all__ = (
+	'ChunkedStreamWriter',
+)
+
+
+class ChunkedStreamWriter(StreamWriter):
+	"""Write directly using HTTP/1.1 chunked framing."""
+
+	chunked = True
+
+	def chunk(
+		self,
+		data: bytes,
+	) -> None:
+		self.write(data)
+		self.flush()
+		return

DefEssential/IO/Writers/ServerSentEvents.py

@@ -0,0 +1,41 @@
+from dataclasses import dataclass
+from typing import Optional
+
+__all__ = (
+	'ServerSentEvents',
+)
+
+
+@dataclass(frozen=True)
+class ServerSentEvents:
+	"""Encode UTF-8 SSE events for both direct writes and iterable output."""
+
+	data: str = ''
+	event: Optional[str] = None
+	id: Optional[str] = None
+	retry: Optional[int] = None
+
+	def retry_line(self) -> str:
+		"""Return a validated reconnection delay field."""
+		if self.retry is None:
+			return ''
+		if isinstance(self.retry, bool) or not isinstance(self.retry, int) or self.retry < 0:
+			raise ValueError('SSE retry must be a non-negative integer')
+		return 'retry: ' + str(self.retry) + '\n'
+
+	def encode(self) -> bytes:
+		lines = []
+		for name, value in (
+			('event', self.event),
+			('id', self.id),
+		):
+			if value is None:
+				continue
+			if '\r' in value or '\n' in value or (name == 'id' and '\x00' in value):
+				raise ValueError('Invalid SSE ' + name)
+			lines.append(name + ': ' + value)
+		if self.retry is not None:
+			lines.append(self.retry_line().rstrip('\n'))
+		data = self.data.replace('\r\n', '\n').replace('\r', '\n')
+		lines.extend('data: ' + line for line in data.split('\n'))
+		return ('\n'.join(lines) + '\n\n').encode('utf-8')

DefEssential/IO/Writers/ServerSentEventsWriter.py

@@ -0,0 +1,67 @@
+from typing import (
+	Mapping,
+	Optional,
+	Union,
+)
+
+from ..ResponseWriter import ResponseWriter
+from .ServerSentEvents import ServerSentEvents
+from .StreamWriter import StreamWriter
+
+__all__ = (
+	'ServerSentEventsWriter',
+)
+
+
+class ServerSentEventsWriter:
+	"""Wrap a plain or chunked stream to send SSE events immediately."""
+
+	def __init__(
+		self,
+		writer: Union[ResponseWriter, StreamWriter],
+	):
+		self.writer = StreamWriter(writer) if isinstance(writer, ResponseWriter) else writer
+		return
+
+	def begin(
+		self,
+		headers: Optional[Mapping[str, str]] = None,
+	) -> None:
+		properties = dict(headers) if headers is not None else {}
+		if not any(key.lower() == 'cache-control' for key in properties):
+			properties['Cache-Control'] = 'no-cache'
+		self.writer.begin(
+			headers=properties,
+			format='text/event-stream',
+			charset='utf-8',
+		)
+		return
+
+	def emit(
+		self,
+		data: str,
+		event: Optional[str] = None,
+		id: Optional[str] = None,
+		retry: Optional[int] = None,
+	) -> None:
+		self.writer.write(ServerSentEvents(
+			data=data,
+			event=event,
+			id=id,
+			retry=retry,
+		).encode())
+		self.writer.flush()
+		return
+
+	def retry(
+		self,
+		ms: int,
+	) -> None:
+		frame = ServerSentEvents(retry=ms).retry_line()
+		self.writer.write((frame + '\n').encode('utf-8'))
+		self.writer.flush()
+		return
+
+	def end(self) -> None:
+		self.writer.end()
+		return