#!/usr/bin/env python3
"""
工作手机 SDK v7 - HTTP REST API 服务
基于 Frida + 微信8.0.56 开发
提供50+个HTTP接口，供存客宝等系统调用

启动方式：
    PHONE_IP=192.168.110.80 python3 workphone_server_v7.py

接口地址：http://localhost:8899/api/v1/...
认证方式：Header: X-API-Key: workphone-sdk-2026
"""

import frida
import json
import time
import threading
import os
import subprocess
from flask import Flask, request, jsonify
from functools import wraps

app = Flask(__name__)

# ==================== 配置 ====================
PHONE_IP = os.environ.get('PHONE_IP', '192.168.110.80')
PHONE_PORT = int(os.environ.get('PHONE_PORT', '5555'))
WECHAT_PKG = 'com.tencent.mm'
API_KEY = os.environ.get('API_KEY', 'workphone-sdk-2026')
SERVER_PORT = int(os.environ.get('SERVER_PORT', '8899'))
JS_PATH = os.path.join(os.path.dirname(__file__), 'sdk/agent/hook/wechat_api_v7.js')

# ==================== Frida管理 ====================
_session = None
_script = None
_lock = threading.Lock()

def get_js_code():
    """读取JS代码"""
    if os.path.exists(JS_PATH):
        with open(JS_PATH, 'r') as f:
            return f.read()
    # 内嵌备用代码
    return open(os.path.join(os.path.dirname(__file__), 'wechat_api_v7.js')).read()

def connect_frida():
    """连接Frida"""
    global _session, _script
    with _lock:
        try:
            if _session and _script:
                # 测试连接是否有效
                try:
                    _script.exports.ping()
                    return True
                except:
                    pass
            
            # 重新连接
            device = frida.get_device_manager().add_remote_device(f'{PHONE_IP}:{PHONE_PORT}')
            _session = device.attach(WECHAT_PKG)
            js_code = get_js_code()
            _script = _session.create_script(js_code)
            _script.load()
            return True
        except Exception as e:
            print(f'[Frida] Connect failed: {e}')
            _session = None
            _script = None
            return False

def call_rpc(method, *args):
    """调用Frida RPC"""
    if not connect_frida():
        return {'error': 'Frida connection failed'}
    try:
        fn = getattr(_script.exports, method)
        result = fn(*args)
        return result
    except Exception as e:
        return {'error': str(e)}

# ==================== 认证中间件 ====================
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({'error': 'Unauthorized', 'code': 401}), 401
        return f(*args, **kwargs)
    return decorated

def api_response(data, status=200):
    return jsonify({'code': 0, 'data': data, 'timestamp': int(time.time())}), status

def api_error(msg, status=400):
    return jsonify({'code': -1, 'error': msg, 'timestamp': int(time.time())}), status

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

# --- 系统 ---
@app.route('/api/v1/ping', methods=['GET'])
def ping():
    return api_response({'status': 'ok', 'version': 'v7', 'phone': PHONE_IP})

@app.route('/api/v1/system/info', methods=['GET'])
@require_api_key
def system_info():
    return api_response(call_rpc('getSystemInfo'))

@app.route('/api/v1/system/db', methods=['GET'])
@require_api_key
def db_info():
    return api_response(call_rpc('getDbInfo'))

@app.route('/api/v1/system/screenshot', methods=['POST'])
@require_api_key
def screenshot():
    data = request.json or {}
    result = call_rpc('takeScreenshot', data.get('path', '/sdcard/wpsdk_screenshot.png'))
    if result.get('success'):
        # 通过ADB拉取截图
        path = result['path']
        local_path = f'/tmp/screenshot_{int(time.time())}.png'
        subprocess.run(['adb', '-s', f'{PHONE_IP}:{PHONE_PORT}', 'pull', path, local_path])
        result['local_path'] = local_path
    return api_response(result)

# --- 账号 ---
@app.route('/api/v1/account', methods=['GET'])
@require_api_key
def get_account():
    return api_response(call_rpc('getAccountInfo'))

# --- 联系人 ---
@app.route('/api/v1/contacts', methods=['GET'])
@require_api_key
def get_contacts():
    limit = int(request.args.get('limit', 100))
    offset = int(request.args.get('offset', 0))
    return api_response(call_rpc('getContacts', limit, offset))

