chuan/WeChatDataAnalysis
1
0
Fork 0
mirror of https://github.com/LifeArchiveProject/WeChatDataAnalysis.git synced 2026-02-19 14:20:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(contacts): 新增联系人列表与导出能力

Browse Source
This commit is contained in:
2977094657
2026-02-09 00:15:07 +08:00
parent 36d9af2b28
commit 62f396e55b
8 changed files with 1999 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
tests/test_contact_type_detection.py Normal file
View File

@@ -0,0 +1,71 @@
import sys
import unittest
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
class TestContactTypeDetection(unittest.TestCase):
def test_infer_group(self):
from wechat_decrypt_tool.routers.chat_contacts import _infer_contact_type
row = {"local_type": 0, "alias": "", "remark": "", "nick_name": ""}
self.assertEqual(_infer_contact_type("123@chatroom", row), "group")
def test_infer_official_by_prefix(self):
from wechat_decrypt_tool.routers.chat_contacts import _infer_contact_type
row = {"local_type": 0, "verify_flag": 0, "alias": "", "remark": "", "nick_name": ""}
self.assertEqual(_infer_contact_type("gh_xxx", row), "official")
def test_infer_official_by_verify_flag(self):
from wechat_decrypt_tool.routers.chat_contacts import _infer_contact_type
row = {"local_type": 1, "verify_flag": 24, "alias": "", "remark": "", "nick_name": ""}
self.assertEqual(_infer_contact_type("wxid_xxx", row), "official")
def test_infer_none_for_local_type_3_without_verify(self):
from wechat_decrypt_tool.routers.chat_contacts import _infer_contact_type
row = {"local_type": 3, "verify_flag": 0, "alias": "", "remark": "", "nick_name": "普通联系人"}
self.assertIsNone(_infer_contact_type("wxid_xxx", row))
def test_infer_none_from_wxid_alias_when_local_type_not_1(self):
from wechat_decrypt_tool.routers.chat_contacts import _infer_contact_type
row = {"local_type": 0, "verify_flag": 0, "alias": "wechat_id", "remark": "", "nick_name": ""}
self.assertIsNone(_infer_contact_type("wxid_xxx", row))
def test_infer_friend_from_local_type_1(self):
from wechat_decrypt_tool.routers.chat_contacts import _infer_contact_type
row = {"local_type": 1, "verify_flag": 0, "alias": "", "remark": "", "nick_name": ""}
self.assertEqual(_infer_contact_type("wxid_xxx", row), "friend")
def test_infer_none_from_local_type_2(self):
from wechat_decrypt_tool.routers.chat_contacts import _infer_contact_type
row = {"local_type": 2, "verify_flag": 0, "alias": "", "remark": "", "nick_name": ""}
self.assertIsNone(_infer_contact_type("wxid_xxx", row))
def test_infer_none_when_empty_type_0(self):
from wechat_decrypt_tool.routers.chat_contacts import _infer_contact_type
row = {"local_type": 0, "verify_flag": 0, "alias": "", "remark": "", "nick_name": ""}
self.assertIsNone(_infer_contact_type("wxid_xxx", row))
def test_valid_contact_username_filters_system_accounts(self):
from wechat_decrypt_tool.routers.chat_contacts import _is_valid_contact_username
self.assertFalse(_is_valid_contact_username("filehelper"))
self.assertFalse(_is_valid_contact_username("notifymessage"))
self.assertFalse(_is_valid_contact_username("fake_abc"))
self.assertTrue(_is_valid_contact_username("weixin"))
self.assertTrue(_is_valid_contact_username("wxid_abc"))
self.assertTrue(_is_valid_contact_username("123@chatroom"))
if __name__ == "__main__":
unittest.main()

546
tests/test_contacts_export.py Normal file
View File

