workphone-sdk/sdk/app/routers/hook_modules.py

"""
Hook 模块管理路由
对标文档：开发文档/5、接口/03-Hook与微信/Hook模块管理接口.md
"""

from __future__ import annotations

from fastapi import APIRouter, File, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field
from typing import List, Optional

from services.adb_device import adb_manager
from services.hook_module_service import hook_module_service
from services.ws_hub import ws_hub

router = APIRouter()


class ModuleUpsertRequest(BaseModel):
    module_id: str
    name: str
    version: str = "1.0.0"
    description: str = ""
    scopes: List[str] = Field(default_factory=lambda: ["com.tencent.mm"])
    capabilities: List[str] = Field(default_factory=list)
    min_frida_version: str = "16.0.0"
    script_content: Optional[str] = None
    enabled: bool = True


class ScopeUpdateRequest(BaseModel):
    scopes: List[str]


class DeployScriptRequest(BaseModel):
    device_ids: List[str]
    auto_reload: bool = True


class ReloadDeviceModulesRequest(BaseModel):
    module_ids: List[str]
    force: bool = False


def _probe_hook_ready(device_id: str) -> dict:
    dev = adb_manager.get_device(device_id)
    if not dev or not dev.is_online():
        return {
            "supports_hook": False,
            "root_status": False,
            "frida_version": "",
            "detail": "device offline",
        }
    root_status = False
    frida_version = ""
    detail = "ok"
    try:
        uid = dev._shell("id -u", timeout=4)
        su_path = dev._shell("which su", timeout=4)
        root_status = (uid.strip() == "0") or bool(su_path.strip())
    except Exception:
        pass
    try:
        frida_version = dev._shell("frida-server --version", timeout=4).strip()
    except Exception:
        frida_version = ""
    # 常见部署把 frida-server 放在 /data/local/tmp，不在默认 PATH
    if not frida_version:
        try:
            frida_version = dev._shell("/data/local/tmp/frida-server --version", timeout=4).strip()
        except Exception:
            frida_version = ""
    supports_hook = bool(root_status and frida_version)
    if not root_status:
        detail = "root unavailable"
    elif not frida_version:
        detail = "frida-server unavailable"
    return {
        "supports_hook": supports_hook,
        "root_status": root_status,
        "frida_version": frida_version,
        "detail": detail,
    }


@router.get("/modules")
async def list_modules(enabled: Optional[bool] = None, scope: Optional[str] = None):
    modules = await hook_module_service.list_modules(enabled=enabled, scope=scope)
    return {"code": 200, "data": {"total": len(modules), "modules": modules}}


@router.post("/modules")
async def create_module(req: ModuleUpsertRequest):
    data = await hook_module_service.upsert_module(req.model_dump())
    return {"code": 200, "data": data}


@router.get("/modules/{module_id}")
async def get_module(module_id: str):
    item = await hook_module_service.get_module(module_id)
    if not item:
        raise HTTPException(status_code=404, detail="模块不存在")
    return {"code": 200, "data": item}


@router.put("/modules/{module_id}")
async def update_module(module_id: str, req: ModuleUpsertRequest):
    payload = req.model_dump()
    payload["module_id"] = module_id
    data = await hook_module_service.upsert_module(payload)
    return {"code": 200, "data": data}


@router.delete("/modules/{module_id}")
async def delete_module(module_id: str):
    ok = await hook_module_service.delete_module(module_id)
    if not ok:
        raise HTTPException(status_code=404, detail="模块不存在")
    return {"code": 200, "data": {"module_id": module_id, "deleted": True}}


@router.put("/modules/{module_id}/scope")
async def set_module_scope(module_id: str, req: ScopeUpdateRequest):
    item = await hook_module_service.set_scope(module_id, req.scopes)
    if not item:
        raise HTTPException(status_code=404, detail="模块不存在")
    return {"code": 200, "data": {"module_id": module_id, "scopes": item.get("scopes", [])}}


@router.post("/modules/{module_id}/enable")
async def enable_module(module_id: str):
    item = await hook_module_service.set_enabled(module_id, True)
    if not item:
        raise HTTPException(status_code=404, detail="模块不存在")
    return {"code": 200, "data": {"module_id": module_id, "enabled": True, "affected_devices": item.get("device_count", 0)}}


