initial commit

zaloha/protocol.py

@@ -0,0 +1,125 @@
+"""Newline-delimited JSON protocol with base64 encoding for binary data."""
+
+import asyncio
+import base64
+import binascii
+import json
+import os
+
+
+def encode_binary(data: bytes) -> str:
+    """Encode bytes to base64 string."""
+    return base64.b64encode(data).decode("ascii")
+
+
+def decode_binary(data: str) -> bytes:
+    """Decode base64 string to bytes."""
+    try:
+        return base64.b64decode(data)
+    except (TypeError, binascii.Error) as e:
+        raise ValueError(f"Invalid base64: {e}")
+
+
+VERSION = "0.8.2"
+MIN_CLIENT_VERSION = "0.8"  # server rejects clients older than this
+
+
+def version_gte(version: str, minimum: str) -> bool:
+    """Return True if version >= minimum (compares numeric tuples, e.g. '0.8.1' >= '0.8')."""
+    def _parse(v: str) -> tuple[int, ...]:
+        try:
+            return tuple(int(x) for x in v.split("."))
+        except (ValueError, AttributeError):
+            return (0,)
+    return _parse(version) >= _parse(minimum)
+
+
+MAX_MESSAGE_BYTES = int(os.getenv("MAX_MESSAGE_BYTES", "65536"))  # 64 KiB default
+MAX_IMAGE_BYTES = int(os.getenv("MAX_IMAGE_BYTES", str(5 * 1024 * 1024)))  # 5 MiB default, 0 = no limit
+MAX_FILE_BYTES = int(os.getenv("MAX_FILE_BYTES", str(50 * 1024 * 1024)))  # 50 MiB default
+IMAGE_CHUNK_SIZE = 32768  # 32 KiB raw chunk size for image upload/download
+
+
+def build_request(msg_type: str, request_id: str | None = None, **kwargs) -> bytes:
+    """Build a protocol message (newline-terminated JSON)."""
+    msg = {"type": msg_type, **kwargs}
+    if request_id:
+        msg["request_id"] = request_id
+    return json.dumps(msg, ensure_ascii=False).encode("utf-8") + b"\n"
+
+
+def build_response(
+    msg_type: str,
+    status: str,
+    data: dict | None = None,
+    request_id: str | None = None,
+) -> bytes:
+    """Build a server response."""
+    msg = {"type": msg_type, "status": status}
+    if data is not None:
+        msg["data"] = data
+    if request_id:
+        msg["request_id"] = request_id
+    return json.dumps(msg, ensure_ascii=False).encode("utf-8") + b"\n"
+
+
+def parse_message(line: bytes) -> dict:
+    """Parse a single protocol message from bytes."""
+    try:
+        return json.loads(line.decode("utf-8"))
+    except (json.JSONDecodeError, UnicodeDecodeError) as e:
+        raise ValueError(f"Invalid message: {e}")
+
+
+class ProtocolReader:
+    """Read newline-delimited JSON messages from an asyncio StreamReader."""
+
+    def __init__(self, reader: asyncio.StreamReader):
+        self._reader = reader
+
+    async def read_message(self) -> dict | None:
+        """Read and parse one message. Returns None on EOF."""
+        try:
+            line = await self._reader.readuntil(b"\n")
+        except (asyncio.IncompleteReadError, ConnectionError):
+            return None
+        except asyncio.LimitOverrunError:
+            # Message exceeded limit — drain the internal buffer and signal error
+            self._reader._buffer.clear()
+            self._reader._maybe_resume_transport()
+            raise ValueError("Message exceeds maximum size")
+        if not line:
+            return None
+        return parse_message(line.strip())
+
+
+class ProtocolWriter:
+    """Write newline-delimited JSON messages to an asyncio StreamWriter."""
+
+    def __init__(self, writer: asyncio.StreamWriter):
+        self._writer = writer
+
+    async def send_request(self, msg_type: str, request_id: str | None = None, **kwargs):
+        """Send a request message."""
+        payload = build_request(msg_type, request_id=request_id, **kwargs)
+        if len(payload) > MAX_MESSAGE_BYTES:
+            raise ValueError(f"Message exceeds limit ({len(payload)} > {MAX_MESSAGE_BYTES})")
+        self._writer.write(payload)
+        await self._writer.drain()
+
+    async def send_response(
+        self,
+        msg_type: str,
+        status: str,
+        data: dict | None = None,
+        request_id: str | None = None,
+    ):
+        """Send a response message."""
+        payload = build_response(msg_type, status, data, request_id=request_id)
+        if len(payload) > MAX_MESSAGE_BYTES:
+            raise ValueError(f"Message exceeds limit ({len(payload)} > {MAX_MESSAGE_BYTES})")
+        self._writer.write(payload)
+        await self._writer.drain()
+
+    def close(self):
+        self._writer.close()