#!/usr/bin/env python3
"""
工作手机 SDK - 完整HTTP REST API服务
版本: v9-complete
功能: 微信全量控制接口，基于Frida无感注入
"""

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

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
logger = logging.getLogger(__name__)

app = Flask(__name__)

# 配置
API_KEY = os.environ.get('API_KEY', 'workphone-sdk-2026')
PHONE_IP = os.environ.get('PHONE_IP', '192.168.110.80')
PHONE_PORT = int(os.environ.get('PHONE_PORT', 27042))
JS_FILE = os.path.join(os.path.dirname(__file__), 'sdk/agent/hook/wechat_api_v8_fix.js')

# 全局状态
class SDKState:
    device = None
    session = None
    script = None
    connected = False
    wechat_pid = None
    lock = threading.Lock()
    last_heartbeat = 0

state = SDKState()

def get_js_code():
    """读取JS脚本"""
    js_paths = [
        JS_FILE,
        os.path.join(os.path.dirname(__file__), 'wechat_api_v8_fix.js'),
        '/Users/karuo/Documents/开发/2、私域银行/工作手机/sdk/agent/hook/wechat_api_v8_fix.js',
    ]
    for p in js_paths:
        if os.path.exists(p):
            with open(p, 'r', encoding='utf-8') as f:
                return f.read()
    raise FileNotFoundError(f"JS文件未找到，已尝试: {js_paths}")

def get_wechat_pid():
    """获取微信主进程PID"""
    try:
        result = subprocess.run(
            ['adb', '-s', f'{PHONE_IP}:5555', 'shell', 'ps -A | grep com.tencent.mm'],
            capture_output=True, text=True, timeout=5
        )
        for line in result.stdout.split('\n'):
            if 'com.tencent.mm' in line and ':push' not in line and ':sandbox' not in line and ':tools' not in line:
                parts = line.split()
                if len(parts) >= 2:
                    return int(parts[1])
    except Exception as e:
        logger.error(f"获取微信PID失败: {e}")
    return None

def ensure_frida_forward():
    """确保ADB端口转发"""
    try:
        subprocess.run(
            ['adb', '-s', f'{PHONE_IP}:5555', 'forward', 'tcp:27042', 'tcp:27042'],
            capture_output=True, timeout=5
        )
    except Exception as e:
        logger.warning(f"端口转发失败: {e}")

def connect_frida():
    """连接Frida并注入脚本"""
    with state.lock:
        try:
            ensure_frida_forward()
            
            # 获取设备
            dm = frida.get_device_manager()
            state.device = dm.add_remote_device('localhost:27042')
            
            # 获取微信PID
            pid = get_wechat_pid()
            if not pid:
                # 尝试启动微信
                subprocess.run(
                    ['adb', '-s', f'{PHONE_IP}:5555', 'shell', 
                     'monkey -p com.tencent.mm -c android.intent.category.LAUNCHER 1'],
                    capture_output=True, timeout=10
                )
                time.sleep(3)
                pid = get_wechat_pid()
            
            if not pid:
                raise Exception("无法获取微信PID")
            
            state.wechat_pid = pid
            logger.info(f"微信PID: {pid}")
            
            # Attach
            state.session = state.device.attach(pid)
            
            # 加载脚本
            js_code = get_js_code()
            state.script = state.session.create_script(js_code)
            
            def on_message(message, data):
                if message.get('type') == 'send':
                    payload = message.get('payload', {})
                    if isinstance(payload, dict):
                        logger.debug(f"Frida: {payload.get('type', '')} - {str(payload.get('data', ''))[:100]}")
                elif message.get('type') == 'error':
                    logger.error(f"Frida错误: {message.get('description', '')}")
            
            state.script.on('message', on_message)
            state.script.load()
            
            # 初始化DB
            time.sleep(1)
            try:
                state.script.exports.initDbs()
            except Exception as e:
                logger.warning(f"initDbs: {e}")
            
            state.connected = True
            state.last_heartbeat = time.time()
            logger.info("✅ Frida连接成功")
            return True
            
        except Exception as e:
            logger.error(f"Frida连接失败: {e}")
            state.connected = False
            return False

def ensure_connected():
    """确保Frida连接有效"""
    if not state.connected or not state.script:
        return connect_frida()
    
    # 心跳检测（每30秒）
    if time.time() - state.last_heartbeat > 30:
        try:
            state.script.exports.ping()
            state.last_heartbeat = time.time()
        except Exception:
            logger.warning("Frida心跳失败，重新连接...")
            state.connected = False
            return connect_frida()
    
    return True

