RRRosmontis
34089d0042
Include utility modules, shared channels, prompts, read-only skills, wardrobe memory and regression tests. Externalize legacy credentials and exclude private runtime state.
351 lines
14 KiB
Python
351 lines
14 KiB
Python
# 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)
|