479 lines
13 KiB
Python
479 lines
13 KiB
Python
|
||
"""web.py - FastAPI 后端
|
||
|
||
目标：
|
||
- 提供基本的 Web UI + API（启动/停止采集、导入用户、登录/验证、日志、SSE）
|
||
- 单账号登录态长期复用（storage_state 保存）
|
||
- txt 导入：支持“数字小红书号”自动解析为 profile_id（并缓存，低频，减少风控）
|
||
- 移除无用/冗余的调试/自更新接口，代码保持清晰可维护
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
from dataclasses import dataclass
|
||
from typing import List, Optional, Dict, Any
|
||
|
||
from fastapi import FastAPI, HTTPException, BackgroundTasks
|
||
from fastapi.responses import HTMLResponse, StreamingResponse, PlainTextResponse
|
||
from fastapi.staticfiles import StaticFiles
|
||
from fastapi.templating import Jinja2Templates
|
||
from fastapi.requests import Request
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from pydantic import BaseModel
|
||
|
||
from playwright.async_api import async_playwright, Browser
|
||
|
||
from app import config
|
||
from app.log_manager import LogManager
|
||
from app.scraper import XiaoHongShuScraper
|
||
from app import auth_manager
|
||
from app import user_resolver
|
||
|
||
|
||
# =========================
|
||
# App State
|
||
# =========================
|
||
|
||
@dataclass
|
||
class AppState:
|
||
is_running: bool = False
|
||
scraper: Optional[XiaoHongShuScraper] = None
|
||
event_queue: Optional[asyncio.Queue] = None
|
||
playwright: Any = None
|
||
browser: Optional[Browser] = None
|
||
|
||
|
||
app_state = AppState(event_queue=asyncio.Queue())
|
||
log_manager = LogManager()
|
||
|
||
|
||
# =========================
|
||
# FastAPI init
|
||
# =========================
|
||
|
||
app = FastAPI(title="小红书数据采集器", version="1.1.0")
|
||
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
expose_headers=["*"],
|
||
)
|
||
|
||
# 静态资源 + 模板
|
||
app.mount("/static", StaticFiles(directory=os.path.join(config.BASE_DIR, "static")), name="static")
|
||
templates = Jinja2Templates(directory=os.path.join(config.BASE_DIR, "templates"))
|
||
|
||
|
||
@app.on_event("startup")
|
||
async def on_startup():
|
||
# 启动 Playwright + Browser（长期复用，减少反复启动带来的抖动）
|
||
app_state.playwright = await async_playwright().start()
|
||
app_state.browser = await app_state.playwright.chromium.launch(
|
||
headless=config.HEADLESS,
|
||
args=config.BROWSER_ARGS,
|
||
)
|
||
log_manager.append("INFO", f"Browser launched (headless={config.HEADLESS})")
|
||
|
||
|
||
@app.on_event("shutdown")
|
||
async def on_shutdown():
|
||
# 安全关闭
|
||
try:
|
||
if app_state.browser:
|
||
await app_state.browser.close()
|
||
finally:
|
||
app_state.browser = None
|
||
|
||
try:
|
||
if app_state.playwright:
|
||
await app_state.playwright.stop()
|
||
finally:
|
||
app_state.playwright = None
|
||
|
||
|
||
# =========================
|
||
# Models
|
||
# =========================
|
||
|
||
class StartRequest(BaseModel):
|
||
mode: str = "single" # single / multi
|
||
max_threads: int = 3
|
||
|
||
|
||
class StatusResponse(BaseModel):
|
||
running: bool
|
||
mode: Optional[str] = None
|
||
user_count: int = 0
|
||
|
||
|
||
class UsersResponse(BaseModel):
|
||
users: List[str]
|
||
count: int
|
||
|
||
|
||
class ImportResponse(BaseModel):
|
||
users: List[str]
|
||
count: int
|
||
total_lines: int
|
||
imported: int
|
||
skipped: int
|
||
resolved: int
|
||
failed_resolve: int
|
||
duplicates: int
|
||
invalid: int
|
||
message: str
|
||
# 可选返回一些失败明细（避免过大，只给前 20 条）
|
||
failed_examples: List[str] = []
|
||
|
||
|
||
class LoginStatusResponse(BaseModel):
|
||
logged_in: bool
|
||
login_time: Optional[str] = None
|
||
last_verified_time: Optional[str] = None
|
||
message: Optional[str] = None
|
||
|
||
|
||
# =========================
|
||
# Pages
|
||
# =========================
|
||
|
||
@app.get("/", response_class=HTMLResponse)
|
||
async def index(request: Request):
|
||
return templates.TemplateResponse("index.html", {"request": request})
|
||
|
||
|
||
# =========================
|
||
# Helpers
|
||
# =========================
|
||
|
||
def _users_file() -> str:
|
||
return os.path.join(config.DATA_DIR, "users.txt")
|
||
|
||
|
||
def read_users_txt() -> List[str]:
|
||
path = _users_file()
|
||
if not os.path.exists(path):
|
||
return []
|
||
users: List[str] = []
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
users.append(line)
|
||
# 去重但保持顺序
|
||
seen = set()
|
||
out = []
|
||
for u in users:
|
||
if u in seen:
|
||
continue
|
||
seen.add(u)
|
||
out.append(u)
|
||
return out
|
||
|
||
|
||
def write_users_txt(users: List[str]) -> None:
|
||
path = _users_file()
|
||
tmp = path + ".tmp"
|
||
with open(tmp, "w", encoding="utf-8") as f:
|
||
for u in users:
|
||
f.write(u + "\n")
|
||
os.replace(tmp, path)
|
||
|
||
|
||
def require_browser() -> Browser:
|
||
if app_state.browser is None:
|
||
raise HTTPException(status_code=500, detail="Browser not ready")
|
||
return app_state.browser
|
||
|
||
|
||
# =========================
|
||
# Status / Logs
|
||
# =========================
|
||
|
||
@app.get("/api/status", response_model=StatusResponse)
|
||
async def api_status():
|
||
return StatusResponse(
|
||
running=app_state.is_running,
|
||
mode=getattr(app_state.scraper, "mode", None) if app_state.scraper else None,
|
||
user_count=len(read_users_txt()),
|
||
)
|
||
|
||
|
||
@app.post("/api/logs/clear")
|
||
async def clear_logs():
|
||
log_manager.clear()
|
||
return {"success": True}
|
||
|
||
|
||
@app.get("/api/logs")
|
||
async def get_logs(limit: int = 200):
|
||
return {"logs": log_manager.get_latest(limit=limit)}
|
||
|
||
|
||
# =========================
|
||
# Users
|
||
# =========================
|
||
|
||
@app.get("/api/users", response_model=UsersResponse)
|
||
async def get_users():
|
||
users = read_users_txt()
|
||
return UsersResponse(users=users, count=len(users))
|
||
|
||
|
||
@app.post("/api/users")
|
||
async def add_user(user_id: str):
|
||
user_id = (user_id or "").strip()
|
||
if not user_id:
|
||
raise HTTPException(status_code=400, detail="user_id 不能为空")
|
||
|
||
users = read_users_txt()
|
||
if user_id in users:
|
||
return {"success": True, "message": "已存在"}
|
||
|
||
users.append(user_id)
|
||
write_users_txt(users)
|
||
return {"success": True}
|
||
|
||
|
||
@app.delete("/api/users/{user_id}")
|
||
async def delete_user(user_id: str):
|
||
users = read_users_txt()
|
||
users = [u for u in users if u != user_id]
|
||
write_users_txt(users)
|
||
return {"success": True}
|
||
|
||
|
||
@app.delete("/api/users")
|
||
async def delete_all_users():
|
||
write_users_txt([])
|
||
return {"success": True}
|
||
|
||
|
||
@app.get("/api/users/export/txt")
|
||
async def export_users_txt():
|
||
return PlainTextResponse("\\n".join(read_users_txt()))
|
||
|
||
|
||
@app.get("/api/users/import/txt", response_model=ImportResponse)
|
||
async def import_users_from_txt():
|
||
"""从 data/users.txt 读取并规范化：
|
||
- 若是 profile_id / profile URL：直接收录
|
||
- 若是纯数字小红书号：通过搜索页解析成 profile_id（Playwright，低频、缓存）
|
||
- 输出：users.txt 最终仅保存 profile_id（稳定可采集）
|
||
"""
|
||
path = _users_file()
|
||
if not os.path.exists(path):
|
||
raise HTTPException(status_code=400, detail="data/users.txt 不存在")
|
||
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
raw_lines = [l.rstrip("\n") for l in f]
|
||
|
||
total_lines = len(raw_lines)
|
||
|
||
browser = require_browser()
|
||
|
||
imported_ids: List[str] = []
|
||
duplicates = 0
|
||
invalid = 0
|
||
skipped = 0
|
||
resolved = 0
|
||
failed_resolve = 0
|
||
failed_examples: List[str] = []
|
||
|
||
seen = set()
|
||
|
||
# 第一轮：先把能直接解析的收集，数字 keyword 单独列表
|
||
keywords: List[tuple[str, str]] = [] # (keyword, raw_line)
|
||
for line in raw_lines:
|
||
kw, pid, _label = user_resolver.parse_line_to_keyword_or_profile(line)
|
||
if pid:
|
||
pid = pid.lower()
|
||
if pid in seen:
|
||
duplicates += 1
|
||
continue
|
||
seen.add(pid)
|
||
imported_ids.append(pid)
|
||
continue
|
||
if kw:
|
||
keywords.append((kw, line))
|
||
else:
|
||
# 空行/注释等视为 skipped；其它无法识别为 invalid
|
||
if not line.strip() or line.strip().startswith("#") or line.strip().startswith("//"):
|
||
skipped += 1
|
||
else:
|
||
invalid += 1
|
||
if len(failed_examples) < 20:
|
||
failed_examples.append(f"invalid: {line}")
|
||
continue
|
||
|
||
# 第二轮：逐个解析 keyword（单线程+慢一点，减少风控）
|
||
for kw, raw in keywords:
|
||
pid, msg = await user_resolver.resolve_keyword_to_profile_id(browser, kw, slow_mode=True)
|
||
if not pid:
|
||
failed_resolve += 1
|
||
if len(failed_examples) < 20:
|
||
failed_examples.append(f"resolve_failed({kw}): {raw} ({msg})")
|
||
continue
|
||
pid = pid.lower()
|
||
if pid in seen:
|
||
duplicates += 1
|
||
continue
|
||
seen.add(pid)
|
||
imported_ids.append(pid)
|
||
resolved += 1
|
||
|
||
# 每解析一个稍微间隔一下（更像真人）
|
||
await asyncio.sleep(0.8)
|
||
|
||
# 写回：只保存 profile_id
|
||
write_users_txt(imported_ids)
|
||
|
||
imported = len(imported_ids)
|
||
message = f"导入完成：{imported} 个 profile_id（解析成功 {resolved}）"
|
||
|
||
log_manager.append("INFO", message)
|
||
|
||
return ImportResponse(
|
||
users=imported_ids,
|
||
count=imported,
|
||
total_lines=total_lines,
|
||
imported=imported,
|
||
skipped=skipped,
|
||
resolved=resolved,
|
||
failed_resolve=failed_resolve,
|
||
duplicates=duplicates,
|
||
invalid=invalid,
|
||
message=message,
|
||
failed_examples=failed_examples,
|
||
)
|
||
|
||
|
||
# =========================
|
||
# Login
|
||
# =========================
|
||
|
||
@app.get("/api/login/status", response_model=LoginStatusResponse)
|
||
async def get_login_status():
|
||
status = await auth_manager.get_login_status(app_state.browser)
|
||
return LoginStatusResponse(
|
||
logged_in=status.logged_in,
|
||
login_time=status.login_time,
|
||
last_verified_time=status.last_verified_time,
|
||
message=status.message,
|
||
)
|
||
|
||
|
||
@app.post("/api/login/qrcode")
|
||
async def open_login():
|
||
"""兼容旧前端：返回 login_url，同时在服务端打开 Playwright 登录页。
|
||
|
||
说明：前端 window.open 打开的页面不会影响服务端 Playwright 的登录态。
|
||
正确使用方式：在部署环境（Docker + VNC）里打开 VNC 浏览器扫码登录。
|
||
"""
|
||
browser = require_browser()
|
||
await auth_manager.open_login_page(browser)
|
||
return {
|
||
"login_url": config.LOGIN_URL,
|
||
"hint": "已在服务端打开登录页。请在 VNC/有界面环境里扫码登录，然后点击“验证登录”。"
|
||
}
|
||
|
||
|
||
@app.post("/api/login/verify")
|
||
async def verify_login():
|
||
browser = require_browser()
|
||
ok, msg = await auth_manager.verify_login(browser)
|
||
meta = auth_manager.load_auth_meta()
|
||
return {
|
||
"success": ok,
|
||
"message": None if ok else msg,
|
||
"login_time": meta.get("login_time"),
|
||
"expire_time": None, # 单账号长期复用：不强制本地过期时间
|
||
"last_verified_time": meta.get("last_verified_time"),
|
||
}
|
||
|
||
|
||
@app.post("/api/logout")
|
||
async def logout():
|
||
auth_manager.clear_login_state()
|
||
return {"success": True}
|
||
|
||
|
||
# =========================
|
||
# Scrape control
|
||
# =========================
|
||
|
||
@app.post("/api/start")
|
||
async def start_scraping(request: StartRequest, background_tasks: BackgroundTasks):
|
||
if app_state.is_running:
|
||
raise HTTPException(status_code=400, detail="采集任务正在运行中")
|
||
|
||
users = read_users_txt()
|
||
if not users:
|
||
raise HTTPException(status_code=400, detail="用户列表为空，请先导入 users.txt")
|
||
|
||
# 启动前校验登录态（单账号）
|
||
browser = require_browser()
|
||
ok, msg = await auth_manager.verify_login(browser)
|
||
if not ok:
|
||
raise HTTPException(status_code=400, detail=f"未登录或登录失效：{msg}")
|
||
|
||
app_state.is_running = True
|
||
app_state.scraper = XiaoHongShuScraper(app_state.event_queue, log_manager, browser)
|
||
|
||
async def run_scraper():
|
||
try:
|
||
await app_state.scraper.run(user_ids=users, mode=request.mode, max_threads=request.max_threads)
|
||
finally:
|
||
app_state.is_running = False
|
||
|
||
background_tasks.add_task(run_scraper)
|
||
return {"success": True}
|
||
|
||
|
||
@app.post("/api/stop")
|
||
async def stop_scraping():
|
||
if not app_state.is_running or not app_state.scraper:
|
||
return {"success": True, "message": "未在运行"}
|
||
|
||
await app_state.scraper.stop()
|
||
app_state.is_running = False
|
||
return {"success": True}
|
||
|
||
|
||
# =========================
|
||
# Data
|
||
# =========================
|
||
|
||
@app.get("/api/data/{user_id}")
|
||
async def get_user_data(user_id: str):
|
||
# 读取该用户数据文件（若存在）
|
||
file_path = os.path.join(config.DATA_DIR, f"{user_id}.json")
|
||
if not os.path.exists(file_path):
|
||
raise HTTPException(status_code=404, detail="数据不存在")
|
||
try:
|
||
with open(file_path, "r", encoding="utf-8") as f:
|
||
return json.load(f)
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
# =========================
|
||
# SSE events
|
||
# =========================
|
||
|
||
@app.get("/api/events")
|
||
async def sse_events():
|
||
async def event_generator():
|
||
while True:
|
||
event = await app_state.event_queue.get()
|
||
yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"
|
||
|
||
return StreamingResponse(event_generator(), media_type="text/event-stream")
|