83 lines
4.1 KiB
Python
83 lines
4.1 KiB
Python
"""自动密钥提取：合成数据库/内存，不读取真实账号的内存。"""
|
|
import ctypes
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
import wxwork_key as scanner
|
|
from test_crypto_roundtrip import build_encrypted_db
|
|
|
|
|
|
class AutomaticKeyScannerTest(unittest.TestCase):
|
|
def setUp(self):
|
|
temporary = tempfile.TemporaryDirectory()
|
|
self.addCleanup(temporary.cleanup)
|
|
self.path = str(Path(temporary.name) / "message.db")
|
|
self.key = bytes(range(1, 17))
|
|
Path(self.path).write_bytes(build_encrypted_db(self.key))
|
|
|
|
def test_binary_patterns_and_chunk_boundaries_return_hex_and_stop_early(self):
|
|
patterns = [self.key + b"\x01\x00\x00\x00sAlT",
|
|
b"\x10\x00\x00\x00" + self.key,
|
|
b"raw:" + self.key, self.key.hex().encode()]
|
|
for pattern in patterns:
|
|
with self.subTest(pattern_length=len(pattern)):
|
|
memory = b"-" * 59 + pattern + b"-" * 512
|
|
pm = mock.Mock()
|
|
pm.regions.return_value = [(0, len(memory))]
|
|
pm.read.side_effect = lambda base, size: memory[base:base+size]
|
|
with mock.patch.object(scanner, "_ProcessMemory", return_value=pm), mock.patch.object(scanner, "_CHUNK", 64):
|
|
result = scanner._scan_pid(123, [self.path])
|
|
self.assertEqual(result, {self.path: [self.key.hex()]})
|
|
self.assertLessEqual(pm.read.call_count, 2 + 2 * ((len(memory) + 63) // 64))
|
|
pm.close.assert_called_once_with()
|
|
|
|
def test_completed_accounts_skip_remaining_processes(self):
|
|
with mock.patch.object(scanner, "find_wxwork_pids", return_value=[1, 2, 3]), \
|
|
mock.patch.object(scanner, "_scan_pid", return_value={self.path: [self.key.hex()]}) as scan:
|
|
result = scanner.extract_keys_from_running([self.path], timeout=50)
|
|
self.assertEqual(result, {self.path: [self.key.hex()]})
|
|
self.assertEqual(scan.call_count, 1)
|
|
|
|
def test_deadline_preserves_partial_results_and_closes_memory(self):
|
|
pm = mock.Mock()
|
|
pm.regions.return_value = [(0, 64), (64, 64)]
|
|
pm.read.return_value = self.key + b"\x01\x00\x00\x00sAlT"
|
|
with mock.patch.object(scanner, "_ProcessMemory", return_value=pm), \
|
|
mock.patch.object(scanner.time, "monotonic", side_effect=[0, 0, 20]):
|
|
result = scanner._scan_pid(123, [self.path, "absent.db"], deadline=10)
|
|
self.assertEqual(result, {self.path: [self.key.hex()]})
|
|
pm.close.assert_called_once_with()
|
|
|
|
def test_expired_budget_does_not_start_scanning(self):
|
|
with mock.patch.object(scanner, "find_wxwork_pids", return_value=[1]), \
|
|
mock.patch.object(scanner, "_scan_pid") as scan:
|
|
self.assertEqual(scanner.extract_keys_from_running([self.path], timeout=0), {})
|
|
scan.assert_not_called()
|
|
|
|
def test_native_memory_api_reads_own_test_buffer_at_full_pointer_width(self):
|
|
payload = b"synthetic-native-memory-probe"
|
|
buffer = ctypes.create_string_buffer(payload)
|
|
pm = scanner._ProcessMemory(os.getpid())
|
|
try:
|
|
address = ctypes.addressof(buffer)
|
|
self.assertEqual(pm.read(address, len(payload)), payload)
|
|
self.assertTrue(any(base <= address < base + size for base, size in pm.regions(rw_only=True)))
|
|
finally:
|
|
pm.close()
|
|
|
|
|
|
def test_fast_salt_pass_reaches_later_regions_before_expensive_fallback(self):
|
|
pm = mock.Mock()
|
|
memory = b"-" * 192 + self.key + b"\x01\x00\x00\x00sAlT"
|
|
pm.regions.return_value = [(0, len(memory))]
|
|
pm.read.side_effect = lambda base, size: memory[base:base + size]
|
|
with mock.patch.object(scanner, "_ProcessMemory", return_value=pm), \
|
|
mock.patch.object(scanner, "_CHUNK", 64), \
|
|
mock.patch.object(scanner, "_candidates_codec", side_effect=AssertionError("fallback ran too early")):
|
|
result = scanner._scan_pid(123, [self.path])
|
|
self.assertEqual(result, {self.path: [self.key.hex()]})
|
|
self.assertEqual(pm.read.call_count, 4)
|