1"""Message protocol for portal communication.
  2
  3Messages are JSON dicts encrypted with NaCl before being sent through
  4the relay. The relay only sees opaque bytes.
  5
  6Message types:
  7  exec           - Execute Python code (local → remote)
  8  exec_stdout    - Streaming stdout chunk (remote → local)
  9  exec_result    - Execution result (remote → local)
 10  error          - Error response (remote → local)
 11  file_pull      - Request a file (local → remote)
 12  file_data      - File contents chunk (remote → local)
 13  file_push      - Send a file chunk (local → remote)
 14  file_push_result - File push confirmation (remote → local)
 15  ping           - Keepalive (either direction)
 16  pong           - Keepalive response (either direction)
 17"""
 18
 19from __future__ import annotations
 20
 21import base64
 22import math
 23
 24# Max chunk size for file transfers (256KB).
 25FILE_CHUNK_SIZE = 256 * 1024
 26
 27# Max file size for transfers (50MB).
 28MAX_FILE_SIZE = 50 * 1024 * 1024
 29
 30# Default exec timeout in seconds.
 31DEFAULT_EXEC_TIMEOUT = 120
 32
 33# Relay WebSocket endpoint.
 34DEFAULT_RELAY_HOST = "portal.plainframework.com"
 35RELAY_PATH = "/__portal__"
 36
 37# Protocol version — bumped on breaking changes.
 38PROTOCOL_VERSION = 1
 39
 40
 41def make_relay_url(relay_host: str, channel: str, side: str) -> str:
 42    """Build the relay WebSocket URL."""
 43    scheme = "ws" if relay_host.startswith(("localhost", "127.0.0.1")) else "wss"
 44    return f"{scheme}://{relay_host}{RELAY_PATH}?v={PROTOCOL_VERSION}&channel={channel}&side={side}"
 45
 46
 47def chunk_count(file_size: int) -> int:
 48    """Number of chunks needed to transfer a file."""
 49    return max(1, math.ceil(file_size / FILE_CHUNK_SIZE))
 50
 51
 52def make_exec(
 53    code: str, json_output: bool = False, timeout: int = DEFAULT_EXEC_TIMEOUT
 54) -> dict:
 55    return {
 56        "type": "exec",
 57        "code": code,
 58        "json_output": json_output,
 59        "timeout": timeout,
 60    }
 61
 62
 63def make_exec_stdout(data: str) -> dict:
 64    return {"type": "exec_stdout", "data": data}
 65
 66
 67def make_exec_result(
 68    return_value: str | None, error: str | None, stdout: str = ""
 69) -> dict:
 70    return {
 71        "type": "exec_result",
 72        "stdout": stdout,
 73        "return_value": return_value,
 74        "error": error,
 75    }
 76
 77
 78def make_error(error: str) -> dict:
 79    return {"type": "error", "error": error}
 80
 81
 82def make_file_pull(remote_path: str) -> dict:
 83    return {"type": "file_pull", "remote_path": remote_path}
 84
 85
 86def make_file_data(name: str, chunk: int, chunks: int, data: bytes) -> dict:
 87    return {
 88        "type": "file_data",
 89        "name": name,
 90        "chunk": chunk,
 91        "chunks": chunks,
 92        "data": base64.b64encode(data).decode("ascii"),
 93    }
 94
 95
 96def make_file_push(remote_path: str, chunk: int, chunks: int, data: bytes) -> dict:
 97    return {
 98        "type": "file_push",
 99        "remote_path": remote_path,
100        "chunk": chunk,
101        "chunks": chunks,
102        "data": base64.b64encode(data).decode("ascii"),
103    }
104
105
106def make_file_push_result(path: str, total_bytes: int) -> dict:
107    return {"type": "file_push_result", "path": path, "bytes": total_bytes}
108
109
110def make_ping() -> dict:
111    return {"type": "ping"}
112
113
114def make_pong() -> dict:
115    return {"type": "pong"}