chuanqi/apilot/wire.py

"""协议线层: 24B头(9×i32)+体, 全帧 Encode6BitBuf。提炼自 protocol_bot/cq2_pure_enter.py(实测验证)。"""
from __future__ import annotations
import socket
import struct
import sys
import time
from dataclasses import dataclass

sys.path.insert(0, "protocol_bot")
from encode6bitbuf_final import encode6, decode6  # type: ignore

HDR = struct.Struct("<9i")  # recog,msg,p1,p2,p3,zero?,len,tail → 对齐 cq2_pure_enter.send()


@dataclass
class Frame:
    recog: int
    msg_id: int
    p1: int
    p2: int
    p3: int
    tail: int
    body: bytes = b""


def pack_frame(f: Frame) -> bytes:
    head = HDR.pack(0, f.msg_id, f.recog, f.p1, f.p2, f.p3, 0, len(f.body), f.tail)
    return encode6(head) + f.body


def unpack_stream(buf: bytes) -> tuple[list[Frame], bytes]:
    out, pos = [], 0
    while pos + 48 <= len(buf):
        try:
            m = HDR.unpack(decode6(buf[pos:pos + 48]))
        except Exception:
            break
        body_len = m[7]
        if pos + 48 + body_len > len(buf):
            break
        out.append(Frame(recog=m[2], msg_id=m[1], p1=m[3], p2=m[4], p3=m[5],
                         tail=m[8], body=buf[pos + 48:pos + 48 + body_len]))
        pos += 48 + body_len
    return out, buf[pos:]


class Wire:
    """同步 TCP 线客户端。持跨 recv 调用的内部 buffer:半包 rest 留存到下次。"""

    def __init__(self, host: str, port: int, timeout: float = 10) -> None:
        self.sock = socket.create_connection((host, port), timeout=timeout)
        self._buf: bytes = b""

    def send(self, f: Frame) -> None:
        self.sock.sendall(pack_frame(f))

    def recv(self, deadline_s: float) -> list[Frame]:
        deadline = time.time() + deadline_s
        while time.time() < deadline:
            self.sock.settimeout(max(0.2, deadline - time.time()))
            try:
                d = self.sock.recv(65536)
                if not d:
                    break
                self._buf += d
            except (socket.timeout, ConnectionResetError, OSError):
                break
        frames, self._buf = unpack_stream(self._buf)
        return frames

    def close(self) -> None:
        try:
            self.sock.close()
        except OSError:
            pass