ros_contorl/scripts/rdk_switch_tool.py

#!/usr/bin/env python3
"""把 Flexiv 当前工具（TCP + 负载）切换成 Studio 里配置好的名字。

必须在 RDK 的 Python（rdk_py310）下运行。机器人必须处于 IDLE；如果
ros2_control 已经占住了 RDK 连接，请先停掉 run_real.sh。

  /home/cc/rdk_py310/bin/python ~/ros2/scripts/rdk_switch_tool.py --list
  /home/cc/rdk_py310/bin/python ~/ros2/scripts/rdk_switch_tool.py --name Flange
  /home/cc/rdk_py310/bin/python ~/ros2/scripts/rdk_switch_tool.py   # site.yaml robot.tool_name
"""

from __future__ import annotations

import argparse
import sys
import time


def _single_arm_groups(robot):
    """RDK 1.9 的 Tool API 是单臂的（没有 JointGroup）。老版本用 groups。"""
    info = robot.info()
    groups = getattr(info, "single_arm_groups", None) or {}
    if isinstance(groups, dict):
        keys = list(groups.keys())
    else:
        keys = list(groups)
    if keys:
        return keys
    try:
        import flexivrdk

        if hasattr(flexivrdk, "JointGroup"):
            return [flexivrdk.JointGroup.ARM_1]
    except Exception:
        pass
    return [None]  # 哨兵：调用 Tool.name()/Switch() 时不带 group


def _tool_name(tool, group):
    if group is None:
        return tool.name()
    return tool.name(group)


def _tool_params(tool, group):
    if group is None:
        return tool.params()
    return tool.params(group)


def _tool_switch(tool, group, target: str) -> None:
    if group is None:
        tool.Switch(target)
        return
    tool.Switch(group, target)


def _group_name(group) -> str:
    if group is None:
        return "单臂"
    try:
        import flexivrdk

        return str(flexivrdk.kJointGroupNames[group])
    except Exception:
        return str(group)


def _fmt_params(params) -> str:
    mass = getattr(params, "mass", None)
    com = list(getattr(params, "CoM", []) or [])
    tcp = list(getattr(params, "tcp_location", []) or [])
    com_s = ",".join(f"{float(v):.4f}" for v in com[:3]) if com else "?"
    tcp_s = ",".join(f"{float(v):.4f}" for v in tcp[:3]) if tcp else "?"
    return f"mass={mass} kg  CoM=[{com_s}] m  TCP_xyz=[{tcp_s}] m"


def _ensure_idle(robot, mode) -> None:
    try:
        robot.SwitchMode(mode.IDLE)
        return
    except Exception:
        pass
    if hasattr(robot, "Enable"):
        robot.Enable()
    else:
        robot.ServoOn()
    for _ in range(40):
        if robot.operational():
            break
        time.sleep(0.25)
    if not robot.operational():
        raise RuntimeError(
            "机器人未处于 operational 状态 —— 请在示教器切 AUTO Remote 并使能，松开急停"
        )
    robot.SwitchMode(mode.IDLE)


def _resolve_name(tool, want: str) -> str:
    want = (want or "").strip()
    if not want:
        return ""
    names = list(tool.list() or [])
    if want in names:
        return want
    lower = {n.lower(): n for n in names}
    hit = lower.get(want.lower())
    if hit:
        return hit
    raise SystemExit(
        "工具 %r 不在 Studio 工具池里。\n  可选: %s"
        % (want, ", ".join(names) or "（空）")
    )


def _iface_ipv4s(name: str) -> list:
    """FastDDS 白名单要的是 IPv4 地址，而不是 enp4s0 这类网卡名。"""
    import socket
    import fcntl
    import struct

    name = (name or "").strip()
    if not name:
        return []
    if name.count(".") == 3:
        return [name]
    addrs = []
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        ifreq = struct.pack("256s", name.encode("utf-8")[:15])
        addrs.append(socket.inet_ntoa(fcntl.ioctl(s.fileno(), 0x8915, ifreq)[20:24]))
        s.close()
    except OSError:
        pass
    try:
        import subprocess

        out = subprocess.check_output(
            ["ip", "-4", "-o", "addr", "show", "dev", name],
            text=True,
            stderr=subprocess.DEVNULL,
        )
        for line in out.splitlines():
            parts = line.split()
            if "inet" in parts:
                cidr = parts[parts.index("inet") + 1]
                addrs.append(cidr.split("/")[0])
    except Exception:
        pass
    # 去重，保持原有顺序
    seen = set()
    uniq = []
    for a in addrs:
        if a and a not in seen:
            seen.add(a)
            uniq.append(a)
    return uniq


