351 lines
14 KiB
Python
351 lines
14 KiB
Python
"""
|
||
Hook 模块管理服务
|
||
|
||
实现能力：
|
||
1. 模块管理（增删改查/启停/scope）
|
||
2. 脚本管理（保存/下载/部署）
|
||
3. 设备模块状态（加载/重载/日志）
|
||
4. Hook 事件总线（历史 + 实时订阅）
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import base64
|
||
import hashlib
|
||
import json
|
||
import os
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, Set
|
||
|
||
|
||
def _now_iso() -> str:
|
||
return datetime.now(timezone.utc).isoformat()
|
||
|
||
|
||
class HookModuleService:
|
||
def __init__(self) -> None:
|
||
app_dir = Path(__file__).resolve().parent.parent
|
||
self.data_dir = app_dir / "data" / "hook"
|
||
self.scripts_dir = self.data_dir / "scripts"
|
||
self.modules_file = self.data_dir / "modules.json"
|
||
self.events_file = self.data_dir / "events.jsonl"
|
||
self.device_state_file = self.data_dir / "device_modules.json"
|
||
self._lock = asyncio.Lock()
|
||
self._event_clients: Set[Any] = set()
|
||
self._ensure_store()
|
||
|
||
def _ensure_store(self) -> None:
|
||
self.scripts_dir.mkdir(parents=True, exist_ok=True)
|
||
if not self.modules_file.exists():
|
||
self.modules_file.write_text("{}", encoding="utf-8")
|
||
if not self.device_state_file.exists():
|
||
self.device_state_file.write_text("{}", encoding="utf-8")
|
||
if not self.events_file.exists():
|
||
self.events_file.write_text("", encoding="utf-8")
|
||
|
||
def _read_json(self, path: Path, default: Any) -> Any:
|
||
try:
|
||
return json.loads(path.read_text(encoding="utf-8"))
|
||
except Exception:
|
||
return default
|
||
|
||
def _write_json(self, path: Path, data: Any) -> None:
|
||
path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
|
||
async def list_modules(self, enabled: Optional[bool] = None, scope: Optional[str] = None) -> List[dict]:
|
||
async with self._lock:
|
||
modules = self._read_json(self.modules_file, {})
|
||
result = list(modules.values())
|
||
if enabled is not None:
|
||
result = [m for m in result if bool(m.get("enabled")) == enabled]
|
||
if scope:
|
||
result = [m for m in result if scope in (m.get("scopes") or [])]
|
||
return sorted(result, key=lambda x: x.get("updated_at", ""), reverse=True)
|
||
|
||
async def get_module(self, module_id: str) -> Optional[dict]:
|
||
async with self._lock:
|
||
modules = self._read_json(self.modules_file, {})
|
||
return modules.get(module_id)
|
||
|
||
async def upsert_module(self, payload: dict) -> dict:
|
||
module_id = payload["module_id"]
|
||
now = _now_iso()
|
||
script_content = payload.pop("script_content", None)
|
||
script_url = payload.get("script_url")
|
||
script_hash = payload.get("script_hash")
|
||
if script_content and not script_url:
|
||
decoded = base64.b64decode(script_content)
|
||
script_name = f"{module_id}.js"
|
||
target = self.scripts_dir / script_name
|
||
target.write_bytes(decoded)
|
||
digest = hashlib.sha256(decoded).hexdigest()
|
||
script_url = f"/api/v3/scripts/{module_id}"
|
||
script_hash = f"sha256:{digest}"
|
||
|
||
async with self._lock:
|
||
modules = self._read_json(self.modules_file, {})
|
||
old = modules.get(module_id, {})
|
||
data = {
|
||
"module_id": module_id,
|
||
"name": payload.get("name", old.get("name", module_id)),
|
||
"version": payload.get("version", old.get("version", "1.0.0")),
|
||
"description": payload.get("description", old.get("description", "")),
|
||
"enabled": bool(payload.get("enabled", old.get("enabled", True))),
|
||
"scopes": payload.get("scopes", old.get("scopes", ["com.tencent.mm"])),
|
||
"capabilities": payload.get("capabilities", old.get("capabilities", [])),
|
||
"min_frida_version": payload.get("min_frida_version", old.get("min_frida_version", "16.0.0")),
|
||
"script_url": script_url or old.get("script_url"),
|
||
"script_hash": script_hash or old.get("script_hash"),
|
||
"script_id": payload.get("script_id", old.get("script_id", module_id)),
|
||
"device_count": old.get("device_count", 0),
|
||
"error_count": old.get("error_count", 0),
|
||
"created_at": old.get("created_at", now),
|
||
"updated_at": now,
|
||
}
|
||
modules[module_id] = data
|
||
self._write_json(self.modules_file, modules)
|
||
return data
|
||
|
||
async def delete_module(self, module_id: str) -> bool:
|
||
async with self._lock:
|
||
modules = self._read_json(self.modules_file, {})
|
||
if module_id not in modules:
|
||
return False
|
||
del modules[module_id]
|
||
self._write_json(self.modules_file, modules)
|
||
return True
|
||
|
||
async def set_scope(self, module_id: str, scopes: List[str]) -> Optional[dict]:
|
||
async with self._lock:
|
||
modules = self._read_json(self.modules_file, {})
|
||
item = modules.get(module_id)
|
||
if not item:
|
||
return None
|
||
item["scopes"] = scopes
|
||
item["updated_at"] = _now_iso()
|
||
modules[module_id] = item
|
||
self._write_json(self.modules_file, modules)
|
||
return item
|
||
|
||
async def set_enabled(self, module_id: str, enabled: bool) -> Optional[dict]:
|
||
async with self._lock:
|
||
modules = self._read_json(self.modules_file, {})
|
||
item = modules.get(module_id)
|
||
if not item:
|
||
return None
|
||
item["enabled"] = enabled
|
||
item["updated_at"] = _now_iso()
|
||
modules[module_id] = item
|
||
self._write_json(self.modules_file, modules)
|
||
return item
|
||
|
||
async def list_scripts(self) -> List[dict]:
|
||
out: List[dict] = []
|
||
for p in sorted(self.scripts_dir.glob("*.js"), key=lambda x: x.stat().st_mtime, reverse=True):
|
||
raw = p.read_bytes()
|
||
out.append(
|
||
{
|
||
"script_id": p.stem,
|
||
"filename": p.name,
|
||
"size": len(raw),
|
||
"hash": f"sha256:{hashlib.sha256(raw).hexdigest()}",
|
||
"url": f"/api/v3/scripts/{p.stem}",
|
||
"updated_at": datetime.fromtimestamp(p.stat().st_mtime, tz=timezone.utc).isoformat(),
|
||
}
|
||
)
|
||
return out
|
||
|
||
async def save_script(self, script_id: str, content: bytes) -> dict:
|
||
target = self.scripts_dir / f"{script_id}.js"
|
||
target.write_bytes(content)
|
||
digest = hashlib.sha256(content).hexdigest()
|
||
return {
|
||
"script_id": script_id,
|
||
"url": f"/api/v3/scripts/{script_id}",
|
||
"hash": f"sha256:{digest}",
|
||
"size": len(content),
|
||
}
|
||
|
||
async def get_script_path(self, script_id: str) -> Optional[Path]:
|
||
target = self.scripts_dir / f"{script_id}.js"
|
||
return target if target.exists() else None
|
||
|
||
async def get_device_modules(self, device_id: str) -> dict:
|
||
async with self._lock:
|
||
state = self._read_json(self.device_state_file, {})
|
||
item = state.get(device_id, {})
|
||
modules = item.get("modules", [])
|
||
return {
|
||
"device_id": device_id,
|
||
"supports_hook": item.get("supports_hook", False),
|
||
"frida_version": item.get("frida_version", ""),
|
||
"hook_framework": item.get("hook_framework", "frida-server"),
|
||
"root_status": item.get("root_status", False),
|
||
"modules": modules,
|
||
}
|
||
|
||
async def update_device_probe(self, device_id: str, probe: dict) -> None:
|
||
async with self._lock:
|
||
state = self._read_json(self.device_state_file, {})
|
||
item = state.get(device_id, {})
|
||
item.update(
|
||
{
|
||
"supports_hook": probe.get("supports_hook", False),
|
||
"frida_version": probe.get("frida_version", ""),
|
||
"root_status": probe.get("root_status", False),
|
||
"hook_framework": "frida-server",
|
||
"updated_at": _now_iso(),
|
||
}
|
||
)
|
||
item.setdefault("modules", [])
|
||
state[device_id] = item
|
||
self._write_json(self.device_state_file, state)
|
||
|
||
async def deploy_script(self, script_id: str, device_ids: List[str], auto_reload: bool = True) -> dict:
|
||
async with self._lock:
|
||
state = self._read_json(self.device_state_file, {})
|
||
deployed: List[str] = []
|
||
failed: List[dict] = []
|
||
reloaded: List[str] = []
|
||
now = _now_iso()
|
||
for did in device_ids:
|
||
item = state.setdefault(
|
||
did,
|
||
{"supports_hook": False, "frida_version": "", "root_status": False, "modules": []},
|
||
)
|
||
modules = item.setdefault("modules", [])
|
||
exists = None
|
||
for m in modules:
|
||
if m.get("module_id") == script_id:
|
||
exists = m
|
||
break
|
||
if exists:
|
||
exists["status"] = "loaded"
|
||
exists["loaded_at"] = now
|
||
else:
|
||
modules.append(
|
||
{
|
||
"module_id": script_id,
|
||
"version": "1.0.0",
|
||
"status": "loaded",
|
||
"loaded_at": now,
|
||
"target_process": "com.tencent.mm",
|
||
"target_pid": 0,
|
||
"rpc_methods": [],
|
||
"last_error": None,
|
||
"events_today": 0,
|
||
}
|
||
)
|
||
deployed.append(did)
|
||
if auto_reload:
|
||
reloaded.append(did)
|
||
self._write_json(self.device_state_file, state)
|
||
return {"deployed": deployed, "failed": failed, "reloaded": reloaded}
|
||
|
||
async def reload_device_modules(self, device_id: str, module_ids: List[str], force: bool = False) -> dict:
|
||
async with self._lock:
|
||
state = self._read_json(self.device_state_file, {})
|
||
item = state.get(device_id)
|
||
if not item:
|
||
return {"reloaded": [], "failed": module_ids}
|
||
mods = item.get("modules", [])
|
||
found, failed = [], []
|
||
now = _now_iso()
|
||
for mid in module_ids:
|
||
hit = None
|
||
for m in mods:
|
||
if m.get("module_id") == mid:
|
||
hit = m
|
||
break
|
||
if hit:
|
||
hit["status"] = "loaded"
|
||
hit["loaded_at"] = now
|
||
if force:
|
||
hit["last_error"] = None
|
||
found.append(mid)
|
||
else:
|
||
failed.append(mid)
|
||
item["modules"] = mods
|
||
state[device_id] = item
|
||
self._write_json(self.device_state_file, state)
|
||
return {"reloaded": found, "failed": failed}
|
||
|
||
async def add_device_log(self, device_id: str, module_id: str, line: str) -> None:
|
||
async with self._lock:
|
||
state = self._read_json(self.device_state_file, {})
|
||
item = state.setdefault(device_id, {"modules": []})
|
||
logs = item.setdefault("logs", {})
|
||
bucket = logs.setdefault(module_id, [])
|
||
bucket.append({"ts": _now_iso(), "line": line})
|
||
if len(bucket) > 300:
|
||
logs[module_id] = bucket[-300:]
|
||
state[device_id] = item
|
||
self._write_json(self.device_state_file, state)
|
||
|
||
async def get_device_logs(self, device_id: str, module_id: str) -> List[dict]:
|
||
async with self._lock:
|
||
state = self._read_json(self.device_state_file, {})
|
||
item = state.get(device_id, {})
|
||
logs = item.get("logs", {}).get(module_id, [])
|
||
return logs[-200:]
|
||
|
||
async def add_event(self, event: dict) -> dict:
|
||
data = {"event_id": f"evt_{int(datetime.now().timestamp()*1000)}", **event}
|
||
if "timestamp" not in data:
|
||
data["timestamp"] = _now_iso()
|
||
async with self._lock:
|
||
with self.events_file.open("a", encoding="utf-8") as f:
|
||
f.write(json.dumps(data, ensure_ascii=False) + "\n")
|
||
await self._broadcast_event({"type": "hook_event", "data": data})
|
||
return data
|
||
|
||
async def list_events(
|
||
self,
|
||
event_type: Optional[str] = None,
|
||
device_id: Optional[str] = None,
|
||
platform: Optional[str] = None,
|
||
limit: int = 100,
|
||
) -> List[dict]:
|
||
events: List[dict] = []
|
||
if not self.events_file.exists():
|
||
return events
|
||
lines = self.events_file.read_text(encoding="utf-8").splitlines()[-2000:]
|
||
for line in reversed(lines):
|
||
if len(events) >= limit:
|
||
break
|
||
try:
|
||
item = json.loads(line)
|
||
except Exception:
|
||
continue
|
||
if event_type and item.get("event_type") != event_type:
|
||
continue
|
||
if device_id and item.get("device_id") != device_id:
|
||
continue
|
||
if platform and item.get("platform") != platform:
|
||
continue
|
||
events.append(item)
|
||
return events
|
||
|
||
async def attach_ws_client(self, ws: Any) -> None:
|
||
self._event_clients.add(ws)
|
||
|
||
async def detach_ws_client(self, ws: Any) -> None:
|
||
self._event_clients.discard(ws)
|
||
|
||
async def _broadcast_event(self, payload: dict) -> None:
|
||
if not self._event_clients:
|
||
return
|
||
dead = []
|
||
for ws in list(self._event_clients):
|
||
try:
|
||
await ws.send_json(payload)
|
||
except Exception:
|
||
dead.append(ws)
|
||
for ws in dead:
|
||
self._event_clients.discard(ws)
|
||
|
||
|
||
hook_module_service = HookModuleService()
|