199 lines
7.6 KiB
Python
199 lines
7.6 KiB
Python
"""
|
|
FridaManager 单元测试 (H30)
|
|
使用 Mock 模拟 frida 库，不需要真机/frida-server。
|
|
测试 FridaManager 的生命周期管理、状态管理和错误处理。
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch, PropertyMock
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "agent"))
|
|
|
|
|
|
class TestFridaManagerInit(unittest.TestCase):
|
|
"""测试 FridaManager 初始化和属性"""
|
|
|
|
def test_initial_state(self):
|
|
from hook.frida_manager import FridaManager
|
|
mgr = FridaManager(device_serial="abc123")
|
|
self.assertEqual(mgr.device_serial, "abc123")
|
|
self.assertFalse(mgr.connected)
|
|
self.assertIsNone(mgr.rpc)
|
|
self.assertFalse(mgr._running)
|
|
|
|
def test_default_script_path(self):
|
|
from hook.frida_manager import FridaManager, DEFAULT_SCRIPT
|
|
mgr = FridaManager()
|
|
self.assertEqual(mgr.script_path, DEFAULT_SCRIPT)
|
|
self.assertTrue(mgr.script_path.endswith("wechat_hook_v2.js"))
|
|
|
|
def test_get_status_disconnected(self):
|
|
from hook.frida_manager import FridaManager
|
|
mgr = FridaManager(device_serial="test_serial")
|
|
status = mgr.get_status()
|
|
self.assertFalse(status["connected"])
|
|
self.assertIsNone(status["device"])
|
|
self.assertEqual(status["device_serial"], "test_serial")
|
|
self.assertFalse(status["running"])
|
|
|
|
|
|
class TestFridaManagerStartNoFrida(unittest.TestCase):
|
|
"""frida 库未安装时应优雅失败"""
|
|
|
|
@patch.dict("sys.modules", {"frida": None})
|
|
def test_start_without_frida_returns_false(self):
|
|
from hook.frida_manager import FridaManager
|
|
mgr = FridaManager()
|
|
with patch("builtins.__import__", side_effect=ImportError("No module named 'frida'")):
|
|
result = mgr.start()
|
|
self.assertFalse(result)
|
|
|
|
|
|
class TestFridaManagerWithMockFrida(unittest.TestCase):
|
|
"""模拟 frida 库测试核心流程"""
|
|
|
|
def setUp(self):
|
|
self.mock_frida = MagicMock()
|
|
self.mock_device = MagicMock()
|
|
self.mock_device.name = "Test Xiaomi"
|
|
self.mock_session = MagicMock()
|
|
self.mock_script = MagicMock()
|
|
self.mock_rpc = MagicMock()
|
|
self.mock_rpc.ping.return_value = "pong from wechat_hook_v2.1"
|
|
|
|
self.mock_device.enumerate_processes.return_value = [
|
|
MagicMock(name="com.tencent.mm", identifier="com.tencent.mm", pid=12345)
|
|
]
|
|
self.mock_device.attach.return_value = self.mock_session
|
|
self.mock_session.create_script.return_value = self.mock_script
|
|
self.mock_script.exports_sync = self.mock_rpc
|
|
|
|
@patch("hook.frida_manager.frida", create=True)
|
|
def test_call_rpc_disconnected(self, mock_frida_module):
|
|
from hook.frida_manager import FridaManager
|
|
mgr = FridaManager()
|
|
result = mgr.call_rpc("ping")
|
|
self.assertFalse(result["success"])
|
|
self.assertIn("未连接", result["error"])
|
|
|
|
def test_call_rpc_unknown_method(self):
|
|
from hook.frida_manager import FridaManager
|
|
mgr = FridaManager()
|
|
mgr._session = self.mock_session
|
|
mgr._script = self.mock_script
|
|
mgr._rpc = self.mock_rpc
|
|
self.mock_rpc.nonexistent = None
|
|
delattr(self.mock_rpc, "nonexistent")
|
|
result = mgr.call_rpc("nonexistent")
|
|
self.assertFalse(result["success"])
|
|
|
|
def test_call_rpc_success(self):
|
|
from hook.frida_manager import FridaManager
|
|
mgr = FridaManager()
|
|
mgr._session = self.mock_session
|
|
mgr._script = self.mock_script
|
|
mgr._rpc = MagicMock()
|
|
mgr._rpc.send_message.return_value = {"success": True, "message_id": "m123"}
|
|
result = mgr.call_rpc("sendMessage", {"to_id": "wxid_abc", "content": "hi"})
|
|
self.assertTrue(result["success"])
|
|
self.assertEqual(result["message_id"], "m123")
|
|
|
|
def test_call_rpc_exception(self):
|
|
from hook.frida_manager import FridaManager
|
|
mgr = FridaManager()
|
|
mgr._session = self.mock_session
|
|
mgr._script = self.mock_script
|
|
mgr._rpc = MagicMock()
|
|
mgr._rpc.send_message.side_effect = Exception("进程已退出")
|
|
result = mgr.call_rpc("sendMessage", {})
|
|
self.assertFalse(result["success"])
|
|
self.assertIn("进程已退出", result["error"])
|
|
|
|
def test_stop_cleans_up(self):
|
|
from hook.frida_manager import FridaManager
|
|
mgr = FridaManager()
|
|
mgr._session = self.mock_session
|
|
mgr._script = self.mock_script
|
|
mgr._running = True
|
|
mgr.stop()
|
|
self.assertFalse(mgr._running)
|
|
self.assertIsNone(mgr._script)
|
|
self.assertIsNone(mgr._session)
|
|
|
|
def test_on_message_log(self):
|
|
from hook.frida_manager import FridaManager
|
|
mgr = FridaManager()
|
|
message = {"type": "send", "payload": {"type": "log", "level": "info", "tag": "test", "message": "hello"}}
|
|
mgr._on_message(message, None)
|
|
|
|
def test_on_message_hook_event(self):
|
|
from hook.frida_manager import FridaManager
|
|
events = []
|
|
mgr = FridaManager(on_event=lambda e: events.append(e))
|
|
message = {"type": "send", "payload": {"type": "hook_event", "event_type": "message_received", "payload": {}}}
|
|
mgr._on_message(message, None)
|
|
self.assertEqual(len(events), 1)
|
|
self.assertEqual(events[0]["event_type"], "message_received")
|
|
|
|
def test_on_message_error(self):
|
|
from hook.frida_manager import FridaManager
|
|
mgr = FridaManager()
|
|
message = {"type": "error", "description": "Script destroyed"}
|
|
mgr._on_message(message, None)
|
|
|
|
|
|
class TestEventReporter(unittest.TestCase):
|
|
"""EventReporter 单元测试"""
|
|
|
|
def test_buffer_events_without_send_fn(self):
|
|
from hook.event_reporter import EventReporter
|
|
reporter = EventReporter()
|
|
reporter.on_hook_event({"event_type": "test", "payload": {"a": 1}})
|
|
reporter.on_hook_event({"event_type": "test2", "payload": {"b": 2}})
|
|
stats = reporter.get_stats()
|
|
self.assertEqual(stats["total"], 2)
|
|
self.assertEqual(stats["buffered"], 2)
|
|
self.assertEqual(stats["sent"], 0)
|
|
|
|
def test_send_events_with_fn(self):
|
|
from hook.event_reporter import EventReporter
|
|
sent = []
|
|
reporter = EventReporter(send_fn=lambda e: sent.append(e))
|
|
reporter.on_hook_event({"event_type": "msg", "payload": {}})
|
|
self.assertEqual(len(sent), 1)
|
|
stats = reporter.get_stats()
|
|
self.assertEqual(stats["sent"], 1)
|
|
|
|
def test_flush_buffer_on_set_send_fn(self):
|
|
from hook.event_reporter import EventReporter
|
|
reporter = EventReporter()
|
|
reporter.on_hook_event({"event_type": "buffered1", "payload": {}})
|
|
reporter.on_hook_event({"event_type": "buffered2", "payload": {}})
|
|
sent = []
|
|
reporter.set_send_fn(lambda e: sent.append(e))
|
|
self.assertEqual(len(sent), 2)
|
|
|
|
def test_get_buffered_events(self):
|
|
from hook.event_reporter import EventReporter
|
|
reporter = EventReporter()
|
|
for i in range(5):
|
|
reporter.on_hook_event({"event_type": f"ev_{i}", "payload": {}})
|
|
events = reporter.get_buffered_events(limit=3)
|
|
self.assertEqual(len(events), 3)
|
|
|
|
def test_send_failure_buffers(self):
|
|
from hook.event_reporter import EventReporter
|
|
def failing_send(e):
|
|
raise ConnectionError("offline")
|
|
reporter = EventReporter(send_fn=failing_send)
|
|
reporter.on_hook_event({"event_type": "fail", "payload": {}})
|
|
stats = reporter.get_stats()
|
|
self.assertEqual(stats["errors"], 1)
|
|
self.assertEqual(stats["buffered"], 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|