Initial import of Shirobana bot backend

Include utility modules, shared channels, prompts, read-only skills, wardrobe memory and regression tests. Externalize legacy credentials and exclude private runtime state.
2026-11-11 00:25:30 +08:00
commit 34089d0042
106 changed files with 28028 additions and 0 deletions
@@ -0,0 +1,350 @@
+# live_monitor.py
+# SRS 直播监控 + 通知模块
+#   群聊关注：群推群（条目按 channel + 群号）
+#   私聊关注：挂在统一身份上（条目按 uid），通知时按 TG + 已绑定 QQ 扇出镜像推送
+# 主站(site=1): 远程 aliyun SRS → /remote_srs/on_publish
+# 副站(site=2): 本地 SRS → /remote_srs/local/on_publish
+import hmac
+import json
+import os
+import queue
+import threading
+from flask import Flask, request
+
+app = Flask(__name__)
+notification_queue = queue.Queue()
+
+RECORD_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "live_monitor")
+RECORD_FILE = os.path.join(RECORD_DIR, "record.json")
+
+websocket = None
+
+SITE_LABEL = {"1": "主站", "2": "副站"}
+
+# Configure SRS http_hooks with ?token=<secret>, or send X-SRS-Token. The route
+# and proxy/frp setup stay unchanged. Missing configuration fails closed.
+SRS_CALLBACK_SECRET = os.environ.get(
+    "BOT_SRS_CALLBACK_SECRET", os.environ.get("SRS_CALLBACK_SECRET", ""))
+_live_sessions = set()       # (site, room_id) currently publishing
+_live_sessions_lock = threading.Lock()
+_auth_warning = False
+
+HELP_TEXT = (
+    "直播监控\n"
+    "━━━━━━━━━━━━\n"
+    "• .live follow {房间号0-100} {1/2/空}\n"
+    "• .live unfollow {房间号} {1/2/空}\n"
+    "• .live status  查看当前监控的房间\n"
+    "  1=主站 2=副站 空=全部\n"
+    "示例：.live follow 42 1\n"
+    "━━━━━━━━━━━━\n"
+    "直播站 https://live.rrrosmontis.icu"
+)
+
+def _ensure_record():
+    os.makedirs(os.path.dirname(RECORD_FILE), exist_ok=True)
+    if not os.path.exists(RECORD_FILE):
+        with open(RECORD_FILE, "w", encoding="utf-8") as f:
+            json.dump({"follows": {}}, f)
+
+def _load_records():
+    _ensure_record()
+    with open(RECORD_FILE, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+def _save_records(data):
+    """原子写：同目录临时文件 + flush/fsync + os.replace，避免写一半留下截断的
+    record.json（订阅全丢）。失败时清理临时文件，原文件保持不动。
+    签名与输出内容与旧实现一致（ensure_ascii=False, indent=2，无尾随换行）。
+    """
+    _ensure_record()
+    tmp = RECORD_FILE + ".tmp"
+    try:
+        with open(tmp, "w", encoding="utf-8") as f:
+            json.dump(data, f, ensure_ascii=False, indent=2)
+            f.flush()
+            os.fsync(f.fileno())
+        os.replace(tmp, RECORD_FILE)
+    except BaseException:
+        try:
+            if os.path.exists(tmp):
+                os.unlink(tmp)
+        except OSError:
+            pass
+        raise
+
+
+def _private_uid(value, channel=None):
+    """把历史/当前私聊标识归一为统一身份 uid。
+
+    兼容三种历史写法:
+      - uid（[UID_BASE, 1e10)，历史 record 里 chat_id 直接存的就是 uid）
+      - TG ns（>= 1e10）-> identity.resolve("tg", ns)
+      - 原始 QQ 号 / TG 裸 id -> identity.resolve(channel|qq, value)
+    解析不出来返回 None（调用方跳过该条目）。
+    """
+    try:
+        number = int(value)
+    except (TypeError, ValueError):
+        return None
+    import identity
+    if identity.UID_BASE <= number < 10 ** 10:
+        return number
+    uid = identity.resolve("tg", number)
+    if uid is not None:
+        return uid
+    return identity.resolve(channel or "qq", number)
+
+
+def _entry_uid(entry):
+    """关注条目的 uid：新记录直接存 uid，旧记录回退解析 chat_id"""
+    if entry.get("uid") is not None:
+        try:
+            return int(entry["uid"])
+        except (TypeError, ValueError):
+            return None
+    return _private_uid(entry.get("chat_id"), entry.get("channel"))
+
+
+def _follow_key(entry):
+    """关注条目的去重/匹配键：私聊按 uid（跨通道合并），群聊按 channel+群号"""
+    if entry.get("chat_type") == "private":
+        uid = _entry_uid(entry)
+        return ("private", uid) if uid is not None else None
+    return ("group", entry.get("channel", "tg"), str(entry.get("chat_id")))
+
+
+def _matches_follow(entry, chat_type, chat_id, channel):
+    """条目是否属于 (chat_type, chat_id, channel)。私聊参数 chat_id 传 uid。"""
+    if entry.get("chat_type") != chat_type:
+        return False
+    if chat_type == "private":
+        target = _private_uid(chat_id, channel)
+        return target is not None and _entry_uid(entry) == target
+    return entry.get("channel", "tg") == channel and str(entry.get("chat_id")) == str(chat_id)
+
+
+def add_follow(room_id, chat_type, chat_id, site, channel="tg"):
+    """关注直播间。site: '1'=主站, '2'=副站, ''=全部; channel: 群聊通知通道 (qq/tg)
+
+    私聊关注统一挂在统一身份上（chat_id 传 uid），通知时按 TG + QQ 扇出；
+    群聊关注仍是群推群（channel + 群号）。
+    """
+    data = _load_records()
+    room_id = str(room_id)
+    follows = data.setdefault("follows", {})
+    room_follows = follows.setdefault(room_id, [])
+    if chat_type == "private":
+        uid = _private_uid(chat_id, channel)
+        if uid is None:
+            return
+        entry = {"chat_type": "private", "uid": uid, "site": site}
+    else:
+        entry = {"channel": channel, "chat_type": chat_type,
+                 "chat_id": chat_id, "site": site}
+    key = _follow_key(entry)
+    if key is None or not any(_follow_key(f) == key and f.get("site", "") == site
+                              for f in room_follows):
+        room_follows.append(entry)
+    _save_records(data)
+
+
+def remove_follow(room_id, chat_type, chat_id, site, channel="tg"):
+    """取消关注"""
+    data = _load_records()
+    room_id = str(room_id)
+    follows = data.get("follows", {})
+    if room_id in follows:
+        follows[room_id] = [
+            f for f in follows[room_id]
+            if not (_matches_follow(f, chat_type, chat_id, channel)
+                    and f.get("site", "") == site)
+        ]
+        if not follows[room_id]:
+            del follows[room_id]
+    _save_records(data)
+
+
+def get_follows(chat_type, chat_id, channel="tg"):
+    """获取某聊天关注的所有房间号和站点（私聊按 uid，群聊按 channel+群号）"""
+    data = _load_records()
+    follows = data.get("follows", {})
+    result = []
+    for room_id, flist in follows.items():
+        for f in flist:
+            if _matches_follow(f, chat_type, chat_id, channel):
+                result.append((room_id, f.get("site", "")))
+                break
+    return result
+
+
+def status_text(chat_type=None, chat_id=None, channel=None):
+    """当前已监控的房间列表；标注会推送到当前会话的房间（不暴露其它会话的 id）"""
+    data = _load_records()
+    follows = data.get("follows") or {}
+    if not follows:
+        return "当前没有监控任何直播间。\n用法：.live follow {房间号0-100} {1/2/空}"
+    lines = ["直播关注（共 %d 个房间）" % len(follows), "━━━━━━━━━━━━"]
+    for room in sorted(follows, key=lambda x: int(x) if str(x).isdigit() else 0):
+        items = follows[room] or []
+        sites = []
+        for item in items:
+            label = SITE_LABEL.get(item.get("site", ""), "全部")
+            if label not in sites:
+                sites.append(label)
+        marks = []
+        if chat_id is not None and any(
+                _matches_follow(f, chat_type, chat_id, channel) for f in items):
+            marks.append("本群" if chat_type == "group" else "本会话")
+        if any(f.get("chat_type") == "private" for f in items):
+            marks.append("私聊")
+        suffix = "（%s）" % "、".join(marks) if marks else ""
+        lines.append("• 房间 %-3s %s %s" % (room, "/".join(sites), suffix))
+    lines.append("━━━━━━━━━━━━")
+    lines.append("关注：.live follow <房间号> [1/2]　取消：.live unfollow <房间号> [1/2]")
+    return "\n".join(lines)
+
+
+def _notification(room_id, site, f):
+    """把关注条目转成一条通知。
+
+    私聊条目归一为 {"chat_type":"private","uid":..}（消费者按 TG + QQ 扇出）；
+    群聊条目保持 {"channel":.., "chat_type":"group", "chat_id":群号}（群推群）；
+    私聊条目若 uid 解析不出来，退回旧写法（按原 channel/target 发送），保证兼容。
+    """
+    if f.get("chat_type") == "private":
+        uid = _entry_uid(f)
+        if uid is not None:
+            return {"room_id": room_id, "site": site,
+                    "chat_type": "private", "uid": uid}
+    return {"room_id": room_id, "site": site, "chat_type": f.get("chat_type"),
+            "chat_id": f.get("chat_id"), "channel": f.get("channel", "tg")}
+
+
+def _srs_callback_authorized():
+    global _auth_warning
+    if not SRS_CALLBACK_SECRET:
+        if not _auth_warning:
+            print("[live_monitor] 拒绝 SRS 回调: 未配置 BOT_SRS_CALLBACK_SECRET")
+            _auth_warning = True
+        return False
+    authorization = request.headers.get("Authorization", "")
+    supplied = (request.headers.get("X-SRS-Token", "")
+                or request.args.get("token", "")
+                or (authorization[7:] if authorization.startswith("Bearer ") else ""))
+    return hmac.compare_digest(str(supplied), str(SRS_CALLBACK_SECRET))
+
+
+def _callback_room_id(data):
+    raw_stream = str(data.get("stream", ""))
+    return (raw_stream.replace("room_", "", 1)
+            if raw_stream.startswith("room_") else raw_stream)
+
+
+def _begin_live_session(site, room_id):
+    """Atomically deduplicate repeated on_publish for one active session."""
+    key = (str(site), str(room_id))
+    with _live_sessions_lock:
+        if key in _live_sessions:
+            return False
+        _live_sessions.add(key)
+        return True
+
+
+def _end_live_session(site, room_id):
+    with _live_sessions_lock:
+        _live_sessions.discard((str(site), str(room_id)))
+
+
+def _handle_on_publish(site):
+    """处理 on_publish 回调，site='1' 主站 '2' 副站"""
+    if not _srs_callback_authorized():
+        return '{"code": 1, "msg": "forbidden"}', 403
+    data = request.get_json(force=True, silent=True) or {}
+    print(f"[live_monitor] on_publish(site={site}) received: {json.dumps(data, ensure_ascii=False)}")
+    raw_stream = str(data.get("stream", ""))
+    if not raw_stream:
+        print(f"[live_monitor] on_publish(site={site}): no stream field")
+        return '{"code": 0}'
+
+    room_id = _callback_room_id(data)
+    if not _begin_live_session(site, room_id):
+        print(f"[live_monitor] on_publish(site={site}): duplicate active session room={room_id}")
+        return '{"code": 0}'
+
+    records = _load_records()
+    room_follows = records.get("follows", {}).get(room_id, [])
+    matched = [f for f in room_follows if f.get("site", "") in ("", site)]
+    print(f"[live_monitor] on_publish(site={site}): raw_stream={raw_stream}, room_id={room_id}, matched={matched}")
+
+    for f in matched:
+        notif = _notification(room_id, site, f)
+        notification_queue.put(notif)
+        print("[live_monitor] on_publish(site=%s): queued for %s"
+              % (site, notif.get("uid") if notif.get("uid") is not None
+                 else "%s/%s/%s" % (notif.get("channel", "qq"), notif.get("chat_type"),
+                                    notif.get("chat_id"))))
+
+    return '{"code": 0}'
+
+
+@app.route("/internal/notify", methods=["POST"])
+def internal_notify():
+    """live-api 内部开播通知（仅本机 127.0.0.1 可达）：按 room+site 直接入队。"""
+    # Flask 绑在 0.0.0.0（SRS 公网回调需要），本端点必须自验来源，否则外部可伪造开播通知
+    if request.remote_addr not in ("127.0.0.1", "::1"):
+        print(f"[live_monitor] internal_notify 拒绝非本机来源: {request.remote_addr}")
+        return '{"code": 1, "msg": "forbidden"}', 403
+    data = request.get_json(force=True, silent=True) or {}
+    room_id = str(data.get("room_id", "")).strip()
+    site = str(data.get("site", "")).strip()
+    title = str(data.get("title", ""))
+    nickname = str(data.get("nickname", ""))
+    if not room_id or site not in ("1", "2"):
+        return '{"code": 1, "msg": "bad params"}'
+    if not _begin_live_session(site, room_id):
+        return '{"code": 0}'
+    records = _load_records()
+    room_follows = records.get("follows", {}).get(room_id, [])
+    matched = [f for f in room_follows if f.get("site", "") in ("", site)]
+    for f in matched:
+        notification_queue.put(_notification(room_id, site, f))
+    print(f"[live_monitor] internal_notify room={room_id} site={site} matched={len(matched)} title={title!r} anchor={nickname!r}")
+    return '{"code": 0}'
+
+
+@app.route("/remote_srs/on_publish", methods=["POST"])
+def on_publish_main():
+    """主站(aliyun 远程) on_publish"""
+    return _handle_on_publish("1")
+
+
+def _handle_on_unpublish(site):
+    if not _srs_callback_authorized():
+        return '{"code": 1, "msg": "forbidden"}', 403
+    data = request.get_json(force=True, silent=True) or {}
+    raw_stream = str(data.get("stream", ""))
+    if raw_stream:
+        _end_live_session(site, _callback_room_id(data))
+    return '{"code": 0}'
+
+
+@app.route("/remote_srs/on_unpublish", methods=["POST"])
+def on_unpublish_main():
+    return _handle_on_unpublish("1")
+
+
+@app.route("/remote_srs/local/on_publish", methods=["POST"])
+def on_publish_local():
+    """副站(本地) on_publish"""
+    return _handle_on_publish("2")
+
+
+@app.route("/remote_srs/local/on_unpublish", methods=["POST"])
+def on_unpublish_local():
+    return _handle_on_unpublish("2")
+
+
+def run_flask():
+    from waitress import serve
+    serve(app, host=os.environ.get("BOT_LIVE_HOST", "127.0.0.1"), port=1986)