@app.route('/api/v1/contacts/count', methods=['GET'])
@require_api_key
def get_contact_count():
    return api_response(call_rpc('getContactCount'))

@app.route('/api/v1/contacts/search', methods=['GET'])
@require_api_key
def search_contacts():
    keyword = request.args.get('q', '')
    if not keyword:
        return api_error('Missing query parameter: q')
    return api_response(call_rpc('searchContact', keyword))

@app.route('/api/v1/contacts/<wxid>', methods=['GET'])
@require_api_key
def get_contact(wxid):
    return api_response(call_rpc('getContactDetail', wxid))

@app.route('/api/v1/contacts/stats', methods=['GET'])
@require_api_key
def contact_stats():
    return api_response(call_rpc('getContactStats'))

# --- 消息 ---
@app.route('/api/v1/messages', methods=['GET'])
@require_api_key
def get_messages():
    talker = request.args.get('talker', 'filehelper')
    limit = int(request.args.get('limit', 50))
    return api_response(call_rpc('getMessages', talker, limit))

@app.route('/api/v1/messages/recent', methods=['GET'])
@require_api_key
def get_recent_messages():
    hours = int(request.args.get('hours', 24))
    return api_response(call_rpc('getRecentMessages', hours))

@app.route('/api/v1/messages/send', methods=['POST'])
@require_api_key
def send_message():
    data = request.json or {}
    to_user = data.get('to_user')
    content = data.get('content')
    if not to_user or not content:
        return api_error('Missing required fields: to_user, content')
    return api_response(call_rpc('sendMessage', to_user, content))

@app.route('/api/v1/messages/mass', methods=['POST'])
@require_api_key
def mass_message():
    data = request.json or {}
    targets = data.get('targets', [])
    content = data.get('content')
    if not targets or not content:
        return api_error('Missing required fields: targets, content')
    return api_response(call_rpc('massMessage', targets, content))

# --- 会话 ---
@app.route('/api/v1/conversations', methods=['GET'])
@require_api_key
def get_conversations():
    limit = int(request.args.get('limit', 30))
    return api_response(call_rpc('getConversations', limit))

# --- 群组 ---
@app.route('/api/v1/groups', methods=['GET'])
@require_api_key
def get_groups():
    limit = int(request.args.get('limit', 50))
    return api_response(call_rpc('getGroups', limit))

@app.route('/api/v1/groups/count', methods=['GET'])
@require_api_key
def get_group_count():
    return api_response(call_rpc('getGroupCount'))

@app.route('/api/v1/groups/<group_id>', methods=['GET'])
@require_api_key
def get_group(group_id):
    return api_response(call_rpc('getGroupDetail', group_id))

@app.route('/api/v1/groups/<group_id>/members', methods=['GET'])
@require_api_key
def get_group_members(group_id):
    return api_response(call_rpc('getGroupMembers', group_id))

@app.route('/api/v1/groups/<group_id>/announcement', methods=['GET'])
@require_api_key
def get_group_announcement(group_id):
    return api_response(call_rpc('getGroupAnnouncement', group_id))

@app.route('/api/v1/groups/<group_id>/remark', methods=['PUT'])
@require_api_key
def set_group_remark(group_id):
    data = request.json or {}
    remark = data.get('remark', '')
    return api_response(call_rpc('setGroupRemark', group_id, remark))

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

@app.route('/api/v1/labels/<label_id>/contacts', methods=['GET'])
@require_api_key
def get_contacts_by_label(label_id):
    return api_response(call_rpc('getContactsByLabel', label_id))

# --- 朋友圈 ---
@app.route('/api/v1/moments', methods=['GET'])
@require_api_key
def get_moments():
    limit = int(request.args.get('limit', 20))
    return api_response(call_rpc('getMoments', limit))

@app.route('/api/v1/moments/count', methods=['GET'])
@require_api_key
def get_moment_count():
    return api_response(call_rpc('getMomentCount'))

@app.route('/api/v1/moments', methods=['POST'])
@require_api_key
def post_moment():
    data = request.json or {}
    content = data.get('content')
    if not content:
        return api_error('Missing required field: content')
    return api_response(call_rpc('postMoment', content))

# --- 视频号 ---
@app.route('/api/v1/finder/info', methods=['GET'])
@require_api_key
def get_finder_info():
    return api_response(call_rpc('getFinderInfo'))