def call_rpc(method, *args):
    """调用Frida RPC方法"""
    if not ensure_connected():
        raise Exception("Frida未连接")
    
    try:
        func = getattr(state.script.exports, method)
        result = func(*args) if args else func()
        state.last_heartbeat = time.time()
        return result
    except Exception as e:
        logger.error(f"RPC调用失败 {method}: {e}")
        # 尝试重连一次
        if 'session' in str(e).lower() or 'detach' in str(e).lower():
            state.connected = False
            if connect_frida():
                func = getattr(state.script.exports, method)
                return func(*args) if args else func()
        raise

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': '无效的API Key'}), 401
        return f(*args, **kwargs)
    return decorated

def success(data=None, **kwargs):
    resp = {'success': True, 'timestamp': int(time.time() * 1000)}
    if data is not None:
        resp['data'] = data
    resp.update(kwargs)
    return jsonify(resp)

def error(msg, code=500):
    return jsonify({'success': False, 'error': str(msg)}), code

# ==================== 系统接口 ====================

@app.route('/api/v1/ping', methods=['GET'])
def ping():
    return jsonify({
        'status': 'ok',
        'version': 'v9-complete',
        'frida_connected': state.connected,
        'wechat_pid': state.wechat_pid,
        'timestamp': int(time.time() * 1000)
    })

@app.route('/api/v1/init', methods=['POST'])
@require_api_key
def init():
    result = connect_frida()
    return success({'connected': result, 'pid': state.wechat_pid})

@app.route('/api/v1/dbs', methods=['GET'])
@require_api_key
def list_dbs():
    try:
        result = call_rpc('getDbInfo')
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 账号接口 ====================

@app.route('/api/v1/account', methods=['GET'])
@require_api_key
def get_account():
    try:
        result = call_rpc('getAccountInfo')
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 联系人接口 ====================

@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))
    try:
        result = call_rpc('getContacts', limit, offset)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/contacts/count', methods=['GET'])
@require_api_key
def get_contacts_count():
    try:
        result = call_rpc('getContactStats')
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/contacts/search', methods=['GET'])
@require_api_key
def search_contacts():
    keyword = request.args.get('keyword', '')
    limit = int(request.args.get('limit', 50))
    try:
        result = call_rpc('searchContacts', keyword, limit)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/contacts/<wxid>', methods=['GET'])
@require_api_key
def get_contact(wxid):
    try:
        result = call_rpc('getContactDetail', wxid)
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 消息接口 ====================

@app.route('/api/v1/messages/send', methods=['POST'])
@require_api_key
def send_message():
    data = request.get_json()
    to_user = data.get('to_user', 'filehelper')
    content = data.get('content', '')
    msg_type = data.get('type', 1)
    
    if not content:
        return error('content不能为空', 400)
    
    try:
        result = call_rpc('sendMessage', to_user, content, msg_type)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/messages/mass', methods=['POST'])
@require_api_key
def mass_send():
    data = request.get_json()
    targets = data.get('targets', [])
    content = data.get('content', '')
    
    if not targets or not content:
        return error('targets和content不能为空', 400)
    
    try:
        result = call_rpc('massMessage', targets, content)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/messages/<talker>', methods=['GET'])
@require_api_key
def get_messages(talker):
    limit = int(request.args.get('limit', 50))
    try:
        result = call_rpc('getMessages', talker, limit)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/messages/recent', methods=['GET'])
@require_api_key
def get_recent_messages():
    hours = int(request.args.get('hours', 24))
    try:
        result = call_rpc('getRecentMessages', hours)
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 会话接口 ====================

@app.route('/api/v1/conversations', methods=['GET'])
@require_api_key
def get_conversations():
    limit = int(request.args.get('limit', 30))
    try:
        result = call_rpc('getConversations', limit)
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 群组接口 ====================

@app.route('/api/v1/groups', methods=['GET'])
@require_api_key
def get_groups():
    limit = int(request.args.get('limit', 50))
    try:
        result = call_rpc('getGroups', limit)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/groups/count', methods=['GET'])
