117 lines
3.9 KiB
Python
117 lines
3.9 KiB
Python
"""一次性工具:捕获目标群的 chat_id 并写回 Secret Block / 发测试消息。
|
|
|
|
用法:
|
|
# 1. 初始化 Secret Block(凭据可从环境变量 WECOM_BOT_ID / WECOM_SECRET 读取)
|
|
uv run python scripts/wecom_capture_chatid.py --init
|
|
|
|
# 2. 监听模式:运行后在目标群里 @机器人 说任意一句话,即可捕获 chat_id
|
|
uv run python scripts/wecom_capture_chatid.py --listen
|
|
|
|
# 3. 捕获到 chat_id 后写回 Secret Block(替换 <chatid>)
|
|
uv run python scripts/wecom_capture_chatid.py --save-chatid <chatid>
|
|
|
|
# 4. 验证收发链路
|
|
uv run python scripts/wecom_capture_chatid.py --test "hello"
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from prefect.blocks.system import Secret # noqa: E402
|
|
|
|
SECRET_BLOCK = "wecom-aibot-credentials"
|
|
|
|
|
|
def _read_block() -> dict[str, str]:
|
|
value = Secret.load(SECRET_BLOCK).get()
|
|
if isinstance(value, str):
|
|
value = json.loads(value)
|
|
return value
|
|
|
|
|
|
def _write_block(creds: dict[str, str]) -> None:
|
|
Secret(value=json.dumps(creds, ensure_ascii=False)).save(SECRET_BLOCK, overwrite=True)
|
|
shown = {k: (v[:8] + "..." if k == "secret" and v else v) for k, v in creds.items()}
|
|
print(f"已保存 Secret Block {SECRET_BLOCK}: {shown}")
|
|
|
|
|
|
def do_init() -> None:
|
|
creds = {
|
|
"bot_id": os.environ.get("WECOM_BOT_ID", ""),
|
|
"secret": os.environ.get("WECOM_SECRET", ""),
|
|
"chat_id": "",
|
|
}
|
|
missing = [k for k in ("bot_id", "secret") if not creds[k]]
|
|
if missing:
|
|
sys.exit(f"缺少环境变量: {missing},请先 export WECOM_BOT_ID=... WECOM_SECRET=...")
|
|
_write_block(creds)
|
|
|
|
|
|
def do_save_chatid(chatid: str) -> None:
|
|
creds = _read_block()
|
|
creds["chat_id"] = chatid
|
|
_write_block(creds)
|
|
|
|
|
|
def do_test(message: str) -> None:
|
|
from flows.shared.wecom import send_markdown
|
|
|
|
ok = send_markdown(f"✅ **机器人通知链路正常**\n\n> {message}")
|
|
sys.exit(0 if ok else 1)
|
|
|
|
|
|
async def _listen(creds: dict[str, str]) -> None:
|
|
from wecom_aibot_sdk import WSClient
|
|
|
|
ws = WSClient(bot_id=creds["bot_id"], secret=creds["secret"])
|
|
|
|
def on_any(frame) -> None:
|
|
body = frame.get("body") or {}
|
|
chatid = body.get("chatid") or ""
|
|
sender = (body.get("from") or {}).get("userid", "")
|
|
msgtype = body.get("msgtype", "")
|
|
text = ((body.get("text") or {}).get("content") or "")[:50]
|
|
print(f"收到帧: msgtype={msgtype} chatid={chatid} from={sender} text={text}")
|
|
if chatid:
|
|
print(f"\n*** chat_id = {chatid} ***\n")
|
|
print(f"写回命令: uv run python scripts/wecom_capture_chatid.py --save-chatid {chatid}")
|
|
|
|
ws.on("message", on_any)
|
|
ws.on("event.enter_chat", on_any)
|
|
ws.on("authenticated", lambda: print("已认证,等待消息(请在目标群 @机器人 发一句话)...Ctrl+C 退出"))
|
|
await ws.connect()
|
|
await asyncio.Event().wait()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--init", action="store_true", help="初始化 Secret Block(需 WECOM_BOT_ID/WECOM_SECRET 环境变量)")
|
|
parser.add_argument("--listen", action="store_true", help="监听消息以捕获 chat_id")
|
|
parser.add_argument("--save-chatid", metavar="CHATID", help="把捕获的 chat_id 写回 Secret Block")
|
|
parser.add_argument("--test", metavar="MSG", help="向已保存的 chat_id 发送测试消息")
|
|
args = parser.parse_args()
|
|
|
|
if args.init:
|
|
do_init()
|
|
elif args.save_chatid:
|
|
do_save_chatid(args.save_chatid)
|
|
elif args.test:
|
|
do_test(args.test)
|
|
elif args.listen:
|
|
# Secret.load 是 sync_compatible,须在进入事件循环前同步调用
|
|
creds = _read_block()
|
|
asyncio.run(_listen(creds))
|
|
else:
|
|
parser.print_help()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|