214 lines
10 KiB
Python
214 lines
10 KiB
Python
import hashlib
|
|
import hmac
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import struct
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from wechat_tool.cipher import (PAGE, RESERVE, aes_cbc, decrypt_database,
|
|
hmac_key, matches, wal_checksum)
|
|
from wechat_tool.mac_connect import MacConnection, parse_keys
|
|
from wechat_tool.wechat_native import message_rows
|
|
from wechat_tool import store
|
|
|
|
|
|
def sqlite_fixture(text):
|
|
db = sqlite3.connect(':memory:')
|
|
db.execute('PRAGMA page_size=4096')
|
|
db.execute('CREATE TABLE t (value TEXT)')
|
|
db.execute('INSERT INTO t VALUES (?)', (text,))
|
|
data = db.serialize()
|
|
db.close()
|
|
return reserved_pages(data)
|
|
|
|
|
|
def reserved_pages(data):
|
|
pages = []
|
|
# Repack these small, single-leaf fixture pages with SQLCipher's reserved tail.
|
|
for index in range(0, len(data), PAGE):
|
|
page = bytearray(data[index:index+PAGE])
|
|
offset = 100 if index == 0 else 0
|
|
assert page[offset] == 13 # no overflow, free blocks or internal nodes
|
|
count = int.from_bytes(page[offset+3:offset+5], 'big')
|
|
start = int.from_bytes(page[offset+5:offset+7], 'big') or PAGE
|
|
cells = bytes(page[start:])
|
|
page[start-RESERVE:PAGE-RESERVE] = cells
|
|
page[PAGE-RESERVE:] = bytes(RESERVE)
|
|
page[offset+5:offset+7] = (start-RESERVE).to_bytes(2, 'big')
|
|
for i in range(count):
|
|
pos = offset+8+2*i
|
|
value = int.from_bytes(page[pos:pos+2], 'big')
|
|
page[pos:pos+2] = (value-RESERVE).to_bytes(2, 'big')
|
|
if not index:
|
|
page[20] = RESERVE
|
|
pages.append(bytes(page))
|
|
return pages
|
|
|
|
|
|
def encrypted_page(plain, number, key, salt):
|
|
start = 16 if number == 1 else 0
|
|
iv = os.urandom(16)
|
|
ciphertext = aes_cbc(plain[start:-RESERVE], key, iv, True)
|
|
mac = hmac.new(hmac_key(key, salt), ciphertext+iv+struct.pack('<I', number), hashlib.sha512).digest()
|
|
return (salt if number == 1 else b'') + ciphertext + iv + mac
|
|
|
|
|
|
def native_fixture(root, invalid=False):
|
|
(root/'contact').mkdir(parents=True)
|
|
(root/'message').mkdir()
|
|
group = 'test-group@chatroom'
|
|
db = sqlite3.connect(root/'contact/contact.db')
|
|
db.execute('CREATE TABLE contact (username TEXT, nick_name TEXT, remark TEXT)')
|
|
db.executemany('INSERT INTO contact VALUES (?,?,?)', [(group,'测试群',''),('sender-a','张工','张工程师')])
|
|
db.commit()
|
|
db.close()
|
|
table = 'Msg_'+hashlib.md5(group.encode()).hexdigest()
|
|
db = sqlite3.connect(root/'message/message_0.db')
|
|
db.execute('CREATE TABLE Name2Id (user_name TEXT)')
|
|
db.executemany('INSERT INTO Name2Id VALUES (?)', [(group,),('sender-a',)])
|
|
db.execute('CREATE TABLE '+table+' (local_id INTEGER, server_id INTEGER, create_time INTEGER, message_content, local_type INTEGER, real_sender_id INTEGER)')
|
|
db.execute('INSERT INTO '+table+' VALUES (1,1001,1790816400,?,1,2)', ('测试比例怎么确定？',))
|
|
if invalid:
|
|
db.execute('INSERT INTO '+table+' VALUES (2,1002,1790816401,?,1,999)', ('不能错误归属',))
|
|
db.commit()
|
|
db.close()
|
|
|
|
|
|
class CipherTests(unittest.TestCase):
|
|
def test_native_aes_nist_vector(self):
|
|
key = bytes.fromhex('603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4')
|
|
iv = bytes.fromhex('000102030405060708090a0b0c0d0e0f')
|
|
plain = bytes.fromhex('6bc1bee22e409f96e93d7e117393172a')
|
|
expected = bytes.fromhex('f58c4c04d6e5f1ba779eabfb5f7bfbd6')
|
|
self.assertEqual(aes_cbc(plain,key,iv,True), expected)
|
|
self.assertEqual(aes_cbc(expected,key,iv), plain)
|
|
|
|
def test_authenticated_decryption_and_corruption(self):
|
|
key, salt = os.urandom(32), os.urandom(16)
|
|
pages = sqlite_fixture('正确内容')
|
|
data = b''.join(encrypted_page(p,i,key,salt) for i,p in enumerate(pages,1))
|
|
self.assertTrue(matches(data[:PAGE], key))
|
|
self.assertFalse(matches(data[:PAGE], os.urandom(32)))
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
src, dst = Path(directory)/'source.db', Path(directory)/'readable.db'
|
|
src.write_bytes(data)
|
|
decrypt_database(src,dst,key)
|
|
db=sqlite3.connect(dst)
|
|
self.assertEqual(db.execute('SELECT value FROM t').fetchone()[0], '正确内容')
|
|
db.close()
|
|
self.assertEqual(src.read_bytes(),data)
|
|
broken=bytearray(data);broken[PAGE+30]^=1;src.write_bytes(broken)
|
|
other=Path(directory)/'broken.db'
|
|
with self.assertRaises(ValueError):
|
|
decrypt_database(src,other,key)
|
|
self.assertFalse(other.exists())
|
|
with self.assertRaises(FileExistsError):
|
|
src.write_bytes(data)
|
|
decrypt_database(src,dst,key)
|
|
self.assertTrue(dst.exists())
|
|
|
|
def test_wal_replays_only_commits_and_rejects_bad_checksum(self):
|
|
key,salt=os.urandom(32),os.urandom(16)
|
|
before,after,last=sqlite_fixture('before'),sqlite_fixture('committed'),sqlite_fixture('uncommitted')
|
|
main=b''.join(encrypted_page(p,i,key,salt) for i,p in enumerate(before,1))
|
|
header=struct.pack('>IIIIII',0x377f0682,3007000,PAGE,0,123,456)
|
|
checksum=wal_checksum(header)
|
|
wal=bytearray(header+struct.pack('>II',*checksum))
|
|
for number,page in list(enumerate(after,1))+[(2,last[1])]:
|
|
count=len(after) if len(wal)==32+(len(after)-1)*(PAGE+24) else 0
|
|
payload=encrypted_page(page,number,key,salt)
|
|
first=struct.pack('>II',number,count)
|
|
checksum=wal_checksum(first+payload,checksum)
|
|
wal.extend(first+struct.pack('>II',123,456)+struct.pack('>II',*checksum)+payload)
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
src,dst=Path(directory)/'source.db',Path(directory)/'out.db'
|
|
src.write_bytes(main);walpath=Path(str(src)+'-wal');walpath.write_bytes(wal)
|
|
result=decrypt_database(src,dst,key)
|
|
self.assertEqual(result['wal_frames'],len(after))
|
|
db=sqlite3.connect(dst)
|
|
self.assertEqual(db.execute('SELECT value FROM t').fetchone()[0],'committed')
|
|
db.close()
|
|
wal[-1]^=1;walpath.write_bytes(wal)
|
|
with self.assertRaises(ValueError):
|
|
decrypt_database(src,Path(directory)/'bad.db',key)
|
|
|
|
|
|
class NativeTests(unittest.TestCase):
|
|
def test_encrypted_native_pipeline(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root=Path(directory)/'source'/'account'/'db_storage'
|
|
native_fixture(root)
|
|
keys, originals = [], {}
|
|
for path in MacConnection.files(root):
|
|
key,salt=os.urandom(32),os.urandom(16)
|
|
keys.append(key)
|
|
pages=reserved_pages(path.read_bytes())
|
|
data=b''.join(encrypted_page(p,i,key,salt) for i,p in enumerate(pages,1))
|
|
path.write_bytes(data)
|
|
originals[path]=data
|
|
service=MacConnection(Path(directory)/'workspace');service.job={'status':'running'}
|
|
with patch('wechat_tool.mac_connect.wechat_pid',return_value=None):
|
|
result=service.import_account(root,keys,False)
|
|
self.assertEqual(result['added'],1)
|
|
for path,data in originals.items():
|
|
self.assertEqual(path.read_bytes(),data)
|
|
db=store.connect(service.data)
|
|
self.assertEqual(db.execute('SELECT sender FROM messages').fetchone()[0],'张工程师')
|
|
db.close()
|
|
|
|
def test_key_json_variants_never_return_raw_keys(self):
|
|
key='ab'*32
|
|
self.assertEqual(parse_keys({'a':{'key_hex':key},'b':{'key':"x'"+key+'cd'*16+"'"}}),[bytes.fromhex(key)])
|
|
with self.assertRaises(ValueError):
|
|
parse_keys({'password':'微信密码不是密钥'})
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
service=MacConnection(directory)
|
|
state=service.load_keys({'key':key})
|
|
self.assertTrue(state['keys_ready'])
|
|
self.assertNotIn(key,json.dumps(state))
|
|
self.assertEqual(list(Path(directory).rglob('*')),[])
|
|
|
|
def test_native_mapping_and_atomic_idempotent_import(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root=Path(directory)/'source'/'account'/'db_storage'
|
|
native_fixture(root)
|
|
rows=list(message_rows(root,'account'))
|
|
self.assertEqual(rows[0]['sender'],'张工程师')
|
|
self.assertEqual(rows[0]['kind'],'group')
|
|
self.assertEqual(rows[0]['conversation'],'测试群')
|
|
service=MacConnection(Path(directory)/'workspace')
|
|
service.job={'status':'running'}
|
|
with patch('wechat_tool.mac_connect.wechat_pid',return_value=None):
|
|
one=service.import_account(root,[],False)
|
|
two=service.import_account(root,[],False)
|
|
self.assertEqual((one['added'],two['duplicates']),(1,1))
|
|
db=store.connect(service.data)
|
|
self.assertEqual(db.execute('SELECT COUNT(*) FROM messages').fetchone()[0],1)
|
|
db.close()
|
|
|
|
def test_unknown_identity_aborts_entire_import(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root=Path(directory)/'source'/'account'/'db_storage'
|
|
native_fixture(root,True)
|
|
service=MacConnection(Path(directory)/'workspace');service.job={'status':'running'}
|
|
with patch('wechat_tool.mac_connect.wechat_pid',return_value=None):
|
|
with self.assertRaises(ValueError):
|
|
service.import_account(root,[],False)
|
|
db=store.connect(service.data)
|
|
self.assertEqual(db.execute('SELECT COUNT(*) FROM messages').fetchone()[0],0)
|
|
db.close()
|
|
|
|
def test_running_wechat_and_non_mac_are_actionable(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
service=MacConnection(directory)
|
|
with patch('wechat_tool.mac_connect.wechat_pid',return_value=123):
|
|
with self.assertRaisesRegex(ValueError,'退出微信'):
|
|
service.import_account(Path(directory),[],False)
|
|
with patch('wechat_tool.mac_connect.platform.system',return_value='Windows'):
|
|
with self.assertRaisesRegex(ValueError,'目标 Mac'):
|
|
service.scan()
|