@@ -0,0 +1,546 @@
import json
import os
import sqlite3
import sys
import unittest
import importlib
from pathlib import Path
from tempfile import TemporaryDirectory
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
class TestContactsExport(unittest.TestCase):
@staticmethod
def _encode_varint(value: int) -> bytes:
v = int(value)
out = bytearray()
while True:
b = v & 0x7F
v >>= 7
if v:
out.append(b | 0x80)
else:
out.append(b)
break
return bytes(out)
@classmethod
def _encode_field_len(cls, field_no: int, raw: bytes) -> bytes:
tag = (int(field_no) << 3) | 2
payload = bytes(raw)
return cls._encode_varint(tag) + cls._encode_varint(len(payload)) + payload
@classmethod
def _encode_field_varint(cls, field_no: int, value: int) -> bytes:
tag = int(field_no) << 3
return cls._encode_varint(tag) + cls._encode_varint(int(value))
@classmethod
def _build_extra_buffer(cls, *, country: str, province: str, city: str, source_scene: int) -> bytes:
return b"".join(
[
cls._encode_field_len(5, country.encode("utf-8")),
cls._encode_field_len(6, province.encode("utf-8")),
cls._encode_field_len(7, city.encode("utf-8")),
cls._encode_field_varint(8, source_scene),
]
)
def _seed_contact_db(self, path: Path) -> None:
conn = sqlite3.connect(str(path))
try:
conn.execute(
"""
CREATE TABLE contact (
username TEXT,
remark TEXT,
nick_name TEXT,
alias TEXT,
local_type INTEGER,
verify_flag INTEGER,
big_head_url TEXT,
small_head_url TEXT,
extra_buffer BLOB
)
"""
)
conn.execute(
"""
CREATE TABLE stranger (
username TEXT,
remark TEXT,
nick_name TEXT,
alias TEXT,
local_type INTEGER,
verify_flag INTEGER,
big_head_url TEXT,
small_head_url TEXT,
extra_buffer BLOB
)
"""
)
friend_extra_buffer = self._build_extra_buffer(
country="CN",
province="Sichuan",
city="Chengdu",
source_scene=14,
)
conn.execute(
"INSERT INTO contact VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
(
"wxid_friend",
"好友备注",
"好友昵称",
"friend_alias",
1,
0,
"https://cdn.example.com/friend_big.jpg",
"https://cdn.example.com/friend_small.jpg",
friend_extra_buffer,
),
)
conn.execute(
"INSERT INTO contact VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
(
"room@chatroom",
"",
"测试群",
"",
0,
0,
"https://cdn.example.com/group_big.jpg",
"",
b"",
),
)
conn.execute(
"INSERT INTO contact VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
(
"gh_official",
"",
"公众号",
"",
4,
8,
"",
"https://cdn.example.com/official_small.jpg",
b"",
),
)
conn.execute(
"INSERT INTO contact VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
(
"wxid_local_type_3",
"",
"不应计入联系人",
"",
3,
0,
"",
"",
b"",
),
)
conn.execute(
"INSERT INTO contact VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
(
"weixin",
"",
"微信团队",
"",
1,
56,
"",
"",
b"",
),
)
conn.execute(
"INSERT INTO contact VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
(
"filehelper",
"",
"文件传输助手",
"",
0,
0,
"",
"",
b"",
),
)
conn.execute(
"INSERT INTO stranger VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
(
"stranger_verified",
"",
"陌生人认证号",
"",
4,
24,
"",
"",
b"",
),
)
conn.commit()
finally:
conn.close()
def _seed_session_db(self, path: Path) -> None:
conn = sqlite3.connect(str(path))
try:
conn.execute(
"""
CREATE TABLE SessionTable (
username TEXT,
sort_timestamp INTEGER,
last_timestamp INTEGER
)
"""
)
conn.execute("INSERT INTO SessionTable VALUES (?, ?, ?)", ("room@chatroom", 300, 300))
conn.execute("INSERT INTO SessionTable VALUES (?, ?, ?)", ("wxid_friend", 200, 200))
conn.execute("INSERT INTO SessionTable VALUES (?, ?, ?)", ("gh_official", 100, 100))
conn.execute("INSERT INTO SessionTable VALUES (?, ?, ?)", ("missing@chatroom", 250, 250))
conn.commit()
finally:
conn.close()
def _seed_contact_db_legacy(self, path: Path) -> None:
conn = sqlite3.connect(str(path))
try:
conn.execute(
"""
CREATE TABLE contact (
username TEXT,
remark TEXT,
nick_name TEXT,
alias TEXT,
local_type INTEGER,
verify_flag INTEGER,
big_head_url TEXT,
small_head_url TEXT
)
"""
)
conn.execute(
"""
CREATE TABLE stranger (
username TEXT,
remark TEXT,
nick_name TEXT,
alias TEXT,
local_type INTEGER,
verify_flag INTEGER,
big_head_url TEXT,
small_head_url TEXT
)
"""
)
conn.execute(
"INSERT INTO contact VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
(
"wxid_legacy_friend",
"旧版好友备注",
"旧版好友昵称",
"legacy_friend_alias",
1,
0,
"",
"",
),
)
conn.commit()
finally:
conn.close()
def test_export_json_and_csv(self):
from fastapi import FastAPI
from fastapi.testclient import TestClient
with TemporaryDirectory() as td:
root = Path(td)
account = "wxid_test"
account_dir = root / "output" / "databases" / account
account_dir.mkdir(parents=True, exist_ok=True)
self._seed_contact_db(account_dir / "contact.db")
self._seed_session_db(account_dir / "session.db")
prev = None
try:
prev = os.environ.get("WECHAT_TOOL_DATA_DIR")
os.environ["WECHAT_TOOL_DATA_DIR"] = str(root)
import wechat_decrypt_tool.chat_helpers as chat_helpers
import wechat_decrypt_tool.routers.chat_contacts as chat_contacts
importlib.reload(chat_helpers)
importlib.reload(chat_contacts)
app = FastAPI()
app.include_router(chat_contacts.router)
client = TestClient(app)
list_resp = client.get(
"/api/chat/contacts",
params={
"account": account,
"include_friends": True,
"include_groups": True,
"include_officials": True,
},
)
self.assertEqual(list_resp.status_code, 200)
list_payload = list_resp.json()
self.assertEqual(list_payload["status"], "success")
self.assertEqual(list_payload["total"], 6)
self.assertEqual(list_payload["counts"]["friends"], 1)
self.assertEqual(list_payload["counts"]["groups"], 2)
self.assertEqual(list_payload["counts"]["officials"], 3)
usernames = {str(x.get("username")) for x in list_payload.get("contacts", [])}
self.assertIn("missing@chatroom", usernames)
self.assertIn("weixin", usernames)
self.assertNotIn("wxid_local_type_3", usernames)
first = list_payload["contacts"][0]
self.assertIn("avatarLink", first)
friend_contact = next(
(x for x in list_payload.get("contacts", []) if str(x.get("username")) == "wxid_friend"),
{},
)
self.assertEqual(friend_contact.get("country"), "CN")
self.assertEqual(friend_contact.get("province"), "Sichuan")
self.assertEqual(friend_contact.get("city"), "Chengdu")
self.assertEqual(friend_contact.get("region"), "中国大陆·Sichuan·Chengdu")
self.assertEqual(friend_contact.get("sourceScene"), 14)
self.assertEqual(friend_contact.get("source"), "通过群聊添加")
export_dir = root / "exports"
export_dir.mkdir(parents=True, exist_ok=True)
json_resp = client.post(
"/api/chat/contacts/export",
json={
"account": account,
"output_dir": str(export_dir),
"format": "json",
"include_avatar_link": True,
"contact_types": {
"friends": True,
"groups": True,
"officials": True,
},
},
)
self.assertEqual(json_resp.status_code, 200)
json_payload = json_resp.json()
self.assertEqual(json_payload["status"], "success")
self.assertEqual(json_payload["count"], 6)
json_path = Path(json_payload["outputPath"])
self.assertTrue(json_path.exists())
data = json.loads(json_path.read_text(encoding="utf-8"))
self.assertEqual(data["count"], 6)
self.assertIn("avatarLink", data["contacts"][0])
self.assertIn("region", data["contacts"][0])
self.assertIn("country", data["contacts"][0])
self.assertIn("province", data["contacts"][0])
self.assertIn("city", data["contacts"][0])
self.assertIn("source", data["contacts"][0])
self.assertIn("sourceScene", data["contacts"][0])
export_usernames = {str(x.get("username")) for x in data.get("contacts", [])}
self.assertIn("missing@chatroom", export_usernames)
self.assertNotIn("wxid_local_type_3", export_usernames)
friend_export = next(
(x for x in data.get("contacts", []) if str(x.get("username")) == "wxid_friend"),
{},
)
self.assertEqual(friend_export.get("region"), "中国大陆·Sichuan·Chengdu")
self.assertEqual(friend_export.get("sourceScene"), 14)
self.assertEqual(friend_export.get("source"), "通过群聊添加")
csv_resp = client.post(
"/api/chat/contacts/export",
json={
"account": account,
"output_dir": str(export_dir),
"format": "csv",
"include_avatar_link": False,
"contact_types": {
"friends": True,
"groups": False,
"officials": False,
},
},
)
self.assertEqual(csv_resp.status_code, 200)
csv_payload = csv_resp.json()
self.assertEqual(csv_payload["count"], 1)
csv_path = Path(csv_payload["outputPath"])
text = csv_path.read_text(encoding="utf-8-sig")
self.assertIn("用户名,显示名称,备注,昵称,微信号,类型,地区,国家/地区码,省份,城市,来源,来源场景码", text.splitlines()[0])
self.assertNotIn("头像链接", text.splitlines()[0])
self.assertIn("wxid_friend", text)
self.assertIn("中国大陆·Sichuan·Chengdu", text)
self.assertIn("通过群聊添加", text)
self.assertIn(",14", text)
self.assertNotIn("wxid_local_type_3", text)
finally:
if prev is None:
os.environ.pop("WECHAT_TOOL_DATA_DIR", None)
else:
os.environ["WECHAT_TOOL_DATA_DIR"] = prev
def test_export_invalid_format_returns_400(self):
from fastapi import FastAPI
from fastapi.testclient import TestClient
with TemporaryDirectory() as td:
root = Path(td)
account = "wxid_test"
account_dir = root / "output" / "databases" / account
account_dir.mkdir(parents=True, exist_ok=True)
self._seed_contact_db(account_dir / "contact.db")
self._seed_session_db(account_dir / "session.db")
prev = None
try:
prev = os.environ.get("WECHAT_TOOL_DATA_DIR")
os.environ["WECHAT_TOOL_DATA_DIR"] = str(root)
import wechat_decrypt_tool.chat_helpers as chat_helpers
import wechat_decrypt_tool.routers.chat_contacts as chat_contacts
importlib.reload(chat_helpers)
importlib.reload(chat_contacts)
app = FastAPI()
app.include_router(chat_contacts.router)
client = TestClient(app)
resp = client.post(
"/api/chat/contacts/export",
json={
"account": account,
"output_dir": str(root / "exports"),
"format": "vcf",
"include_avatar_link": True,
"contact_types": {
"friends": True,
"groups": True,
"officials": True,
},
},
)
self.assertEqual(resp.status_code, 400)
finally:
if prev is None:
os.environ.pop("WECHAT_TOOL_DATA_DIR", None)
else:
os.environ["WECHAT_TOOL_DATA_DIR"] = prev
def test_missing_contact_db_returns_404(self):
from fastapi import FastAPI
from fastapi.testclient import TestClient
with TemporaryDirectory() as td:
root = Path(td)
account = "wxid_test"
account_dir = root / "output" / "databases" / account
account_dir.mkdir(parents=True, exist_ok=True)
# only session.db exists
self._seed_session_db(account_dir / "session.db")
prev = None
try:
prev = os.environ.get("WECHAT_TOOL_DATA_DIR")
os.environ["WECHAT_TOOL_DATA_DIR"] = str(root)
import wechat_decrypt_tool.chat_helpers as chat_helpers
import wechat_decrypt_tool.routers.chat_contacts as chat_contacts
importlib.reload(chat_helpers)
importlib.reload(chat_contacts)
app = FastAPI()
app.include_router(chat_contacts.router)
client = TestClient(app)
resp = client.get("/api/chat/contacts", params={"account": account})
self.assertEqual(resp.status_code, 404)
finally:
if prev is None:
os.environ.pop("WECHAT_TOOL_DATA_DIR", None)
else:
os.environ["WECHAT_TOOL_DATA_DIR"] = prev
def test_legacy_schema_without_extra_buffer_is_compatible(self):
from fastapi import FastAPI
from fastapi.testclient import TestClient
with TemporaryDirectory() as td:
root = Path(td)
account = "wxid_legacy"
account_dir = root / "output" / "databases" / account
account_dir.mkdir(parents=True, exist_ok=True)
self._seed_contact_db_legacy(account_dir / "contact.db")
self._seed_session_db(account_dir / "session.db")
prev = None
try:
prev = os.environ.get("WECHAT_TOOL_DATA_DIR")
os.environ["WECHAT_TOOL_DATA_DIR"] = str(root)
import wechat_decrypt_tool.chat_helpers as chat_helpers
import wechat_decrypt_tool.routers.chat_contacts as chat_contacts
importlib.reload(chat_helpers)
importlib.reload(chat_contacts)
app = FastAPI()
app.include_router(chat_contacts.router)
client = TestClient(app)
resp = client.get(
"/api/chat/contacts",
params={
"account": account,
"include_friends": True,
"include_groups": False,
"include_officials": False,
},
)
self.assertEqual(resp.status_code, 200)
payload = resp.json()
self.assertEqual(payload.get("status"), "success")
self.assertEqual(int(payload.get("total", 0)), 1)
contact = payload.get("contacts", [])[0]
self.assertEqual(contact.get("username"), "wxid_legacy_friend")
self.assertEqual(contact.get("country"), "")
self.assertEqual(contact.get("province"), "")
self.assertEqual(contact.get("city"), "")
self.assertEqual(contact.get("region"), "")
self.assertIsNone(contact.get("sourceScene"))
self.assertEqual(contact.get("source"), "")
finally:
if prev is None:
os.environ.pop("WECHAT_TOOL_DATA_DIR", None)
else:
os.environ["WECHAT_TOOL_DATA_DIR"] = prev
if __name__ == "__main__":
unittest.main()