@router.post("/modules/{module_id}/disable")
async def disable_module(module_id: str):
    item = await hook_module_service.set_enabled(module_id, False)
    if not item:
        raise HTTPException(status_code=404, detail="模块不存在")
    return {"code": 200, "data": {"module_id": module_id, "enabled": False}}


@router.get("/devices/{device_id}/modules")
async def get_device_modules(device_id: str):
    probe = _probe_hook_ready(device_id)
    await hook_module_service.update_device_probe(device_id, probe)
    data = await hook_module_service.get_device_modules(device_id)
    data["supports_hook"] = probe["supports_hook"]
    data["frida_version"] = probe["frida_version"]
    data["root_status"] = probe["root_status"]
    data["probe_detail"] = probe["detail"]
    return {"code": 200, "data": data}


@router.post("/devices/{device_id}/modules/reload")
async def reload_device_modules(device_id: str, req: ReloadDeviceModulesRequest):
    result = await hook_module_service.reload_device_modules(device_id, req.module_ids, force=req.force)
    if ws_hub.is_online(device_id):
        result["agent_reload"] = await ws_hub.send_command(
            device_id,
            {
                "type": "frida_reload",
                "data": {"module_ids": req.module_ids, "force": req.force},
            },
            timeout=30,
        )
    else:
        result["agent_reload"] = {"code": 503, "message": "设备 WebSocket 不在线"}
    return {"code": 200, "data": result}


@router.get("/devices/{device_id}/modules/{module_id}/logs")
async def get_device_module_logs(device_id: str, module_id: str):
    logs = await hook_module_service.get_device_logs(device_id, module_id)
    return {"code": 200, "data": {"device_id": device_id, "module_id": module_id, "logs": logs}}


@router.get("/scripts")
async def list_scripts():
    scripts = await hook_module_service.list_scripts()
    return {"code": 200, "data": {"total": len(scripts), "scripts": scripts}}


@router.post("/scripts")
async def upload_script(
    file: UploadFile = File(...),
    module_id: str = "wechat_hook_v1",
    version: str = "1.0.0",
    description: str = "",
):
    raw = await file.read()
    script_id = f"{module_id}_{version}"
    saved = await hook_module_service.save_script(script_id, raw)
    await hook_module_service.upsert_module(
        {
            "module_id": module_id,
            "name": "微信Hook模块" if "wechat" in module_id else module_id,
            "version": version,
            "description": description or "Hook脚本上传",
            "scopes": ["com.tencent.mm"] if "wechat" in module_id else [],
            "capabilities": ["send_message", "get_messages", "get_contacts"] if "wechat" in module_id else [],
            "min_frida_version": "16.0.0",
            "enabled": True,
            "script_id": script_id,
            "script_url": saved["url"],
            "script_hash": saved["hash"],
        }
    )
    return {"code": 200, "data": saved}


@router.get("/scripts/{script_id}")
async def download_script(script_id: str):
    path = await hook_module_service.get_script_path(script_id)
    if not path:
        raise HTTPException(status_code=404, detail="脚本不存在")
    return FileResponse(path, media_type="application/javascript", filename=path.name)


@router.post("/scripts/{script_id}/deploy")
async def deploy_script(script_id: str, req: DeployScriptRequest):
    result = await hook_module_service.deploy_script(script_id, req.device_ids, auto_reload=req.auto_reload)
    for did in result["deployed"]:
        await hook_module_service.add_device_log(did, script_id, f"deploy {script_id} ok")
    return {"code": 200, "data": result}


@router.get("/hook/events")
async def list_hook_events(
    event_type: Optional[str] = None,
    device_id: Optional[str] = None,
    platform: Optional[str] = None,
    limit: int = 100,
):
    events = await hook_module_service.list_events(
        event_type=event_type,
        device_id=device_id,
        platform=platform,
        limit=max(1, min(limit, 500)),
    )
    return {"code": 200, "data": {"total": len(events), "events": events}}


@router.post("/hook/events")
async def ingest_hook_event(payload: dict):
    event = await hook_module_service.add_event(payload)
    return {"code": 200, "data": event}


@router.websocket("/hook/events/stream")
async def hook_event_stream(websocket: WebSocket):
    await websocket.accept()
    await hook_module_service.attach_ws_client(websocket)
    try:
        while True:
            # 订阅接口支持客户端心跳/控制消息，当前不做强校验
            _ = await websocket.receive_text()
    except WebSocketDisconnect:
        await hook_module_service.detach_ws_client(websocket)
    except Exception:
        await hook_module_service.detach_ws_client(websocket)