@require_api_key
def get_groups_count():
    try:
        result = call_rpc('getGroupStats')
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/groups/<chatroom_id>/members', methods=['GET'])
@require_api_key
def get_group_members(chatroom_id):
    try:
        result = call_rpc('getGroupMembers', chatroom_id)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/groups/<chatroom_id>/remark', methods=['PUT'])
@require_api_key
def set_group_remark(chatroom_id):
    data = request.get_json()
    remark = data.get('remark', '')
    try:
        result = call_rpc('setGroupRemark', chatroom_id, remark)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/groups/<chatroom_id>/announcement', methods=['PUT'])
@require_api_key
def set_group_announcement(chatroom_id):
    data = request.get_json()
    announcement = data.get('announcement', '')
    try:
        result = call_rpc('setGroupAnnouncement', chatroom_id, announcement)
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 标签接口 ====================

@app.route('/api/v1/labels', methods=['GET'])
@require_api_key
def get_labels():
    try:
        result = call_rpc('getLabels')
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/labels/<label_id>/contacts', methods=['GET'])
@require_api_key
def get_label_contacts(label_id):
    limit = int(request.args.get('limit', 100))
    try:
        result = call_rpc('getContactsByLabel', label_id, limit)
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 朋友圈接口 ====================

@app.route('/api/v1/moments', methods=['GET'])
@require_api_key
def get_moments():
    limit = int(request.args.get('limit', 20))
    try:
        result = call_rpc('getMoments', limit)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/moments', methods=['POST'])
@require_api_key
def post_moment():
    data = request.get_json()
    content = data.get('content', '')
    if not content:
        return error('content不能为空', 400)
    try:
        result = call_rpc('postMoment', content)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/moments/dbinfo', methods=['GET'])
@require_api_key
def moments_dbinfo():
    try:
        result = call_rpc('getMomentsDbInfo')
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 视频号接口 ====================

@app.route('/api/v1/finder/info', methods=['GET'])
@require_api_key
def finder_info():
    try:
        result = call_rpc('getFinderInfo')
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/finder/contacts', methods=['GET'])
@require_api_key
def finder_contacts():
    limit = int(request.args.get('limit', 20))
    try:
        result = call_rpc('getFinderContacts', limit)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/finder/accounts', methods=['GET'])
@require_api_key
def finder_accounts():
    try:
        result = call_rpc('getFinderAccounts')
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/finder/media', methods=['GET'])
@require_api_key
def finder_media():
    limit = int(request.args.get('limit', 20))
    try:
        result = call_rpc('getFinderMedia', limit)
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 收藏接口 ====================

@app.route('/api/v1/favorites', methods=['GET'])
@require_api_key
def get_favorites():
    limit = int(request.args.get('limit', 20))
    try:
        result = call_rpc('getFavorites', limit)
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 钱包接口 ====================

@app.route('/api/v1/wallet/info', methods=['GET'])
@require_api_key
def wallet_info():
    try:
        result = call_rpc('getWalletInfo')
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/wallet/ledger', methods=['GET'])
@require_api_key
def wallet_ledger():
    limit = int(request.args.get('limit', 20))
    try:
        result = call_rpc('getWalletLedger', limit)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/wallet/lucky-money', methods=['GET'])
@require_api_key
def lucky_money():
    limit = int(request.args.get('limit', 20))
    try:
        result = call_rpc('getLuckyMoney', limit)
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 媒体接口 ====================

@app.route('/api/v1/media/images', methods=['GET'])
@require_api_key
def get_images():
    limit = int(request.args.get('limit', 20))
    try:
        result = call_rpc('getImages', limit)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/media/images/count', methods=['GET'])
@require_api_key
def images_count():
    try:
        result = call_rpc('getImageCount')
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/media/videos', methods=['GET'])
@require_api_key
def get_videos():
    limit = int(request.args.get('limit', 20))
    try:
        result = call_rpc('getVideos', limit)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/media/files', methods=['GET'])
@require_api_key
def get_files():
    limit = int(request.args.get('limit', 20))
    try:
        result = call_rpc('getFiles', limit)
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 小程序接口 ====================

@app.route('/api/v1/appbrand', methods=['GET'])
@require_api_key
def get_appbrand():
    try:
        result = call_rpc('getAppBrandInfo')
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 搜索接口 ====================

@app.route('/api/v1/search', methods=['GET'])
@require_api_key
def search():
    keyword = request.args.get('keyword', '')
    if not keyword:
        return error('keyword不能为空', 400)
    try:
        result = call_rpc('globalSearch', keyword)
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 原始SQL接口 ====================

