zesdex/security-sidecar/zesdex_sec_daemon/protocol.py

"""Newline-delimited JSON frame protocol for the security sidecar.

Frame format: one JSON object per line, terminated by LF.
Request:  {"id": "<req_id>", "op": "call"|"health"|"install", "tool": "<name>", "args": {...}, "timeout": <ms>}
Response: {"id": "<req_id>", "ok": true, "output": "..."} | {"id": "<req_id>", "ok": false, "error": "..."}
"""

import json
import sys
import traceback
from typing import TextIO, Optional
from dataclasses import dataclass, field
from datetime import datetime, timezone


class SecProtocolError(Exception):
    """Raised on malformed frames or protocol violations."""


@dataclass
class SecRequest:
    req_id: str
    op: str  # "call" | "health" | "install"
    tool: str = ""
    args: dict = field(default_factory=dict)
    timeout_ms: int = 30_000

    @classmethod
    def parse(cls, line: str) -> "SecRequest":
        line = line.strip()
        if not line:
            raise SecProtocolError("empty line")
        try:
            data = json.loads(line)
        except json.JSONDecodeError as e:
            raise SecProtocolError(f"invalid JSON: {e}") from e
        req_id = data.get("id")
        op = data.get("op")
        if not isinstance(req_id, str) or not req_id:
            raise SecProtocolError("missing or invalid 'id'")
        if op not in ("call", "health", "install"):
            raise SecProtocolError(f"unknown op: {op!r}")
        return cls(
            req_id=req_id,
            op=op,
            tool=data.get("tool", ""),
            args=data.get("args", {}),
            timeout_ms=data.get("timeout", 30_000),
        )


@dataclass
class SecResponse:
    req_id: str
    ok: bool
    output: str = ""
    error: str = ""
    duration_ms: int = 0

    def to_json(self) -> str:
        obj = {"id": self.req_id, "ok": self.ok, "ts": datetime.now(timezone.utc).isoformat()}
        if self.ok:
            obj["output"] = self.output
        else:
            obj["error"] = self.error
        obj["duration_ms"] = self.duration_ms
        return json.dumps(obj, ensure_ascii=False)


class FrameReader:
    """Reads one JSON line from a buffered stream, enforcing a 64 MiB hard cap."""

    MAX_FRAME_BYTES = 64 * 1024 * 1024

    def __init__(self, stream: TextIO):
        self._stream = stream

    def read_line(self) -> Optional[str]:
        line = self._stream.readline()
        if not line:
            return None
        if len(line) > self.MAX_FRAME_BYTES:
            raise SecProtocolError(f"frame exceeds {self.MAX_FRAME_BYTES} byte limit")
        return line


def make_handshake_frame(token: str) -> str:
    return json.dumps({"op": "handshake", "token": token})


def verify_handshake(line: str, expected_token: str) -> bool:
    try:
        data = json.loads(line.strip())
        return data.get("op") == "handshake" and data.get("token") == expected_token
    except (json.JSONDecodeError, KeyError):
        return False


def run_daemon(stdin: TextIO, stdout: TextIO, token: str = "", registry=None, timeout_cap_ms: int = 300_000):
    """Read requests from stdin, dispatch to tool registry, write responses to stdout."""
    from .tools import ToolRegistry
    registry = registry or ToolRegistry()
    reader = FrameReader(stdin)
    first = reader.read_line()
    if first is None:
        return
    if token:
        if not verify_handshake(first, token):
            err = json.dumps({"ok": False, "error": "handshake failed"})
            stdout.write(err + "\n")
            stdout.flush()
            return
    else:
        try:
            req = SecRequest.parse(first)
            _handle_request(req, registry, stdout, timeout_cap_ms)
        except SecProtocolError as e:
            _write_error("init", str(e), stdout)

    while True:
        line = reader.read_line()
        if line is None:
            break
        if not line.strip():
            continue
        try:
            req = SecRequest.parse(line)
        except SecProtocolError as e:
            _write_error("unknown", str(e), stdout)
            continue
        _handle_request(req, registry, stdout, timeout_cap_ms)


def _handle_request(req: SecRequest, registry, stdout: TextIO, global_cap_ms: int):
    from .tools import ToolResult
    start = datetime.now(timezone.utc)
    effective_timeout = min(req.timeout_ms, global_cap_ms)
    try:
        if req.op == "health":
            health = registry.health_check()
            resp = SecResponse(req_id=req.req_id, ok=True, output=json.dumps(health))
        elif req.op == "install":
            installer_cls = None
            try:
                from .installer import TieredInstaller
                installer_cls = TieredInstaller
            except ImportError:
                pass
            if installer_cls:
                installer = installer_cls()
                result = installer.install_tool(req.tool)
                resp = SecResponse(req_id=req.req_id, ok=result["ok"], output=result.get("message", ""))
            else:
                resp = SecResponse(req_id=req.req_id, ok=False, error="installer not available")
        else:
            result: ToolResult = registry.run(req.tool, req.args, timeout_ms=effective_timeout)
            if result.ok:
                resp = SecResponse(req_id=req.req_id, ok=True, output=result.output)
            else:
                resp = SecResponse(req_id=req.req_id, ok=False, error=result.error)
    except Exception as exc:
        resp = SecResponse(req_id=req.req_id, ok=False, error=f"dispatch error: {exc}")
        traceback.print_exc(file=sys.stderr)
    elapsed = int((datetime.now(timezone.utc) - start).total_seconds() * 1000)
    resp.duration_ms = elapsed
    stdout.write(resp.to_json() + "\n")
    stdout.flush()


def _write_error(req_id: str, msg: str, stdout: TextIO):
    resp = SecResponse(req_id=req_id, ok=False, error=msg)
    stdout.write(resp.to_json() + "\n")
    stdout.flush()