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, ""