xiaoli-cli-plugins/onebot/onebot.py

"""
OneBot v11 插件 — 让 xiaoli-cli 成为 QQ 机器人
通过 WebSocket 连接 OneBot 实现（go-cqhttp / Lagrange / LLOneBot 等）

功能:
- WebSocket 正向/反向连接
- 接收私聊和群消息，转发给 AI 处理
- AI 回复自动发回 QQ
- 支持 @机器人 触发 / 全部消息 / 前缀触发
- 群消息可配置是否响应
"""

import json
import time
import threading
import re
import os

try:
    import websocket
    HAS_WS = True
except ImportError:
    HAS_WS = False


class Liugin:
    """OneBot v11 QQ 机器人插件"""

    def __init__(self):
        self.ws = None
        self.connected = False
        self.running = False
        self.ws_url = ""
        self.mode = ""  # "forward" 正向 / "reverse" 反向

        # 配置
        self.trigger_mode = "at"  # at: 只响应@ / prefix: 前缀触发 / all: 全部消息
        self.trigger_prefix = "/"
        self.group_whitelist = []  # 群白名单，空=全部响应
        self.admin_users = []  # 管理员 QQ 号
        self.listen_port = 8765  # 反向 WebSocket 监听端口
        self.self_id = ""  # 机器人自身 QQ 号

        # 消息队列 (用于 AI 回复)
        self._pending_replies = {}  # {message_id: {"context": ..., "reply": ...}}
        self._msg_counter = 0

        self.cli = None
        self._listener_thread = None

        self.usage = """OneBot v11 QQ 机器人插件 使用说明：

/onebot connect <ws_url>     连接 OneBot (正向WebSocket)
/onebot listen [port]        启动反向WebSocket监听 (默认8765)
/onebot disconnect           断开连接
/onebot status               查看连接状态
/onebot config <key> <val>   配置:
    trigger_mode at|prefix|all  触发模式
    trigger_prefix <前缀>       前缀触发的前缀
    group_whitelist <群号>      添加群白名单
    admin <QQ号>                添加管理员
    self_id <QQ号>              设置机器人QQ号
/onebot send <type> <id> <消息>  手动发消息
    type: private/group
    id: QQ号/群号
/onebot help                 显示帮助
"""

    def set_cli(self, cli):
        self.cli = cli
        if hasattr(cli, 'register_liugin_command'):
            cli.register_liugin_command('onebot', self._cmd_handler)

    def _cmd_handler(self, args):
        parts = args.strip().split(maxsplit=1)
        if not parts:
            return self.handle("help")
        return self.handle(args)

    def get_tool_info(self):
        return {
            "name": "onebot",
            "description": "OneBot v11 QQ 机器人插件，连接 go-cqhttp/Lagrange/LLOneBot 等，让 AI 接管 QQ 消息",
            "keywords": [
                "onebot", "qq", "机器人", "bot", "go-cqhttp",
                "lagrange", "消息", "群聊", "私聊", "cqhttp"
            ],
            "usage": self.usage
        }

    def get_mcp_definition(self):
        return {
            "name": "onebot",
            "description": "OneBot v11 QQ bot plugin",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": ["connect", "listen", "disconnect", "status", "config", "send"]
                    },
                    "args": {"type": "string"}
                },
                "required": ["action"]
            }
        }

    def convert_mcp_args(self, arguments):
        return f"{arguments.get('action', '')} {arguments.get('args', '')}"

    # ── WebSocket 连接 ──

    def _on_message(self, ws, raw):
        try:
            data = json.loads(raw)
        except json.JSONDecodeError:
            return
        self._dispatch(data)

    def _on_error(self, ws, error):
        print(f"\n[OneBot] ❌ WebSocket 错误: {error}")

    def _on_close(self, ws, code, reason):
        self.connected = False
        self.running = False
        print(f"\n[OneBot] 🔌 连接断开 (code={code})")

    def _on_open(self, ws):
        self.connected = True
        self.running = True
        print(f"\n[OneBot] ✅ 已连接到 OneBot 服务")

    def _connect_forward(self, url):
        """正向 WebSocket 连接"""
        if not HAS_WS:
            return "❌ 需要 websocket-client: pip install websocket-client"

        self.ws_url = url
        self.mode = "forward"

        try:
            self.ws = websocket.WebSocketApp(
                url,
                on_message=self._on_message,
                on_error=self._on_error,
                on_close=self._on_close,
                on_open=self._on_open
            )
            self._listener_thread = threading.Thread(target=self.ws.run_forever, daemon=True)
            self._listener_thread.start()

            for _ in range(30):
                if self.connected:
                    break
                time.sleep(0.1)

            if self.connected:
                return f"✅ 已连接到 {url}"
            return f"❌ 连接超时: {url}"
        except Exception as e:
            return f"❌ 连接失败: {e}"

    def _start_reverse_listener(self, port):
        """反向 WebSocket 监听"""
        try:
            from http.server import HTTPServer, BaseHTTPRequestHandler
            import hashlib
            import struct
        except ImportError:
            return "❌ 缺少标准库组件"

        self.listen_port = port
        self.mode = "reverse"
        plugin_ref = self

        class ReverseWSHandler(BaseHTTPRequestHandler):
            def do_GET(self):
                # WebSocket 握手
                if self.headers.get('Upgrade', '').lower() != 'websocket':
                    self.send_error(400)
                    return

                key = self.headers.get('Sec-WebSocket-Key', '')
                accept = hashlib.sha1(
                    (key + "258EAFA5-E914-47DA-95CA-5AB5DC65C740").encode()
                ).digest()

                self.send_response(101, 'Switching Protocols')
                self.send_header('Upgrade', 'websocket')
                self.send_header('Connection', 'Upgrade')
                self.send_header('Sec-WebSocket-Accept',
                                 __import__('base64').b64encode(accept).decode())
                self.end_headers()

                # 简化的 WebSocket 帧处理
                plugin_ref._handle_reverse_ws(self.rfile, self.wfile, self)

            def log_message(self, format, *args):
                pass  # 静默日志

        class ReverseHTTPServer(HTTPServer):
            allow_reuse_address = True

        def _run_server():
            try:
                server = ReverseHTTPServer(('0.0.0.0', port), ReverseWSHandler)
                self.connected = True
                self.running = True
                print(f"\n[OneBot] ✅ 反向 WebSocket 监听中: ws://0.0.0.0:{port}")
                server.serve_forever()
            except Exception as e:
                print(f"\n[OneBot] ❌ 监听失败: {e}")

        self._listener_thread = threading.Thread(target=_run_server, daemon=True)
        self._listener_thread.start()
        return f"✅ 反向 WebSocket 已启动: ws://0.0.0.0:{port}"

    def _handle_reverse_ws(self, rfile, wfile, handler):
        """处理反向 WebSocket 连接"""
        self.connected = True
        self.running = True
        print(f"\n[OneBot] ✅ OneBot 已连接")

        try:
            while self.running:
                # 读取 WebSocket 帧
                header = rfile.read(2)
                if len(header) < 2:
                    break

                opcode = header[0] & 0x0F
                masked = bool(header[1] & 0x80)
                length = header[1] & 0x7F

                if length == 126:
                    length = struct.unpack('>H', rfile.read(2))[0]
                elif length == 127:
                    length = struct.unpack('>Q', rfile.read(8))[0]

                mask_key = rfile.read(4) if masked else b''
                payload = rfile.read(length)

                if masked:
                    payload = bytes(
                        b ^ mask_key[i % 4]
                        for i, b in enumerate(payload)
                    )

                if opcode == 1:  # Text
                    try:
                        data = json.loads(payload.decode('utf-8'))
                        self._dispatch(data)
                    except json.JSONDecodeError:
                        pass
                elif opcode == 8:  # Close
                    break
                elif opcode == 9:  # Ping
                    # 发 Pong
                    self._ws_send_pong(wfile, payload)

        except Exception as e:
            print(f"\n[OneBot] 连接断开: {e}")
        finally:
            self.connected = False

    def _ws_send_pong(self, wfile, payload):
        """发送 Pong 帧"""
        try:
            frame = bytes([0x8A, len(payload)]) + payload
            wfile.write(frame)
            wfile.flush()
        except Exception:
            pass

    # ── 消息分发 ──

    def _dispatch(self, data):
        """分发 OneBot 事件"""
        post_type = data.get("post_type", "")

        if post_type == "message":
            self._handle_message(data)
        elif post_type == "meta_event":
            sub_type = data.get("meta_event_type", "")
            if sub_type == "heartbeat":
                pass
            elif sub_type == "lifecycle":
                if data.get("sub_type") == "connect":
                    self.self_id = str(data.get("self_id", ""))
                    print(f"\n[OneBot] 🤖 机器人已上线: {self.self_id}")
        elif post_type == "request":
            self._handle_request(data)

    def _handle_message(self, data):
        """处理收到的消息"""
        message_type = data.get("message_type", "")
        raw_msg = data.get("raw_message", data.get("message", ""))
        user_id = str(data.get("user_id", ""))
        group_id = str(data.get("group_id", ""))
        message_id = data.get("message_id", 0)
        sender = data.get("sender", {})
        nickname = sender.get("nickname", user_id)

        # 群白名单检查
        if message_type == "group" and self.group_whitelist:
            if group_id not in self.group_whitelist:
                return

        # 解析纯文本
        text = self._extract_text(raw_msg)

        # 判断是否需要响应
        should_respond = False

        if message_type == "private":
            should_respond = True
        elif message_type == "group":
            if self.trigger_mode == "all":
                should_respond = True
            elif self.trigger_mode == "at":
                if f"[CQ:at,qq={self.self_id}]" in raw_msg:
                    should_respond = True
                    text = text.replace(f"@{self.self_id}", "").strip()
            elif self.trigger_mode == "prefix":
                if text.startswith(self.trigger_prefix):
                    should_respond = True
                    text = text[len(self.trigger_prefix):].strip()

        if not should_respond or not text:
            return

        # 打印收到的消息
        if message_type == "group":
            print(f"\n[OneBot] 💬 [{group_id}] {nickname}: {text[:80]}")
        else:
            print(f"\n[OneBot] 💬 {nickname}: {text[:80]}")

        # 交给 AI 处理
        self._process_with_ai(text, message_type, user_id, group_id, message_id, nickname)

    def _extract_text(self, raw_msg):
        """从 CQ 码消息中提取纯文本"""
        # 移除 CQ 码
        text = re.sub(r'\[CQ:[^\]]+\]', '', raw_msg)
        return text.strip()

    def _handle_request(self, data):
        """处理请求事件（好友申请、群邀请等）"""
        request_type = data.get("request_type", "")
        if request_type == "friend":
            # 自动同意好友请求
            self._send_api("set_friend_add_request", {
                "flag": data.get("flag", ""),
                "approve": True
            })
            print(f"\n[OneBot] ✅ 已自动同意好友请求: {data.get('user_id')}")

    def _process_with_ai(self, text, msg_type, user_id, group_id, msg_id, nickname):
        """将消息交给 AI 处理并发回回复"""
        if not self.cli:
            return

        # 构建上下文
        if msg_type == "group":
            context = f"[QQ群消息] 群{group_id} 用户{nickname}({user_id}): {text}"
        else:
            context = f"[QQ私聊] 用户{nickname}({user_id}): {text}"

        def _ai_worker():
            try:
                # 调用 AI 处理
                response = self.cli._generate_response_with_animation(context)
                if response:
                    # 清理回复
                    reply = self._clean_reply(response)
                    if reply:
                        if msg_type == "group":
                            self._send_group_msg(group_id, reply, at_user=user_id)
                        else:
                            self._send_private_msg(user_id, reply)
            except Exception as e:
                print(f"\n[OneBot] ❌ AI 处理失败: {e}")

        threading.Thread(target=_ai_worker, daemon=True).start()

    def _clean_reply(self, response):
        """清理 AI 回复，去掉 JSON 指令等"""
        # 去掉工具调用 JSON
        response = re.sub(r'\{"action":\s*"use_tool"[^}]*\}', '', response)
        response = re.sub(r'\{"action":\s*"continue"[^}]*\}', '', response)
        response = re.sub(r'<tool_call>[^<]*', '', response)
        response = response.strip()

        # 限制长度
        if len(response) > 2000:
            response = response[:1997] + "..."

        return response

    # ── OneBot API 调用 ──

    def _send_api(self, action, params=None):
        """调用 OneBot API"""
        if not self.connected or not self.ws:
            return None

        self._msg_counter += 1
        payload = {
            "action": action,
            "params": params or {},
            "echo": self._msg_counter
        }

        try:
            if self.mode == "forward":
                self.ws.send(json.dumps(payload))
            return self._msg_counter
        except Exception as e:
            print(f"\n[OneBot] ❌ API 调用失败: {e}")
            return None

    def _send_private_msg(self, user_id, message):
        """发送私聊消息"""
        return self._send_api("send_private_msg", {
            "user_id": int(user_id),
            "message": message
        })

    def _send_group_msg(self, group_id, message, at_user=None):
        """发送群消息"""
        if at_user:
            message = f"[CQ:at,qq={at_user}] {message}"
        return self._send_api("send_group_msg", {
            "group_id": int(group_id),
            "message": message
        })

    # ── 命令入口 ──

    def handle(self, args):
        parts = args.strip().split(maxsplit=1)
        if not parts:
            return self.usage

        action = parts[0].lower()
        rest = parts[1] if len(parts) > 1 else ""

        if action == "help":
            return self.usage
        elif action == "connect":
            return self._do_connect(rest.strip())
        elif action == "listen":
            port = int(rest.strip()) if rest.strip().isdigit() else 8765
            return self._start_reverse_listener(port)
        elif action == "disconnect":
            return self._do_disconnect()
        elif action == "status":
            return self._do_status()
        elif action == "config":
            return self._do_config(rest)
        elif action == "send":
            return self._do_send(rest)
        else:
            return f"未知命令: {action}\n{self.usage}"

    def _do_connect(self, url):
        if not url:
            return "请提供 WebSocket 地址: /onebot connect ws://127.0.0.1:6700"
        return self._connect_forward(url)

    def _do_disconnect(self):
        self.running = False
        if self.ws:
            try:
                self.ws.close()
            except Exception:
                pass
        self.connected = False
        return "已断开连接"

    def _do_status(self):
        lines = ["🤖 OneBot 状态:"]
        lines.append(f"  连接: {'✅ 已连接' if self.connected else '❌ 未连接'}")
        lines.append(f"  模式: {self.mode or '未设置'}")
        if self.ws_url:
            lines.append(f"  地址: {self.ws_url}")
        lines.append(f"  端口: {self.listen_port}")
        lines.append(f"  机器人QQ: {self.self_id or '未知'}")
        lines.append(f"  触发模式: {self.trigger_mode}")
        if self.trigger_mode == "prefix":
            lines.append(f"  触发前缀: {self.trigger_prefix}")
        if self.group_whitelist:
            lines.append(f"  群白名单: {', '.join(self.group_whitelist)}")
        if self.admin_users:
            lines.append(f"  管理员: {', '.join(self.admin_users)}")
        return "\n".join(lines)

    def _do_config(self, args):
        parts = args.split(maxsplit=1)
        if len(parts) < 2:
            return "用法: /onebot config <key> <value>\n可用: trigger_mode, trigger_prefix, group_whitelist, admin, self_id"

        key, val = parts[0].lower(), parts[1].strip()

        if key == "trigger_mode":
            if val not in ("at", "prefix", "all"):
                return "触发模式: at / prefix / all"
            self.trigger_mode = val
            return f"✅ 触发模式: {val}"
        elif key == "trigger_prefix":
            self.trigger_prefix = val
            return f"✅ 触发前缀: {val}"
        elif key == "group_whitelist":
            if val not in self.group_whitelist:
                self.group_whitelist.append(val)
            return f"✅ 群白名单: {self.group_whitelist}"
        elif key == "admin":
            if val not in self.admin_users:
                self.admin_users.append(val)
            return f"✅ 管理员: {self.admin_users}"
        elif key == "self_id":
            self.self_id = val
            return f"✅ 机器人QQ: {val}"
        else:
            return f"未知配置: {key}"

    def _do_send(self, args):
        parts = args.split(maxsplit=2)
        if len(parts) < 3:
            return "用法: /onebot send <private|group> <ID> <消息>"

        msg_type, target_id, message = parts[0].lower(), parts[1], parts[2]

        if msg_type in ("private", "p", "私聊"):
            self._send_private_msg(target_id, message)
            return f"✅ 已发送私聊消息给 {target_id}"
        elif msg_type in ("group", "g", "群"):
            self._send_group_msg(target_id, message)
            return f"✅ 已发送群消息到 {target_id}"
        else:
            return "类型: private / group"