@app.route('/api/v1/finder/contacts', methods=['GET'])
@require_api_key
def get_finder_contacts():
    return api_response(call_rpc('getFinderContacts'))

@app.route('/api/v1/finder/media', methods=['GET'])
@require_api_key
def get_finder_media():
    return api_response(call_rpc('getFinderMedia'))

# --- 收藏 ---
@app.route('/api/v1/favorites', methods=['GET'])
@require_api_key
def get_favorites():
    limit = int(request.args.get('limit', 20))
    return api_response(call_rpc('getFavorites', limit))

# --- 钱包/红包 ---
@app.route('/api/v1/wallet/info', methods=['GET'])
@require_api_key
def get_wallet_info():
    return api_response(call_rpc('getWalletInfo'))

@app.route('/api/v1/wallet/ledger', methods=['GET'])
@require_api_key
def get_wallet_ledger():
    limit = int(request.args.get('limit', 20))
    return api_response(call_rpc('getWalletLedger', limit))

@app.route('/api/v1/wallet/lucky-money', methods=['GET'])
@require_api_key
def get_lucky_money():
    limit = int(request.args.get('limit', 20))
    return api_response(call_rpc('getLuckyMoney', limit))

@app.route('/api/v1/wallet/balance', methods=['GET'])
@require_api_key
def get_wallet_balance():
    return api_response(call_rpc('getWalletBalance'))

# --- 图片/文件/视频 ---
@app.route('/api/v1/media/images', methods=['GET'])
@require_api_key
def get_images():
    limit = int(request.args.get('limit', 20))
    return api_response(call_rpc('getImageList', limit))

@app.route('/api/v1/media/images/count', methods=['GET'])
@require_api_key
def get_image_count():
    return api_response(call_rpc('getImageCount'))

@app.route('/api/v1/media/videos', methods=['GET'])
@require_api_key
def get_videos():
    limit = int(request.args.get('limit', 20))
    return api_response(call_rpc('getVideoList', limit))

@app.route('/api/v1/media/videos/count', methods=['GET'])
@require_api_key
def get_video_count():
    return api_response(call_rpc('getVideoCount'))

@app.route('/api/v1/media/files', methods=['GET'])
@require_api_key
def get_files():
    limit = int(request.args.get('limit', 20))
    return api_response(call_rpc('getFileList', limit))

# --- 小程序 ---
@app.route('/api/v1/appbrand', methods=['GET'])
@require_api_key
def get_appbrand():
    return api_response(call_rpc('getAppBrandList'))

# --- 搜索 ---
@app.route('/api/v1/search', methods=['GET'])
@require_api_key
def global_search():
    keyword = request.args.get('q', '')
    if not keyword:
        return api_error('Missing query parameter: q')
    return api_response(call_rpc('globalSearch', keyword))

# --- 调试 ---
@app.route('/api/v1/debug/sql', methods=['POST'])
@require_api_key
def raw_sql():
    data = request.json or {}
    db_name = data.get('db', 'main')
    sql = data.get('sql', '')
    if not sql:
        return api_error('Missing required field: sql')
    return api_response(call_rpc('rawQuery', db_name, sql))

# ==================== 健康检查 ====================
@app.route('/health', methods=['GET'])
def health():
    return jsonify({'status': 'ok', 'version': 'v7'})

@app.route('/api/v1/endpoints', methods=['GET'])
def list_endpoints():
    """列出所有API端点"""
    endpoints = []
    for rule in app.url_map.iter_rules():
        if rule.endpoint != 'static':
            endpoints.append({
                'path': str(rule),
                'methods': list(rule.methods - {'HEAD', 'OPTIONS'})
            })
    return jsonify({'endpoints': sorted(endpoints, key=lambda x: x['path']), 'count': len(endpoints)})

if __name__ == '__main__':
    print(f'[*] WorkPhone SDK v7 starting on port {SERVER_PORT}')
    print(f'[*] Phone: {PHONE_IP}:{PHONE_PORT}')
    print(f'[*] API Key: {API_KEY}')
    print(f'[*] Connecting to Frida...')
    if connect_frida():
        print('[+] Frida connected successfully!')
    else:
        print('[-] Frida connection failed, will retry on first request')
    app.run(host='0.0.0.0', port=SERVER_PORT, debug=False, threaded=True)