fws_code/api_migu.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
咪咕短剧 flows.cdyylkj.com/miguSM/home 接口：
- POST /api/submit_phone  传入电话号码 → 新建 TgeBrowser → 输入手机号、点击发送验证码 → 返回成功
- POST /api/submit_code   传入验证码 → 在对应会话中填写验证码并提交

每个电话号码对应一个新建的 TgeBrowser 浏览器会话。
"""
from __future__ import annotations

import threading
import uuid
from typing import Any, Optional

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

from migu_miguSM_dp import (
    MIGU_HOME_URL,
    connect_dp_to_tgebrowser,
    connect_dp_to_ws,
    input_code_and_submit,
    input_phone_and_send_code,
)
from tgebrowser_client import TgeBrowserClient

app = FastAPI(title="咪咕短剧 miguSM 自动化 API", description="TgeBrowser + DrissionPage 手机号与验证码提交流程")

# 会话存储：session_id -> {page, env_id, client}
_sessions: dict[str, dict[str, Any]] = {}
_sessions_lock = threading.Lock()


class SubmitPhoneRequest(BaseModel):
    phone: str = Field(..., description="手机号码")
    url: str = Field(MIGU_HOME_URL, description="目标页面 URL")


class SubmitPhoneResponse(BaseModel):
    success: bool = True
    session_id: str = Field(..., description="会话 ID，submit_code 时需传入")
    message: str = "输入电话号码成功"


class SubmitCodeRequest(BaseModel):
    session_id: str = Field(..., description="submit_phone 返回的会话 ID")
    code: str = Field(..., description="短信验证码")


class SubmitCodeResponse(BaseModel):
    success: bool = True
    message: str = ""


@app.post("/api/submit_phone", response_model=SubmitPhoneResponse)
def api_submit_phone(req: SubmitPhoneRequest):
    """
    步骤一：传入电话号码。
    新建 TgeBrowser 浏览器 → 打开 miguSM 页面 → 输入手机号 → 点击发送验证码。
    返回 session_id，用于后续 submit_code。
    """
    try:
        client = TgeBrowserClient()
        # 每次输入电话号码新建一个浏览器
        start_data = client.create_and_start(
            browser_name=f"miguSM_{req.phone[-4:]}",
            start_page_url=req.url,
        )
        port = start_data.get("port")
        ws = start_data.get("ws")
        env_id = start_data.get("envId")

        # 优先用端口连接（DrissionPage 对端口支持更好）
        if port:
            page = connect_dp_to_tgebrowser(port)
        elif ws:
            page = connect_dp_to_ws(ws)
        else:
            raise HTTPException(status_code=500, detail="TgeBrowser 未返回 port 或 ws")

        input_phone_and_send_code(page, req.phone, url=req.url)

        session_id = str(uuid.uuid4())
        with _sessions_lock:
            _sessions[session_id] = {
                "page": page,
                "env_id": env_id,
                "client": client,
            }

        return SubmitPhoneResponse(success=True, session_id=session_id, message="输入电话号码成功")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/submit_code", response_model=SubmitCodeResponse)
def api_submit_code(req: SubmitCodeRequest):
    """
    步骤二：传入验证码。
    在 submit_phone 创建的会话中填写验证码并提交。
    """
    with _sessions_lock:
        sess = _sessions.get(req.session_id)
        if not sess:
            raise HTTPException(status_code=404, detail=f"会话不存在或已过期: {req.session_id}")

    page = sess["page"]
    try:
        result = input_code_and_submit(page, req.code)
        return SubmitCodeResponse(success=True, message=result.get("message", "验证码已填写"))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.delete("/api/session/{session_id}")
def close_session(session_id: str):
    """关闭并清理指定会话（可选，用于释放浏览器）。"""
    with _sessions_lock:
        sess = _sessions.pop(session_id, None)
    if not sess:
        return {"success": False, "message": "会话不存在"}
    try:
        client = sess.get("client")
        env_id = sess.get("env_id")
        if client and env_id is not None:
            client.stop_browser(env_id=env_id)
    except Exception:
        pass
    return {"success": True, "message": "会话已关闭"}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8001)