#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
工作手机SDK - HTTP REST API服务
版本: v5.0
功能: 将所有Frida微信控制接口封装为HTTP API，供存客宝等系统调用

启动: python3 workphone_api_server.py
默认端口: 8899
API文档: http://localhost:8899/docs
"""
import frida, time, json, subprocess, os, threading, logging
from flask import Flask, request, jsonify
from datetime import datetime
from functools import wraps

# ============================================================
# 配置
# ============================================================
PHONE_IP = os.environ.get("PHONE_IP", "192.168.110.80")
FRIDA_PORT = int(os.environ.get("FRIDA_PORT", "27042"))
ADB_SERIAL = os.environ.get("ADB_SERIAL", "192.168.110.80:5555")
API_PORT = int(os.environ.get("API_PORT", "8899"))
API_KEY = os.environ.get("API_KEY", "workphone-sdk-2026")
BASE = os.path.dirname(os.path.abspath(__file__))
HOOK_JS = os.path.join(BASE, "sdk/agent/hook/wechat_full_control_v5.js")

logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')
logger = logging.getLogger("WorkphoneAPI")

app = Flask(__name__)

# ============================================================
# Frida连接管理
# ============================================================
class FridaManager:
    def __init__(self):
        self.device = None
        self.session = None
        self.script = None
        self.rpc = None
        self.connected = False
        self.wechat_pid = None
        self.lock = threading.Lock()
        self.events = []

    def get_wechat_pid(self):
        try:
            result = subprocess.run(
                ["adb", "-s", ADB_SERIAL, "shell", "ps -A | grep 'com.tencent.mm$'"],
                capture_output=True, text=True, timeout=5
            )
            line = result.stdout.strip()
            if line:
                return int(line.split()[1])
        except Exception as e:
            logger.error(f"获取PID失败: {e}")
        return None

    def connect(self):
        with self.lock:
            try:
                self.wechat_pid = self.get_wechat_pid()
                if not self.wechat_pid:
                    return False, "无法获取微信PID"

                dm = frida.get_device_manager()
                self.device = dm.add_remote_device(f"{PHONE_IP}:{FRIDA_PORT}")
                self.session = self.device.attach(self.wechat_pid)

                def on_msg(m, d):
                    if m.get("type") == "send":
                        p = m["payload"]
                        self.events.append({**p, "time": datetime.now().isoformat()})
                        if len(self.events) > 1000:
                            self.events = self.events[-500:]
                        logger.info(f"[EVENT:{p.get('type','')}] {json.dumps(p, ensure_ascii=False)[:100]}")

                with open(HOOK_JS, encoding="utf-8") as f:
                    src = f.read()
                self.script = self.session.create_script(src)
                self.script.on("message", on_msg)
                self.script.load()
                time.sleep(3)
                self.rpc = self.script.exports_sync
                self.connected = True
                logger.info(f"Frida连接成功 PID={self.wechat_pid}")
                return True, f"连接成功 PID={self.wechat_pid}"
            except Exception as e:
                self.connected = False
                logger.error(f"Frida连接失败: {e}")
                return False, str(e)

    def reconnect(self):
        try:
            if self.session:
                try: self.session.detach()
                except: pass
        except: pass
        self.connected = False
        return self.connect()

    def call(self, method, params=None):
        if not self.connected or not self.rpc:
            ok, msg = self.reconnect()
            if not ok:
                return {"success": False, "error": f"Frida未连接: {msg}"}
        try:
            fn = getattr(self.rpc, method)
            result = fn(params) if params is not None else fn()
            return result
        except frida.InvalidOperationError as e:
            logger.warning(f"Frida会话失效，重连中: {e}")
            ok, msg = self.reconnect()
            if ok:
                try:
                    fn = getattr(self.rpc, method)
                    return fn(params) if params is not None else fn()
                except Exception as e2:
                    return {"success": False, "error": str(e2)}
            return {"success": False, "error": f"重连失败: {msg}"}
        except Exception as e:
            return {"success": False, "error": str(e)}

frida_mgr = FridaManager()

# ============================================================
# 中间件
# ============================================================
def require_api_key(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        key = request.headers.get("X-API-Key") or request.args.get("api_key")
        if key != API_KEY:
            return jsonify({"success": False, "error": "Invalid API Key", "code": 401}), 401
        return f(*args, **kwargs)
    return decorated

def api_response(data, status=200):
    data["timestamp"] = datetime.now().isoformat()
    return jsonify(data), status

def get_params():
    if request.method == "GET":
        return dict(request.args)
    try:
        return request.get_json(force=True) or {}
    except:
        return {}

# ============================================================
# 路由
# ============================================================

# --- 系统 ---
@app.route("/api/v1/system/info", methods=["GET"])
@require_api_key
def system_info():
    return api_response(frida_mgr.call("getSystemInfo"))

@app.route("/api/v1/system/connect", methods=["POST"])
@require_api_key
def system_connect():
    ok, msg = frida_mgr.connect() if not frida_mgr.connected else (True, "已连接")
    return api_response({"success": ok, "message": msg, "pid": frida_mgr.wechat_pid})

@app.route("/api/v1/system/reconnect", methods=["POST"])
@require_api_key
def system_reconnect():
    ok, msg = frida_mgr.reconnect()
    return api_response({"success": ok, "message": msg, "pid": frida_mgr.wechat_pid})

@app.route("/api/v1/system/events", methods=["GET"])
@require_api_key
def system_events():
    limit = int(request.args.get("limit", 50))
    return api_response({"success": True, "events": frida_mgr.events[-limit:], "total": len(frida_mgr.events)})

@app.route("/api/v1/system/screenshot", methods=["POST"])
@require_api_key
def system_screenshot():
    params = get_params()
    path = params.get("path", "/sdcard/sdk_screenshot.png")
    r = frida_mgr.call("takeScreenshot", {"path": path})
    if r.get("success"):
        local = f"/tmp/sdk_screenshot_{int(time.time())}.png"
        subprocess.run(["adb", "-s", ADB_SERIAL, "pull", path, local], capture_output=True)
        r["local_path"] = local
    return api_response(r)

# --- 账号 ---
@app.route("/api/v1/account/info", methods=["GET"])
@require_api_key
def account_info():
    return api_response(frida_mgr.call("getAccountInfo"))

# --- 联系人 ---
@app.route("/api/v1/contacts", methods=["GET"])
@require_api_key
def get_contacts():
    params = {
        "limit": int(request.args.get("limit", 200)),
        "offset": int(request.args.get("offset", 0)),
        "keyword": request.args.get("keyword")
    }
    return api_response(frida_mgr.call("getContacts", params))

@app.route("/api/v1/contacts/count", methods=["GET"])
@require_api_key
def contacts_count():
    return api_response(frida_mgr.call("getContactsCount"))

@app.route("/api/v1/contacts/<wxid>", methods=["GET"])
@require_api_key
def contact_detail(wxid):
    return api_response(frida_mgr.call("getContactDetail", {"wxid": wxid}))

@app.route("/api/v1/contacts/<wxid>/remark", methods=["PUT"])
@require_api_key
def set_remark(wxid):
    params = get_params()
    return api_response(frida_mgr.call("setRemark", {"wxid": wxid, "remark": params.get("remark")}))

@app.route("/api/v1/contacts/add", methods=["POST"])
@require_api_key
def add_friend():
    params = get_params()
    return api_response(frida_mgr.call("addFriend", {"wxid": params.get("wxid")}))

@app.route("/api/v1/contacts/search", methods=["GET"])
@require_api_key
def search_contacts():
    return api_response(frida_mgr.call("searchContacts", {"keyword": request.args.get("keyword")}))

# --- 消息 ---
@app.route("/api/v1/messages/send", methods=["POST"])
@require_api_key
def send_message():
    params = get_params()
    return api_response(frida_mgr.call("sendMessage", {
        "to_id": params.get("to_id"),
        "content": params.get("content"),
        "type": params.get("type", 1)
    }))

@app.route("/api/v1/messages/mass", methods=["POST"])
@require_api_key
def mass_message():
    params = get_params()
    return api_response(frida_mgr.call("massMessage", {
        "targets": params.get("targets", []),
        "content": params.get("content")
    }))

@app.route("/api/v1/messages", methods=["GET"])
@require_api_key
def get_messages():
    params = {
        "talker": request.args.get("talker"),
        "limit": int(request.args.get("limit", 50))
    }
    return api_response(frida_mgr.call("getMessages", params))

@app.route("/api/v1/messages/search", methods=["GET"])
@require_api_key
def search_messages():
    return api_response(frida_mgr.call("searchMessages", {"keyword": request.args.get("keyword")}))

@app.route("/api/v1/conversations", methods=["GET"])
@require_api_key
def get_conversations():
    return api_response(frida_mgr.call("getConversations", {"limit": int(request.args.get("limit", 20))}))

# --- 朋友圈 ---
@app.route("/api/v1/moments", methods=["GET"])
@require_api_key
def get_moments():
    params = {
        "limit": int(request.args.get("limit", 20)),
        "offset": int(request.args.get("offset", 0))
    }
    return api_response(frida_mgr.call("getMoments", params))

@app.route("/api/v1/moments/post", methods=["POST"])
@require_api_key
def post_moment():
    params = get_params()
    return api_response(frida_mgr.call("postMoment", {"content": params.get("content")}))

@app.route("/api/v1/moments/<sns_id>", methods=["GET"])
@require_api_key
def moment_detail(sns_id):
    return api_response(frida_mgr.call("getMomentDetail", {"sns_id": sns_id}))

@app.route("/api/v1/moments/<sns_id>/comments", methods=["GET"])
@require_api_key
def moment_comments(sns_id):
    return api_response(frida_mgr.call("getMomentComments", {"sns_id": sns_id}))

# --- 视频号 ---
@app.route("/api/v1/finder/contacts", methods=["GET"])
@require_api_key
def finder_contacts():
    return api_response(frida_mgr.call("getFinderContacts", {"limit": int(request.args.get("limit", 20))}))

@app.route("/api/v1/finder/videos", methods=["GET"])
@require_api_key
def finder_videos():
    return api_response(frida_mgr.call("getFinderVideos", {"limit": int(request.args.get("limit", 20))}))

@app.route("/api/v1/finder/open", methods=["POST"])
@require_api_key
def open_finder():
    return api_response(frida_mgr.call("openFinder"))

# --- 收藏 ---
@app.route("/api/v1/favorites", methods=["GET"])
@require_api_key
def get_favorites():
    return api_response(frida_mgr.call("getFavorites", {"limit": int(request.args.get("limit", 50))}))

# --- 红包/钱包 ---
@app.route("/api/v1/wallet/redpackets", methods=["GET"])
@require_api_key
def get_redpackets():
    return api_response(frida_mgr.call("getRedPackets", {"limit": int(request.args.get("limit", 50))}))

@app.route("/api/v1/wallet/info", methods=["GET"])
@require_api_key
def wallet_info():
    return api_response(frida_mgr.call("getWalletInfo"))

@app.route("/api/v1/wallet/ledger", methods=["GET"])
@require_api_key
def wallet_ledger():
    return api_response(frida_mgr.call("getWalletLedger", {"limit": int(request.args.get("limit", 100))}))

# --- 文件/媒体 ---
@app.route("/api/v1/media/files", methods=["GET"])
@require_api_key
def get_files():
    return api_response(frida_mgr.call("getFiles", {"limit": int(request.args.get("limit", 50))}))

@app.route("/api/v1/media/images", methods=["GET"])
@require_api_key
def get_images():
    return api_response(frida_mgr.call("getImages", {"limit": int(request.args.get("limit", 50))}))

@app.route("/api/v1/media/videos", methods=["GET"])
@require_api_key
def get_videos():
    return api_response(frida_mgr.call("getVideos", {"limit": int(request.args.get("limit", 50))}))

@app.route("/api/v1/media/voices", methods=["GET"])
@require_api_key
def get_voices():
    return api_response(frida_mgr.call("getVoices", {"limit": int(request.args.get("limit", 50))}))

# --- 群组 ---
@app.route("/api/v1/groups", methods=["GET"])
@require_api_key
def get_groups():
    return api_response(frida_mgr.call("getGroups", {"limit": int(request.args.get("limit", 100))}))

@app.route("/api/v1/groups/<group_id>", methods=["GET"])
@require_api_key
def group_detail(group_id):
    return api_response(frida_mgr.call("getGroupDetail", {"group_id": group_id}))

@app.route("/api/v1/groups/<group_id>/members", methods=["GET"])
@require_api_key
def group_members(group_id):
    return api_response(frida_mgr.call("getGroupMembers", {"group_id": group_id}))

@app.route("/api/v1/groups/<group_id>/notice", methods=["PUT"])
@require_api_key
def set_group_notice(group_id):
    params = get_params()
    return api_response(frida_mgr.call("setGroupNotice", {"group_id": group_id, "notice": params.get("notice")}))

@app.route("/api/v1/groups/<group_id>/name", methods=["PUT"])
@require_api_key
def set_group_name(group_id):
    params = get_params()
    return api_response(frida_mgr.call("setGroupName", {"group_id": group_id, "name": params.get("name")}))

@app.route("/api/v1/groups/<group_id>/nickname", methods=["PUT"])
@require_api_key
def set_group_nickname(group_id):
    params = get_params()
    return api_response(frida_mgr.call("setGroupNickname", {"group_id": group_id, "nickname": params.get("nickname")}))

@app.route("/api/v1/groups/<group_id>/kick", methods=["POST"])
@require_api_key
def kick_member(group_id):
    params = get_params()
    return api_response(frida_mgr.call("kickGroupMember", {"group_id": group_id, "member_id": params.get("member_id")}))

@app.route("/api/v1/groups/<group_id>/message", methods=["POST"])
@require_api_key
def send_group_message(group_id):
    params = get_params()
    return api_response(frida_mgr.call("sendMessage", {"to_id": group_id, "content": params.get("content")}))

# --- 标签 ---
@app.route("/api/v1/labels", methods=["GET"])
@require_api_key
def get_labels():
    return api_response(frida_mgr.call("getLabels"))

@app.route("/api/v1/labels/<label_id>/members", methods=["GET"])
@require_api_key
def label_members(label_id):
    return api_response(frida_mgr.call("getLabelMembers", {"label_id": label_id}))

# --- 小程序 ---
@app.route("/api/v1/miniapps", methods=["GET"])
@require_api_key
def get_miniapps():
    return api_response(frida_mgr.call("getMiniApps", {"limit": int(request.args.get("limit", 50))}))

@app.route("/api/v1/miniapps/open", methods=["POST"])
@require_api_key
def open_miniapp():
    params = get_params()
    return api_response(frida_mgr.call("openMiniApp", {"app_id": params.get("app_id")}))

# --- 导航 ---
@app.route("/api/v1/navigate/chat", methods=["POST"])
@require_api_key
def navigate_chat():
    params = get_params()
    return api_response(frida_mgr.call("navigateToChat", {"wxid": params.get("wxid")}))

@app.route("/api/v1/navigate/moments", methods=["POST"])
@require_api_key
def navigate_moments():
    return api_response(frida_mgr.call("navigateToMoments"))

@app.route("/api/v1/navigate/contacts", methods=["POST"])
@require_api_key
def navigate_contacts():
    return api_response(frida_mgr.call("navigateToContacts"))

# --- 高级SQL ---
@app.route("/api/v1/db/query", methods=["POST"])
@require_api_key
def db_query():
    params = get_params()
    return api_response(frida_mgr.call("rawQuery", {"sql": params.get("sql"), "db": params.get("db")}))

@app.route("/api/v1/db/info", methods=["GET"])
@require_api_key
def db_info():
    return api_response(frida_mgr.call("getDbInfo"))

@app.route("/api/v1/db/tables", methods=["GET"])
@require_api_key
def db_tables():
    return api_response(frida_mgr.call("listAllTables"))

# --- 健康检查 ---
@app.route("/health", methods=["GET"])
def health():
    return jsonify({
        "status": "ok",
        "connected": frida_mgr.connected,
        "pid": frida_mgr.wechat_pid,
        "timestamp": datetime.now().isoformat()
    })

@app.route("/docs", methods=["GET"])
def docs():
    routes = []
    for rule in app.url_map.iter_rules():
        if rule.endpoint not in ("static", "health", "docs"):
            routes.append({
                "path": str(rule),
                "methods": list(rule.methods - {"HEAD", "OPTIONS"}),
                "endpoint": rule.endpoint
            })
    routes.sort(key=lambda x: x["path"])
    return jsonify({
        "title": "工作手机SDK API文档",
        "version": "v5.0",
        "base_url": f"http://{{host}}:{API_PORT}",
        "auth": "Header: X-API-Key: workphone-sdk-2026",
        "total_endpoints": len(routes),
        "endpoints": routes
    })

# ============================================================
# 启动
# ============================================================
if __name__ == "__main__":
    logger.info("=== 工作手机SDK API服务 v5.0 ===")
    logger.info(f"手机IP: {PHONE_IP}:{FRIDA_PORT}")
    logger.info("正在连接Frida...")
    ok, msg = frida_mgr.connect()
    logger.info(f"Frida连接: {'✅ '+msg if ok else '❌ '+msg}")
    logger.info(f"API服务启动: http://0.0.0.0:{API_PORT}")
    logger.info(f"API文档: http://localhost:{API_PORT}/docs")
    logger.info(f"健康检查: http://localhost:{API_PORT}/health")
    app.run(host="0.0.0.0", port=API_PORT, debug=False, threaded=True)