multinetspeeder/protocol.py

from __future__ import annotations

import json
import struct
from dataclasses import dataclass
from typing import Any

HEADER = struct.Struct("!BIIQII")
AUTH = 1
TCP_OPEN = 2
TCP_STATUS = 3
TCP_DATA = 4
TCP_CLOSE = 5
UDP_SEND = 6
UDP_RECV = 7
PING = 8
PONG = 9

STATUS_OK = 0
STATUS_ERR = 1


@dataclass
class Frame:
    kind: int
    session_id: int
    stream_id: int
    seq: int
    packet_id: int
    payload: bytes = b""


def encode_json(data: Any) -> bytes:
    return json.dumps(data, ensure_ascii=False, separators=(",", ":")).encode("utf-8")


def decode_json(data: bytes) -> Any:
    return json.loads(data.decode("utf-8"))


async def write_frame(writer, frame: Frame) -> None:
    writer.write(HEADER.pack(frame.kind, frame.session_id, frame.stream_id, frame.seq, frame.packet_id, len(frame.payload)))
    if frame.payload:
        writer.write(frame.payload)
    await writer.drain()


async def read_frame(reader) -> Frame:
    header = await reader.readexactly(HEADER.size)
    kind, session_id, stream_id, seq, packet_id, length = HEADER.unpack(header)
    payload = await reader.readexactly(length) if length else b""
    return Frame(kind=kind, session_id=session_id, stream_id=stream_id, seq=seq, packet_id=packet_id, payload=payload)