ai_wht_wechat/backend/main.py

from fastapi import FastAPI, HTTPException, File, UploadFile, Form
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional, Dict, Any, List
import asyncio
from datetime import datetime
import os
import shutil
from pathlib import Path

from xhs_login import XHSLoginService

app = FastAPI(title="小红书登录API")

# CORS配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境应该限制具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 全局登录服务实例
login_service = XHSLoginService()

# 临时文件存储目录
TEMP_DIR = Path("temp_uploads")
TEMP_DIR.mkdir(exist_ok=True)

# 请求模型
class SendCodeRequest(BaseModel):
    phone: str
    country_code: str = "+86"

class LoginRequest(BaseModel):
    phone: str
    code: str
    country_code: str = "+86"

class PublishNoteRequest(BaseModel):
    title: str
    content: str
    images: Optional[list] = None
    topics: Optional[list] = None

class InjectCookiesRequest(BaseModel):
    cookies: list

# 响应模型
class BaseResponse(BaseModel):
    code: int
    message: str
    data: Optional[Dict[str, Any]] = None

@app.on_event("startup")
async def startup_event():
    """启动时不初始化浏览器，等待第一次请求时再初始化"""
    pass

@app.on_event("shutdown")
async def shutdown_event():
    """关闭时清理浏览器"""
    await login_service.close_browser()

@app.post("/api/xhs/send-code", response_model=BaseResponse)
async def send_code(request: SendCodeRequest):
    """
    发送验证码
    通过playwright访问小红书官网，输入手机号并触发验证码发送
    """
    try:
        # 调用登录服务发送验证码
        result = await login_service.send_verification_code(
            phone=request.phone,
            country_code=request.country_code
        )
        
        if result["success"]:
            return BaseResponse(
                code=0,
                message="验证码已发送，请在小红书APP中查看",
                data={"sent_at": datetime.now().isoformat()}
            )
        else:
            return BaseResponse(
                code=1,
                message=result.get("error", "发送验证码失败"),
                data=None
            )
    
    except Exception as e:
        print(f"发送验证码异常: {str(e)}")
        return BaseResponse(
            code=1,
            message=f"发送验证码失败: {str(e)}",
            data=None
        )

@app.post("/api/xhs/login", response_model=BaseResponse)
async def login(request: LoginRequest):
    """
    登录验证
    用户填写验证码后，完成登录并获取小红书返回的数据
    """
    try:
        # 调用登录服务进行登录
        result = await login_service.login(
            phone=request.phone,
            code=request.code,
            country_code=request.country_code
        )
        
        if result["success"]:
            return BaseResponse(
                code=0,
                message="登录成功",
                data={
                    "user_info": result.get("user_info"),
                    "cookies": result.get("cookies"),           # 键值对格式（前端展示）
                    "cookies_full": result.get("cookies_full"),  # Playwright完整格式（数据库存储/脚本使用）
                    "login_time": datetime.now().isoformat()
                }
            )
        else:
            return BaseResponse(
                code=1,
                message=result.get("error", "登录失败"),
                data=None
            )
    
    except Exception as e:
        print(f"登录异常: {str(e)}")
        return BaseResponse(
            code=1,
            message=f"登录失败: {str(e)}",
            data=None
        )

@app.get("/")
async def root():
    """健康检查"""
    return {"status": "ok", "message": "小红书登录服务运行中"}

@app.post("/api/xhs/inject-cookies", response_model=BaseResponse)
async def inject_cookies(request: InjectCookiesRequest):
    """
    注入Cookies并验证登录状态
    允许使用之前保存的Cookies跳过登录
    """
    try:
        # 关闭旧的浏览器（如果有）
        if login_service.browser:
            await login_service.close_browser()
        
        # 使用Cookies初始化浏览器
        await login_service.init_browser(cookies=request.cookies)
        
        # 验证登录状态
        result = await login_service.verify_login_status()
        
        if result.get("logged_in"):
            return BaseResponse(
                code=0,
                message="Cookie注入成功，已登录",
                data={
                    "logged_in": True,
                    "user_info": result.get("user_info"),
                    "cookies": result.get("cookies"),           # 键值对格式
                    "cookies_full": result.get("cookies_full"),  # Playwright完整格式
                    "url": result.get("url")
                }
            )
        else:
            return BaseResponse(
                code=1,
                message=result.get("message", "Cookie已失效，请重新登录"),
                data={
                    "logged_in": False
                }
            )
    
    except Exception as e:
        print(f"注入Cookies异常: {str(e)}")
        return BaseResponse(
            code=1,
            message=f"注入失败: {str(e)}",
            data=None
        )

@app.post("/api/xhs/publish", response_model=BaseResponse)
async def publish_note(request: PublishNoteRequest):
    """
    发布笔记
    登录后可以发布图文笔记到小红书
    """
    try:
        # 调用登录服务发布笔记
        result = await login_service.publish_note(
            title=request.title,
            content=request.content,
            images=request.images,
            topics=request.topics
        )
        
        if result["success"]:
            return BaseResponse(
                code=0,
                message="笔记发布成功",
                data={
                    "url": result.get("url"),
                    "publish_time": datetime.now().isoformat()
                }
            )
        else:
            return BaseResponse(
                code=1,
                message=result.get("error", "发布失败"),
                data=None
            )
    
    except Exception as e:
        print(f"发布笔记异常: {str(e)}")
        return BaseResponse(
            code=1,
            message=f"发布失败: {str(e)}",
            data=None
        )

@app.post("/api/xhs/upload-images")
async def upload_images(files: List[UploadFile] = File(...)):
    """
    上传图片到服务器临时目录
    返回图片的服务器路径
    """
    try:
        uploaded_paths = []
        
        for file in files:
            # 检查文件类型
            if not file.content_type.startswith('image/'):
                return {
                    "code": 1,
                    "message": f"文件 {file.filename} 不是图片类型",
                    "data": None
                }
            
            # 生成唯一文件名
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            file_ext = os.path.splitext(file.filename)[1]
            safe_filename = f"{timestamp}{file_ext}"
            file_path = TEMP_DIR / safe_filename
            
            # 保存文件
            with open(file_path, "wb") as buffer:
                shutil.copyfileobj(file.file, buffer)
            
            # 使用绝对路径
            abs_path = str(file_path.absolute())
            uploaded_paths.append(abs_path)
            print(f"已上传图片: {abs_path}")
        
        return {
            "code": 0,
            "message": f"成功上传 {len(uploaded_paths)} 张图片",
            "data": {
                "paths": uploaded_paths,
                "count": len(uploaded_paths)
            }
        }
    
    except Exception as e:
        print(f"上传图片异常: {str(e)}")
        return {
            "code": 1,
            "message": f"上传失败: {str(e)}",
            "data": None
        }

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