diff --git a/src/wechat_decrypt_tool/chat_helpers.py b/src/wechat_decrypt_tool/chat_helpers.py index a0f1abec..e02e3edb 100644 --- a/src/wechat_decrypt_tool/chat_helpers.py +++ b/src/wechat_decrypt_tool/chat_helpers.py @@ -2473,10 +2473,17 @@ def looks_like_username(s: str) -> bool: return False return True - def pick_display(strings: list[tuple[int, str]], target: str) -> str: + def pick_display( + strings: list[tuple[int, str]], + target: str, + *, + allowed_fields: set[int], + ) -> str: best_score = -1 best = "" for i, (fno, value) in enumerate(strings): + if int(fno) not in allowed_fields: + continue v = str(value or "").strip() if (not v) or v == target: continue @@ -2525,7 +2532,7 @@ def pick_display(strings: list[tuple[int, str]], target: str) -> str: if not ext_buf: return {} - out: dict[str, str] = {} + member_records: list[list[tuple[int, str]]] = [] for _, wire_type, chunk in iter_fields(ext_buf): if wire_type != 2 or (not chunk): continue @@ -2553,13 +2560,64 @@ def pick_display(strings: list[tuple[int, str]], target: str) -> str: if not strings: continue + member_records.append(strings) + + # Layout A stores the member username in field 1 and its optional group + # nickname in field 2. A chatroom with no group nicknames can therefore + # omit field 2 entirely, while field 4 may still reference another + # member. Treat a requested username in field 1 as positive evidence of + # the current layout; do not infer the legacy layout merely because the + # whole chatroom lacks field 2. + uses_current_layout = any( + field_no == 2 or (field_no == 1 and value in target_set) + for strings in member_records + for field_no, value in strings + ) - present = [v for _, v in strings if v in target_set and v not in out] - if not present: - continue - - for target in present: - disp = pick_display(strings, target) + out: dict[str, str] = {} + if uses_current_layout: + for strings in member_records: + primary_targets = [ + value + for field_no, value in strings + if field_no == 1 and value in target_set and value not in out + ] + for target in primary_targets: + disp = pick_display(strings, target, allowed_fields={2}) + if disp: + out[target] = disp + if len(out) >= len(target_set): + break + return out + + # Legacy field 4 -> field 1 mapping is allowed only with positive + # evidence that every relevant field-1 value is display text rather + # than an alias-style username. Ambiguous ASCII values fail closed and + # let the existing caller fallback resolve the display name. + legacy_records = [ + strings + for strings in member_records + if any(field_no == 4 and value in target_set for field_no, value in strings) + ] + if not legacy_records: + return {} + for strings in legacy_records: + field1_values = [ + str(value or "").strip() + for field_no, value in strings + if field_no == 1 and str(value or "").strip() + ] + if not field1_values or any(looks_like_username(value) for value in field1_values): + return {} + + for strings in legacy_records: + legacy_targets = [ + value + for field_no, value in strings + if field_no == 4 and value in target_set and value not in out + ] + for target in legacy_targets: + disp = pick_display(strings, target, allowed_fields={1}) if disp: out[target] = disp if len(out) >= len(target_set): diff --git a/tests/test_group_nickname_ext_buffer_parsing.py b/tests/test_group_nickname_ext_buffer_parsing.py index 40558c4e..2f9cef93 100644 --- a/tests/test_group_nickname_ext_buffer_parsing.py +++ b/tests/test_group_nickname_ext_buffer_parsing.py @@ -67,10 +67,10 @@ def test_parse_pattern_a_field1_username_field2_display(self): out = _load_group_nickname_map_from_contact_db(contact_db_path, chatroom, [username]) self.assertEqual(out.get(username), display) - def test_parse_pattern_b_field4_username_field1_display(self): + def test_parse_pattern_b_requires_unambiguous_field1_display(self): chatroom = "demo2@chatroom" username = "wxid_demo_abcdef" - display = "hjlbingo" + display = "旧版群昵称" inner = _enc_len(4, username.encode("utf-8")) + _enc_len(1, display.encode("utf-8")) ext_buffer = _member_entry(inner=inner) @@ -93,6 +93,115 @@ def test_parse_pattern_b_field4_username_field1_display(self): out = _load_group_nickname_map_from_contact_db(contact_db_path, chatroom, [username]) self.assertEqual(out.get(username), display) + def test_no_field2_field4_cross_reference_fails_closed(self): + chatroom = "demo-no-field2@chatroom" + member_alias = "member_alias_123" + target_username = "wxid_target_123456" + + first_inner = _enc_len(1, member_alias.encode("utf-8")) + _enc_len( + 4, target_username.encode("utf-8") + ) + second_inner = _enc_len(1, target_username.encode("utf-8")) + ext_buffer = _member_entry(inner=first_inner) + _member_entry(inner=second_inner) + + with TemporaryDirectory() as td: + contact_db_path = Path(td) / "contact.db" + conn = sqlite3.connect(str(contact_db_path)) + try: + conn.execute( + "CREATE TABLE chat_room(id INTEGER PRIMARY KEY, username TEXT, owner TEXT, ext_buffer BLOB)" + ) + conn.execute( + "INSERT INTO chat_room(id, username, owner, ext_buffer) VALUES (?, ?, ?, ?)", + (1, chatroom, "", ext_buffer), + ) + conn.commit() + finally: + conn.close() + + out = _load_group_nickname_map_from_contact_db( + contact_db_path, + chatroom, + [target_username], + ) + self.assertEqual(out, {}) + + def test_ambiguous_ascii_field1_does_not_confirm_legacy_layout(self): + chatroom = "demo-ambiguous-legacy@chatroom" + username = "wxid_demo_abcdef" + ambiguous_display = "hjlbingo" + + inner = _enc_len(4, username.encode("utf-8")) + _enc_len( + 1, ambiguous_display.encode("utf-8") + ) + ext_buffer = _member_entry(inner=inner) + + with TemporaryDirectory() as td: + contact_db_path = Path(td) / "contact.db" + conn = sqlite3.connect(str(contact_db_path)) + try: + conn.execute( + "CREATE TABLE chat_room(id INTEGER PRIMARY KEY, username TEXT, owner TEXT, ext_buffer BLOB)" + ) + conn.execute( + "INSERT INTO chat_room(id, username, owner, ext_buffer) VALUES (?, ?, ?, ?)", + (1, chatroom, "", ext_buffer), + ) + conn.commit() + finally: + conn.close() + + out = _load_group_nickname_map_from_contact_db( + contact_db_path, + chatroom, + [username], + ) + self.assertEqual(out, {}) + + def test_field4_reference_does_not_steal_another_members_nickname(self): + chatroom = "demo3@chatroom" + first_username = "wxid_first_123456" + second_username = "wxid_second_123456" + first_display = "第一位群昵称" + second_display = "第二位群昵称" + + first_inner = ( + _enc_len(1, first_username.encode("utf-8")) + + _enc_len(2, first_display.encode("utf-8")) + + _enc_len(4, second_username.encode("utf-8")) + ) + second_inner = _enc_len(1, second_username.encode("utf-8")) + _enc_len( + 2, second_display.encode("utf-8") + ) + ext_buffer = _member_entry(inner=first_inner) + _member_entry(inner=second_inner) + + with TemporaryDirectory() as td: + contact_db_path = Path(td) / "contact.db" + conn = sqlite3.connect(str(contact_db_path)) + try: + conn.execute( + "CREATE TABLE chat_room(id INTEGER PRIMARY KEY, username TEXT, owner TEXT, ext_buffer BLOB)" + ) + conn.execute( + "INSERT INTO chat_room(id, username, owner, ext_buffer) VALUES (?, ?, ?, ?)", + (1, chatroom, "", ext_buffer), + ) + conn.commit() + finally: + conn.close() + + out = _load_group_nickname_map_from_contact_db( + contact_db_path, + chatroom, + [second_username], + ) + self.assertEqual( + out, + { + second_username: second_display, + }, + ) + def test_non_chatroom_returns_empty(self): with TemporaryDirectory() as td: contact_db_path = Path(td) / "contact.db"