598 lines
22 KiB
Python
598 lines
22 KiB
Python
import gzip
|
|
import json
|
|
import logging
|
|
import logging.handlers
|
|
import os
|
|
import queue
|
|
import re
|
|
import threading
|
|
from typing import Any, Optional
|
|
|
|
from .message_content import (
|
|
MSG_TYPE_IMAGE,
|
|
MSG_TYPE_LINK_CARD,
|
|
MSG_TYPE_STICKER,
|
|
MSG_TYPE_TEXT,
|
|
MSG_TYPE_VIDEO,
|
|
MSG_TYPE_VOICE,
|
|
_coerce_message_type,
|
|
format_im_message,
|
|
message_preview,
|
|
parse_incoming_message,
|
|
parse_stored_content,
|
|
serialize_message_content,
|
|
)
|
|
|
|
logger = logging.getLogger("douyin_im.protocol")
|
|
|
|
from datetime import datetime
|
|
|
|
|
|
def _is_control_payload(content_json: Any, msg_type: int = 0) -> bool:
|
|
"""判断是否为「会话控制/状态更新」等非聊天内容帧。
|
|
|
|
例如 command_type=6 的连续互动统计(consecutive_chat_data)、ext_data 元数据更新、
|
|
message_type>=50000 的系统通知等——这些不是用户发的消息，不应记录/展示成聊天气泡，
|
|
更不应触发自动回复。
|
|
"""
|
|
try:
|
|
if msg_type and int(msg_type) >= 50000:
|
|
return True
|
|
except (TypeError, ValueError):
|
|
pass
|
|
if isinstance(content_json, dict) and ("command_type" in content_json or "ext_data" in content_json):
|
|
return True
|
|
return False
|
|
|
|
_WS_DEBUG_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "ws_media_debug.log")
|
|
_WS_DEBUG_WRITER_LOCK = threading.Lock()
|
|
_WS_DEBUG_LOGGER: Optional[logging.Logger] = None
|
|
|
|
|
|
def _bounded_env_int(name: str, default: int, minimum: int, maximum: int) -> int:
|
|
try:
|
|
value = int(os.getenv(name, str(default)) or default)
|
|
except (TypeError, ValueError):
|
|
value = default
|
|
return max(minimum, min(maximum, value))
|
|
|
|
|
|
class _DroppingQueueHandler(logging.handlers.QueueHandler):
|
|
"""Never let optional diagnostics block the IM event loop."""
|
|
|
|
def enqueue(self, record) -> None:
|
|
try:
|
|
self.queue.put_nowait(record)
|
|
except queue.Full:
|
|
# Debug output is intentionally lossy under pressure. Receiving
|
|
# and replying to messages must always take precedence.
|
|
return
|
|
|
|
|
|
def _get_ws_debug_logger() -> logging.Logger:
|
|
global _WS_DEBUG_LOGGER
|
|
if _WS_DEBUG_LOGGER is not None:
|
|
return _WS_DEBUG_LOGGER
|
|
with _WS_DEBUG_WRITER_LOCK:
|
|
if _WS_DEBUG_LOGGER is not None:
|
|
return _WS_DEBUG_LOGGER
|
|
|
|
max_bytes = _bounded_env_int(
|
|
"KEFU_WS_DEBUG_MAX_BYTES", 10 * 1024 * 1024, 1024 * 1024, 100 * 1024 * 1024
|
|
)
|
|
backup_count = _bounded_env_int(
|
|
"KEFU_WS_DEBUG_BACKUP_COUNT", 2, 1, 10
|
|
)
|
|
queue_size = _bounded_env_int(
|
|
"KEFU_WS_DEBUG_QUEUE_SIZE", 1000, 100, 10000
|
|
)
|
|
records: queue.Queue = queue.Queue(maxsize=queue_size)
|
|
rotating = logging.handlers.RotatingFileHandler(
|
|
_WS_DEBUG_PATH,
|
|
maxBytes=max_bytes,
|
|
backupCount=backup_count,
|
|
encoding="utf-8",
|
|
delay=True,
|
|
)
|
|
rotating.setFormatter(logging.Formatter("%(message)s"))
|
|
listener = logging.handlers.QueueListener(
|
|
records,
|
|
rotating,
|
|
respect_handler_level=True,
|
|
)
|
|
listener.start()
|
|
|
|
debug_logger = logging.getLogger("douyin_im.ws_raw_debug")
|
|
debug_logger.handlers.clear()
|
|
debug_logger.addHandler(_DroppingQueueHandler(records))
|
|
debug_logger.setLevel(logging.INFO)
|
|
debug_logger.propagate = False
|
|
# Keep strong references for the lifetime of the logger/listener.
|
|
debug_logger._kefu_queue_listener = listener # type: ignore[attr-defined]
|
|
debug_logger._kefu_rotating_handler = rotating # type: ignore[attr-defined]
|
|
_WS_DEBUG_LOGGER = debug_logger
|
|
return debug_logger
|
|
|
|
|
|
def _should_emit_ws_message(
|
|
conversation_id: str,
|
|
msg_type: int,
|
|
) -> bool:
|
|
"""判断 WS 帧是否为用户聊天消息（控制帧已在 _is_control_payload 过滤）。"""
|
|
if not conversation_id:
|
|
return False
|
|
try:
|
|
if int(msg_type) >= 50000:
|
|
return False
|
|
except (TypeError, ValueError):
|
|
pass
|
|
return True
|
|
|
|
|
|
def _dump_ws_message(msg_type: int, conversation_id: str, content_str: str, msg: Any = None) -> None:
|
|
"""把每条 WS 消息的全部字段落到调试文件，便于排查媒体字段结构。
|
|
|
|
默认关闭，仅当设置环境变量 KEFU_WS_DEBUG=1 时写盘，避免生产环境无界增长 / 泄露聊天内容。
|
|
content 为空时（如 type=26 瘦推送）会额外打印 protobuf 其余字段，确保「接收到的全部信息」可见。
|
|
"""
|
|
if os.getenv("KEFU_WS_DEBUG", "") not in ("1", "true", "True"):
|
|
return
|
|
try:
|
|
extra = ""
|
|
if msg is not None:
|
|
fields = {}
|
|
try:
|
|
for f, v in msg.ListFields():
|
|
if f.name == "content":
|
|
continue
|
|
fields[f.name] = v
|
|
except Exception:
|
|
pass
|
|
if fields:
|
|
extra = " | fields=" + json.dumps(fields, ensure_ascii=False, default=str)
|
|
line = (
|
|
f"{datetime.now().isoformat()} type={msg_type} "
|
|
f"conv={conversation_id} content={content_str}{extra}\n"
|
|
)
|
|
record_limit = _bounded_env_int(
|
|
"KEFU_WS_DEBUG_RECORD_MAX_CHARS", 16384, 1024, 262144
|
|
)
|
|
if len(line) > record_limit:
|
|
marker = "...[单条调试记录过长，已截断]\n"
|
|
line = line[: max(0, record_limit - len(marker))] + marker
|
|
_get_ws_debug_logger().info(line.rstrip("\n"))
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _walk_strings(data: Any, depth: int = 0, max_depth: int = 10):
|
|
if depth > max_depth:
|
|
return
|
|
if isinstance(data, dict):
|
|
for v in data.values():
|
|
yield from _walk_strings(v, depth + 1, max_depth)
|
|
elif isinstance(data, list):
|
|
for item in data:
|
|
yield from _walk_strings(item, depth + 1, max_depth)
|
|
elif isinstance(data, str) and data.strip():
|
|
yield data.strip()
|
|
|
|
|
|
def extract_json_objects(raw: bytes | str) -> list[dict]:
|
|
"""从二进制帧中尽量提取 JSON 对象"""
|
|
if isinstance(raw, bytes):
|
|
for codec in ("utf-8", "latin-1"):
|
|
try:
|
|
text = raw.decode(codec, errors="ignore")
|
|
break
|
|
except Exception:
|
|
text = ""
|
|
else:
|
|
text = ""
|
|
else:
|
|
text = raw
|
|
|
|
results = []
|
|
for match in re.finditer(r"\{[^{}]{0,2000}\}", text):
|
|
chunk = match.group(0)
|
|
try:
|
|
obj = json.loads(chunk)
|
|
if isinstance(obj, dict):
|
|
results.append(obj)
|
|
except Exception:
|
|
continue
|
|
return results
|
|
|
|
|
|
def parse_ws_payload(raw: bytes | str) -> list[dict]:
|
|
"""解析 WebSocket 二进制帧，返回标准化消息 dict 列表"""
|
|
messages = []
|
|
|
|
# 尝试 Protobuf 解包
|
|
if isinstance(raw, bytes):
|
|
try:
|
|
from .static import Live_pb2, Response_pb2
|
|
frame = Live_pb2.PushFrame()
|
|
frame.ParseFromString(raw)
|
|
if frame.payloadType == 'pb':
|
|
response = Response_pb2.Response()
|
|
response.ParseFromString(frame.payload)
|
|
body = response.body
|
|
if body.HasField("new_message_notify"):
|
|
notify = body.new_message_notify
|
|
if notify.HasField("message"):
|
|
msg = notify.message
|
|
sender = str(msg.sender)
|
|
msg_type = msg.message_type
|
|
conversation_id = msg.conversation_id
|
|
content_str = msg.content
|
|
server_message_id = str(getattr(msg, "server_message_id", "") or "")
|
|
|
|
text_content = ""
|
|
media_msg: dict = {}
|
|
content_json: dict = {}
|
|
try:
|
|
content_json = json.loads(content_str) if content_str else {}
|
|
if not isinstance(content_json, dict):
|
|
content_json = {}
|
|
media_msg = format_im_message(content_json, msg_type)
|
|
text_content = media_msg.get("text") or ""
|
|
except Exception:
|
|
media_msg = format_im_message(content_str or "", msg_type)
|
|
text_content = media_msg.get("text") or content_str or ""
|
|
|
|
# 跳过会话控制/状态更新等非聊天内容帧（不记录、不展示、不触发自动回复）
|
|
if _is_control_payload(content_json, msg_type):
|
|
logger.debug(
|
|
"Skip control WS frame: type=%s conv=%s", msg_type, conversation_id
|
|
)
|
|
return messages
|
|
|
|
# Raw diagnostics are optional and intentionally run
|
|
# only after control/status frames have been filtered.
|
|
# The writer itself is queued and rotating, so it can
|
|
# never block message parsing or grow without bound.
|
|
_dump_ws_message(msg_type, conversation_id, content_str, msg)
|
|
|
|
if _should_emit_ws_message(conversation_id, msg_type):
|
|
sender_uid = str(msg.sender)
|
|
if media_msg and (
|
|
media_msg.get("text")
|
|
or media_msg.get("type") not in (None, "text", "")
|
|
):
|
|
display_content = serialize_message_content(media_msg)
|
|
else:
|
|
display_content = text_content or content_str
|
|
payload = {
|
|
"sender_name": sender_uid,
|
|
"sender_uid": sender_uid,
|
|
"content": display_content,
|
|
"raw_content": content_str,
|
|
"conversation_id": conversation_id,
|
|
"unread_count": 1,
|
|
"server_message_id": server_message_id,
|
|
"message_type": msg_type,
|
|
}
|
|
if msg_type in (
|
|
MSG_TYPE_IMAGE,
|
|
MSG_TYPE_STICKER,
|
|
MSG_TYPE_VOICE,
|
|
MSG_TYPE_VIDEO,
|
|
MSG_TYPE_LINK_CARD,
|
|
):
|
|
if not media_msg.get("url") and not media_msg.get("uri"):
|
|
logger.warning(
|
|
"Media WS message missing url: type=%s content=%s",
|
|
msg_type,
|
|
(content_str or "")[:800],
|
|
)
|
|
elif not media_msg.get("url"):
|
|
logger.info(
|
|
"Media WS message resolved via uri: type=%s uri=%s",
|
|
msg_type,
|
|
media_msg.get("uri"),
|
|
)
|
|
messages.append(payload)
|
|
logger.info(
|
|
"Protobuf WS message parsed: sender=%s type=%s content=%s conv=%s",
|
|
sender,
|
|
msg_type,
|
|
text_content,
|
|
conversation_id,
|
|
)
|
|
return messages
|
|
except Exception as e:
|
|
logger.debug(f"Protobuf WS parse failed: {e}")
|
|
|
|
if isinstance(raw, str):
|
|
payloads = [raw.encode("utf-8", errors="ignore")]
|
|
else:
|
|
payloads = [raw]
|
|
# 尝试 gzip 解压（frontier 常见）
|
|
try:
|
|
payloads.append(gzip.decompress(raw))
|
|
except Exception:
|
|
pass
|
|
|
|
for payload in payloads:
|
|
# 1) 直接 JSON
|
|
if isinstance(payload, bytes):
|
|
text = payload.decode("utf-8", errors="ignore").strip()
|
|
else:
|
|
text = str(payload).strip()
|
|
if text.startswith("{") or text.startswith("["):
|
|
try:
|
|
data = json.loads(text)
|
|
messages.extend(normalize_im_payload(data))
|
|
continue
|
|
except Exception:
|
|
pass
|
|
|
|
# 2) 嵌入 JSON
|
|
for obj in extract_json_objects(payload):
|
|
messages.extend(normalize_im_payload(obj))
|
|
|
|
# 3) 纯文本兜底
|
|
if isinstance(payload, bytes):
|
|
text = payload.decode("utf-8", errors="ignore")
|
|
plain = _extract_plain_text(text)
|
|
if plain:
|
|
messages.append({"content": plain, "sender_name": "", "raw_content": plain, "raw": True})
|
|
|
|
return messages
|
|
|
|
|
|
def normalize_im_payload_from_bytes(raw: bytes) -> list[dict]:
|
|
"""Try to extract conversation/message payloads from binary IM API responses."""
|
|
results: list[dict] = []
|
|
for obj in extract_json_objects(raw):
|
|
results.extend(normalize_im_payload(obj))
|
|
if results:
|
|
return results
|
|
|
|
try:
|
|
from .static import Response_pb2
|
|
response = Response_pb2.Response()
|
|
response.ParseFromString(raw)
|
|
body = response.body
|
|
for field in (
|
|
"get_conversation_info_list_v2_response_body",
|
|
"create_conversation_v2_body",
|
|
):
|
|
if body.HasField(field):
|
|
conv_body = getattr(body, field)
|
|
for conv in conv_body.conversation_info_list:
|
|
conv_id = conv.conversation_id
|
|
peer_uid = ""
|
|
parts = conv_id.split(":")
|
|
if len(parts) >= 4:
|
|
peer_uid = parts[-1]
|
|
label = f"用户{peer_uid[-6:]}" if peer_uid else conv_id
|
|
results.append({
|
|
"conversation_id": conv_id,
|
|
"sender_name": label,
|
|
"content": "",
|
|
"unread_count": 0,
|
|
"peer_uid": peer_uid,
|
|
"conversation_short_id": str(conv.conversation_short_id),
|
|
"ticket": conv.ticket,
|
|
})
|
|
except Exception:
|
|
pass
|
|
return results
|
|
|
|
|
|
def normalize_im_payload(data: Any, depth: int = 0) -> list[dict]:
|
|
"""递归标准化 IM JSON 为 {sender_name, content, conversation_id, unread_count}"""
|
|
if depth > 12:
|
|
return []
|
|
results = []
|
|
|
|
if isinstance(data, list):
|
|
for item in data:
|
|
results.extend(normalize_im_payload(item, depth + 1))
|
|
return results
|
|
|
|
if not isinstance(data, dict):
|
|
return results
|
|
|
|
sender = (
|
|
_pick_str(data, "sender_name", "senderName", "nickname", "nick_name", "userName", "peerName")
|
|
or _pick_nested(data, ("core_info", "user_info", "peer_info"), "nick_name", "nickname", "name")
|
|
)
|
|
sender_avatar = _pick_avatar_url(data)
|
|
content = _pick_message_text(data)
|
|
msg_type = _coerce_message_type(
|
|
data.get("message_type") or data.get("messageType") or data.get("msg_type"),
|
|
MSG_TYPE_TEXT,
|
|
)
|
|
# 会话控制/状态更新帧（command_type / ext_data / 系统通知）直接忽略，不当作聊天消息
|
|
if _is_control_payload(data, msg_type) or _is_control_payload(data.get("content"), msg_type):
|
|
return results
|
|
if msg_type != MSG_TYPE_TEXT or (isinstance(data.get("content"), dict)):
|
|
parsed = parse_incoming_message(data)
|
|
if parsed:
|
|
content = parsed
|
|
elif content and content.startswith("{"):
|
|
try:
|
|
parsed = format_im_message(json.loads(content), msg_type)
|
|
content = serialize_message_content(parsed)
|
|
except Exception:
|
|
pass
|
|
elif content in _NON_TEXT_MESSAGE_MARKERS:
|
|
parsed = parse_stored_content(content)
|
|
content = serialize_message_content(parsed)
|
|
conv_id = _pick_str(
|
|
data,
|
|
"conversation_id",
|
|
"conversationId",
|
|
"conv_id",
|
|
"cid",
|
|
)
|
|
unread = data.get("unread_count") or data.get("unreadCount") or data.get("unread_cnt") or 0
|
|
try:
|
|
unread = int(unread or 0)
|
|
except (TypeError, ValueError):
|
|
unread = 0
|
|
|
|
if content and len(content) < 500:
|
|
from_self = data.get("is_self") or data.get("isSelf") or data.get("fromSelf") or data.get("self")
|
|
if not from_self:
|
|
raw_content = _extract_raw_content(data) or content
|
|
results.append({
|
|
"sender_name": sender or "未知用户",
|
|
"sender_avatar": sender_avatar or None,
|
|
"content": content,
|
|
"raw_content": raw_content,
|
|
"conversation_id": conv_id or "",
|
|
"unread_count": unread,
|
|
"message_type": msg_type,
|
|
})
|
|
|
|
if sender and unread > 0 and not content:
|
|
results.append({
|
|
"sender_name": sender,
|
|
"sender_avatar": sender_avatar or None,
|
|
"content": "[未读消息]",
|
|
"raw_content": "[未读消息]",
|
|
"conversation_id": conv_id or "",
|
|
"unread_count": unread,
|
|
"message_type": msg_type,
|
|
})
|
|
|
|
for key in ("conversations", "conversation_list", "data", "messages", "messagesList", "body"):
|
|
nested = data.get(key)
|
|
if nested is not None:
|
|
results.extend(normalize_im_payload(nested, depth + 1))
|
|
|
|
for value in data.values():
|
|
if isinstance(value, (dict, list)):
|
|
results.extend(normalize_im_payload(value, depth + 1))
|
|
|
|
return results
|
|
|
|
|
|
def _extract_raw_content(data: dict) -> str:
|
|
for key in ("content", "message", "msg", "lastMessage", "last_msg", "preview", "brief"):
|
|
val = data.get(key)
|
|
if isinstance(val, str) and val.strip():
|
|
return val.strip()
|
|
if isinstance(val, dict):
|
|
return json.dumps(val, ensure_ascii=False, separators=(",", ":"))
|
|
return ""
|
|
|
|
|
|
def _pick_str(data: dict, *keys: str) -> str:
|
|
for key in keys:
|
|
val = data.get(key)
|
|
if isinstance(val, str) and val.strip():
|
|
return val.strip()
|
|
return ""
|
|
|
|
|
|
def _pick_nested(data: dict, parent_keys: tuple, *child_keys: str) -> str:
|
|
for pk in parent_keys:
|
|
nested = data.get(pk)
|
|
if isinstance(nested, dict):
|
|
val = _pick_str(nested, *child_keys)
|
|
if val:
|
|
return val
|
|
return ""
|
|
|
|
|
|
def _avatar_from_value(val: Any) -> str:
|
|
if isinstance(val, str) and val.strip().startswith("http"):
|
|
return val.strip()
|
|
if isinstance(val, dict):
|
|
direct = val.get("url")
|
|
if isinstance(direct, str) and direct.startswith("http"):
|
|
return direct.strip()
|
|
for list_key in ("url_list", "urls"):
|
|
urls = val.get(list_key)
|
|
if isinstance(urls, list):
|
|
for item in urls:
|
|
if isinstance(item, str) and item.startswith("http"):
|
|
return item.strip()
|
|
return ""
|
|
|
|
|
|
def _pick_avatar_url(data: dict) -> str:
|
|
for key in ("avatar_url", "avatarUrl", "head_url", "headUrl", "avatar"):
|
|
url = _avatar_from_value(data.get(key))
|
|
if url:
|
|
return url
|
|
for thumb_key in ("avatar_thumb", "avatar_medium", "avatar_larger", "avatarThumb"):
|
|
url = _avatar_from_value(data.get(thumb_key))
|
|
if url:
|
|
return url
|
|
for parent_key in ("core_info", "user_info", "peer_info", "target_user", "conversation_core_info"):
|
|
nested = data.get(parent_key)
|
|
if isinstance(nested, dict):
|
|
url = _pick_avatar_url(nested)
|
|
if url:
|
|
return url
|
|
return ""
|
|
|
|
|
|
def _pick_message_text(data: dict) -> str:
|
|
for key in (
|
|
"text",
|
|
"content",
|
|
"message",
|
|
"msg",
|
|
"lastMessage",
|
|
"last_msg",
|
|
"preview",
|
|
"brief",
|
|
):
|
|
val = data.get(key)
|
|
if isinstance(val, str) and val.strip():
|
|
return val.strip()
|
|
if isinstance(val, dict):
|
|
inner = _pick_str(val, "text", "content", "message")
|
|
if inner:
|
|
return inner
|
|
return ""
|
|
|
|
|
|
def _extract_plain_text(text: str) -> Optional[str]:
|
|
text = (text or "").strip()
|
|
if not text or len(text) > 200:
|
|
return None
|
|
if text.startswith("{") or text.startswith("["):
|
|
return None
|
|
# 过滤明显二进制垃圾
|
|
printable = sum(1 for c in text if c.isprintable() or c in "\n\r\t")
|
|
if printable / max(len(text), 1) < 0.8:
|
|
return None
|
|
return text
|
|
|
|
|
|
_NON_TEXT_MESSAGE_MARKERS = {
|
|
"[表情包]",
|
|
"[语音]",
|
|
"[图片]",
|
|
"[视频]",
|
|
"[未读消息]",
|
|
}
|
|
|
|
|
|
def should_skip_auto_reply(content: str) -> tuple[bool, str]:
|
|
"""判断收到的内容是否不适合触发自动回复（如纯点赞/表情互动）。"""
|
|
from .message_content import is_media_message, message_preview
|
|
|
|
text = (content or "").strip()
|
|
if not text:
|
|
return True, "空消息"
|
|
if is_media_message(text):
|
|
return True, f"非文本消息({message_preview(text)})"
|
|
if text in _NON_TEXT_MESSAGE_MARKERS:
|
|
return True, f"非文本消息({text})"
|
|
if re.fullmatch(r"(\[赞\])+", text):
|
|
return True, "表情互动消息(点赞)，抖音通常不允许对此类消息自动回复"
|
|
if re.fullmatch(r"\[[^\]]+\](\[[^\]]+\])*", text) and "http" not in text:
|
|
inner = re.sub(r"[\[\]]", "", text)
|
|
if len(inner) <= 20 and not any(ch.isalnum() for ch in inner):
|
|
return True, f"非文本互动消息({text})"
|
|
return False, ""
|