DeepRoboticsLab-S10-proto/tests/fake_robot.py

#!/usr/bin/env python3
"""本地假机器人：用于在没有真机时跑通 demo。

启动一个 TLS 服务端，按《山猫S10 软件开发指南》的协议响应控制指令，并
周期性上报基础状态。证书用 openssl 现场生成（自签名，客户端需 verify=False）。

用法::

    uv run python tests/fake_robot.py --port 30003

    # 另一个终端
    uv run python demo_stand_rotate_lie.py --host 127.0.0.1 --seconds 3

行为模拟：
* 心跳 -> ErrorCode 0，并开始 5Hz 上报 BasicStatus
* 使用模式切换 -> 更新 ControlUsageMode
* 站立 -> 0.8 秒后 MotionState 变为 17（RL 控制）
* 趴下 -> 0.8 秒后 MotionState 变为 4
* 步态切换 -> 更新 Gait
* 轴指令/速度指令 -> 不回包，仅记录最后一帧
"""

from __future__ import annotations

import argparse
import logging
import socket
import ssl
import subprocess
import sys
import tempfile
import threading
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import s10_robot as s10

log = logging.getLogger("fake")


def make_cert(directory: Path) -> tuple[Path, Path]:
    crt = directory / "server.crt"
    key = directory / "server.key"
    subprocess.run(
        [
            "openssl", "req", "-x509", "-newkey", "rsa:2048", "-nodes",
            "-keyout", str(key), "-out", str(crt), "-days", "1",
            "-subj", "/CN=server.local",
            "-addext", "subjectAltName=DNS:server.local,IP:127.0.0.1",
        ],
        check=True,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    return crt, key


class RobotState:
    def __init__(self):
        self.lock = threading.Lock()
        self.motion_state = 0
        self.gait = 0x1001
        self.mode = 0
        self.hes = 0
        self.last_axis = {}
        self.axis_frames = 0
        self.reporting = False

    def snapshot(self) -> dict:
        with self.lock:
            return {
                "MotionState": self.motion_state,
                "Gait": self.gait,
                "ControlUsageMode": self.mode,
                "HES": self.hes,
                "PowerManagement": 0,
                "Sleep": 0,
                "Model": "CD1",
                "Version": "STD",
                "Charge": 0,
            }


def handle_command(state: RobotState, frame: s10.Frame) -> dict | None:
    if frame.type == s10.T_HEARTBEAT and frame.command == s10.C_HEARTBEAT:
        state.reporting = True
        return {"ErrorCode": 0}

    if frame.type == s10.T_MODE and frame.command == s10.C_MODE_SET:
        with state.lock:
            state.mode = int(frame.items.get("Mode", 0))
        log.info("使用模式 -> %s", state.mode)
        return {"ErrorCode": 0}

    if frame.type == s10.T_MOTION and frame.command == s10.C_MOTION_STATE:
        target = int(frame.items.get("MotionParam", 0))
        log.info("运动状态指令 -> %s", target)

        def transition():
            time.sleep(0.8)
            with state.lock:
                # ASDU 站立完成后会自动进入 RL 控制状态
                state.motion_state = 17 if target == 1 else target
            log.info("运动状态已变为 %s", state.motion_state)

        threading.Thread(target=transition, daemon=True).start()
        return {"ErrorCode": 0}

    if frame.type == s10.T_MOTION and frame.command == s10.C_GAIT:
        with state.lock:
            state.gait = int(frame.items.get("GaitParam", 0))
        log.info("步态 -> 0x%04X", state.gait)
        return {"ErrorCode": 0}

    if frame.type == s10.T_MOTION and frame.command in (s10.C_AXIS, s10.C_VELOCITY):
        with state.lock:
            state.last_axis = dict(frame.items)
            state.axis_frames += 1
            n = state.axis_frames
        if n % 10 == 1:
            log.info("轴指令 #%d: %s", n, state.last_axis)
        return None  # 轴指令无响应帧

    if frame.type == s10.T_MODE and frame.command == s10.C_SLEEP:
        return {"ErrorCode": 0}

    return {"ErrorCode": 0}


def serve_connection(conn: ssl.SSLSocket, addr, state: RobotState, stop: threading.Event,
                     send_lock: threading.Lock) -> None:
    log.info("客户端接入 %s", addr)
    buf = bytearray()
    try:
        while not stop.is_set():
            conn.settimeout(0.2)
            try:
                chunk = conn.recv(65535)
            except socket.timeout:
                continue
            if not chunk:
                break
            buf.extend(chunk)
            for frame in s10.parse_frames(buf):
                response = handle_command(state, frame)
                if response is not None:
                    reply = s10.build_frame(
                        s10.build_asdu(frame.type, frame.command, response),
                        report_id=frame.report_id,
                    )
                    with send_lock:
                        conn.sendall(reply)
    except (OSError, ssl.SSLError) as exc:
        log.info("连接结束: %s", exc)
    finally:
        try:
            conn.close()
        except OSError:
            pass
        log.info("客户端断开")


def reporter(conn: ssl.SSLSocket, state: RobotState, stop: threading.Event,
             send_lock: threading.Lock, period: float = 0.2) -> None:
    next_t = time.monotonic()
    while not stop.is_set():
        if not state.reporting:
            stop.wait(0.05)
            next_t = time.monotonic()
            continue
        payload = s10.build_asdu(
            s10.T_HEARTBEAT, s10.C_STATUS, {"BasicStatus": state.snapshot()}
        )
        try:
            with send_lock:
                conn.sendall(s10.build_frame(payload))
        except (OSError, ssl.SSLError):
            return
        next_t += period
        stop.wait(max(0.0, next_t - time.monotonic()))


def main() -> int:
    parser = argparse.ArgumentParser(description="山猫 S10 本地假机器人")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=30003)
    parser.add_argument("--once", action="store_true", help="处理一个连接后退出")
    parser.add_argument("-v", "--verbose", action="store_true")
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s.%(msecs)03d %(levelname)-7s %(name)s: %(message)s",
        datefmt="%H:%M:%S",
    )

    with tempfile.TemporaryDirectory() as tmp:
        crt, key = make_cert(Path(tmp))
        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        ctx.load_cert_chain(crt, key)
        ctx.verify_mode = ssl.CERT_NONE  # 不校验客户端证书，便于本地联调

        state = RobotState()
        stop = threading.Event()

        with socket.create_server((args.host, args.port), reuse_port=False) as server:
            log.info("假机器人监听 %s:%d（TLS）", args.host, args.port)
            while not stop.is_set():
                raw, addr = server.accept()
                try:
                    conn = ctx.wrap_socket(raw, server_side=True)
                except ssl.SSLError as exc:
                    log.warning("TLS 握手失败: %s", exc)
                    raw.close()
                    continue
                send_lock = threading.Lock()
                rep = threading.Thread(
                    target=reporter, args=(conn, state, stop, send_lock), daemon=True
                )
                rep.start()
                serve_connection(conn, addr, state, stop, send_lock)
                rep.join(timeout=1.0)
                if args.once:
                    break
    return 0


if __name__ == "__main__":
    sys.exit(main())