@app.route('/api/v1/raw/query', methods=['POST'])
@require_api_key
def raw_query():
    data = request.get_json()
    db = data.get('db', 'main')
    sql = data.get('sql', '')
    if not sql:
        return error('sql不能为空', 400)
    try:
        result = call_rpc('rawSql', db, sql)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/raw/write', methods=['POST'])
@require_api_key
def raw_write():
    data = request.get_json()
    db = data.get('db', 'main')
    sql = data.get('sql', '')
    if not sql:
        return error('sql不能为空', 400)
    # 安全检查
    sql_upper = sql.strip().upper()
    if sql_upper.startswith('DROP') or sql_upper.startswith('TRUNCATE'):
        return error('不允许DROP/TRUNCATE操作', 403)
    try:
        result = call_rpc('rawSqlWrite', db, sql)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/raw/tables', methods=['GET'])
@require_api_key
def list_tables():
    db = request.args.get('db', 'main')
    try:
        result = call_rpc('listTables', db)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v1/raw/table', methods=['GET'])
@require_api_key
def table_info():
    db = request.args.get('db', 'main')
    table = request.args.get('table', '')
    if not table:
        return error('table不能为空', 400)
    try:
        result = call_rpc('getTableInfo', db, table)
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 存客宝对接接口 ====================

@app.route('/api/v3/cunke-bao/contacts/sync', methods=['POST'])
@require_api_key
def ckb_sync_contacts():
    """批量同步联系人到存客宝"""
    data = request.get_json()
    limit = data.get('limit', 100)
    offset = data.get('offset', 0)
    try:
        contacts = call_rpc('getContacts', limit, offset)
        # 格式化为存客宝格式
        ckb_contacts = []
        for c in (contacts.get('contacts', []) if isinstance(contacts, dict) else []):
            ckb_contacts.append({
                'wxid': c.get('username', ''),
                'nickname': c.get('nickname', ''),
                'remark': c.get('remark', ''),
                'avatar': c.get('headImgUrl', ''),
                'labels': c.get('contactLabelIds', ''),
                'sex': c.get('sex', 0),
                'region': f"{c.get('province', '')} {c.get('city', '')}".strip(),
            })
        return success({'contacts': ckb_contacts, 'count': len(ckb_contacts), 'total': contacts.get('total', 0) if isinstance(contacts, dict) else 0})
    except Exception as e:
        return error(e)

@app.route('/api/v3/cunke-bao/message/send', methods=['POST'])
@require_api_key
def ckb_send_message():
    """存客宝发送消息"""
    data = request.get_json()
    to_user = data.get('wxid', data.get('to_user', ''))
    content = data.get('content', data.get('message', ''))
    if not to_user or not content:
        return error('wxid和content不能为空', 400)
    try:
        result = call_rpc('sendMessage', to_user, content, 1)
        return success(result)
    except Exception as e:
        return error(e)

@app.route('/api/v3/cunke-bao/groups', methods=['GET'])
@require_api_key
def ckb_get_groups():
    """存客宝获取群列表"""
    try:
        result = call_rpc('getGroups', 100)
        return success(result)
    except Exception as e:
        return error(e)

# ==================== 健康检查和自动重连 ====================

def heartbeat_thread():
    """后台心跳线程，保持Frida连接"""
    while True:
        try:
            time.sleep(60)
            if state.connected:
                try:
                    call_rpc('ping')
                    logger.debug("心跳正常")
                except Exception as e:
                    logger.warning(f"心跳失败: {e}，尝试重连...")
                    state.connected = False
                    connect_frida()
        except Exception as e:
            logger.error(f"心跳线程异常: {e}")

if __name__ == '__main__':
    print("=" * 60)
    print("工作手机 SDK API 服务 v9-complete")
    print(f"手机IP: {PHONE_IP}")
    print(f"API Key: {API_KEY}")
    print("=" * 60)
    
    # 初始连接
    logger.info("正在连接Frida...")
    if connect_frida():
        logger.info("✅ Frida连接成功")
    else:
        logger.warning("⚠️ Frida初始连接失败，将在请求时重试")
    
    # 启动心跳线程
    t = threading.Thread(target=heartbeat_thread, daemon=True)
    t.start()
    
    # 启动API服务
    app.run(host='0.0.0.0', port=8899, debug=False, threaded=True)