def _rdk_whitelist(iface: str) -> list[list[str]]:
    """RDK 手册要的是「连机器人的那张网卡」。enp4s0 上还有办公网 IP，不能整张卡的地址都塞进白名单。"""
    import os

    lists: list[list[str]] = []
    pc = (os.environ.get("SITE_PC_IP") or "").strip()
    if pc:
        lists.append([pc])
    if iface:
        lists.append([iface])
        robot_lan = [a for a in _iface_ipv4s(iface) if a.startswith("192.168.2.")]
        if robot_lan:
            lists.append(robot_lan)
    # 去重
    seen: set[tuple[str, ...]] = set()
    out: list[list[str]] = []
    for item in lists:
        key = tuple(item)
        if key not in seen:
            seen.add(key)
            out.append(item)
    return out


def connect_rdk(sn: str, iface: str = ""):
    import flexivrdk

    for listen in _rdk_whitelist(iface):
        print("RDK 白名单:", ", ".join(listen))
        for args in ((sn, listen), (sn, listen, True), (sn, listen, True, False)):
            try:
                return flexivrdk.Robot(*args)
            except TypeError:
                continue
            except RuntimeError as exc:
                msg = str(exc)
                if "whitelist" in msg.lower() or "filtered out" in msg.lower():
                    print("  白名单被拒，试下一个")
                    break
                if "Failed to connect" in msg or "failed to connect" in msg.lower():
                    print("  用这个白名单搜索超时")
                    break
                raise
    print("不带 iface 过滤再试一次")
    return flexivrdk.Robot(sn)


def main() -> int:
    ap = argparse.ArgumentParser(description="切换 Flexiv Studio 工具 / TCP")
    ap.add_argument("--sn", default="", help="机械臂序列号（默认取 site.yaml）")
    ap.add_argument("--iface", default="", help="RDK 用的网卡白名单（如 enp4s0）")
    ap.add_argument(
        "--name",
        default="",
        help="Studio 工具名（Flange = 不装工具）。留空 = 用 site.yaml 的 robot.tool_name",
    )
    ap.add_argument("--list", action="store_true", help="列出工具后退出")
    args = ap.parse_args()

    sn = (args.sn or "").strip()
    want = (args.name or "").strip()
    if not sn:
        print("错误：必须提供 --sn（请通过 run_real.sh / rdk_switch_tool.sh 运行）", file=sys.stderr)
        return 2

    import flexivrdk

    mode = flexivrdk.Mode
    iface = (args.iface or "").strip()
    print(f"=== RDK 工具  SN={sn}  iface={iface or 'auto'} ===")
    robot = None
    try:
        robot = connect_rdk(sn, iface)
        if robot.fault():
            print("检测到故障，正在 ClearFault ...")
            if not robot.ClearFault():
                print("错误：故障无法清除", file=sys.stderr)
                return 1
        tool = flexivrdk.Tool(robot)
        names = list(tool.list() or [])
        groups = _single_arm_groups(robot)
        print("Studio 工具:", ", ".join(names) or "（无）")
        for g in groups:
            cur = _tool_name(tool, g)
            extra = ""
            try:
                extra = "  " + _fmt_params(_tool_params(tool, g))
            except Exception:
                pass
            print(f"  当前 [{_group_name(g)}]: {cur}{extra}")

        if args.list:
            return 0
        if not want:
            print("tool_name 为空 —— 保持当前工具（要切换请在 site.yaml 里设置 robot.tool_name）")
            print("RDK 正常")
            return 0

        target = _resolve_name(tool, want)
        already = all(_tool_name(tool, g) == target for g in groups)
        if already:
            print(f"已经处于 [{target}]")
            print("RDK 正常")
            return 0

        print(f"正在切换到 → [{target}]（需要 IDLE）")
        _ensure_idle(robot, mode)
        for g in groups:
            _tool_switch(tool, g, target)
        for g in groups:
            cur = _tool_name(tool, g)
            extra = ""
            try:
                extra = "  " + _fmt_params(_tool_params(tool, g))
            except Exception:
                pass
            print(f"  现在 [{_group_name(g)}]: {cur}{extra}")
            if cur != target:
                print(f"错误：切换失败，当前仍是 {cur!r}", file=sys.stderr)
                return 1
        print("RDK 正常")
        return 0
    except Exception as exc:
        print(f"错误：{exc}", file=sys.stderr)
        print("  请先停掉 run_real.sh（RDK 是独占的）。", file=sys.stderr)
        print("  示教器：切 AUTO Remote，并使能 Ethernet/RDK。", file=sys.stderr)
        return 1
    finally:
        robot = None


if __name__ == "__main__":
    raise SystemExit(main())