924 lines
28 KiB
Python
924 lines
28 KiB
Python
"""
|
|
Web Backend - FastAPI 应用
|
|
提供 API 和 SSE 事件推送
|
|
"""
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import subprocess
|
|
from datetime import datetime, timedelta
|
|
from typing import Dict, List, Optional
|
|
from fastapi import FastAPI, HTTPException, BackgroundTasks
|
|
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
from fastapi.templating import Jinja2Templates
|
|
from fastapi.requests import Request
|
|
from pydantic import BaseModel
|
|
from typing import List as ListType
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
|
|
from app import config
|
|
from app.log_manager import get_log_manager
|
|
from app.scraper import XiaoHongShuScraper
|
|
|
|
|
|
# ========== Pydantic 模型 ==========
|
|
class StartRequest(BaseModel):
|
|
mode: str = "single" # single / multi
|
|
max_threads: int = 3
|
|
|
|
|
|
class StatusResponse(BaseModel):
|
|
is_running: bool
|
|
mode: str
|
|
max_threads: int
|
|
total_users: int
|
|
completed_users: int
|
|
failed_users: int
|
|
active_threads: int
|
|
users_status: Dict[str, Dict]
|
|
|
|
|
|
# ========== 全局状态 ==========
|
|
class AppState:
|
|
def __init__(self):
|
|
self.is_running = False
|
|
self.mode = "single"
|
|
self.max_threads = 1
|
|
self.total_users = 0
|
|
self.completed_users = 0
|
|
self.failed_users = 0
|
|
self.active_threads = 0
|
|
self.users_status: Dict[str, Dict] = {}
|
|
self.thread_status: Dict[str, Dict] = {}
|
|
self.event_queue: asyncio.Queue = None
|
|
self.scraper: Optional[XiaoHongShuScraper] = None
|
|
self.scraper_task: Optional[asyncio.Task] = None
|
|
|
|
|
|
app_state = AppState()
|
|
log_manager = get_log_manager()
|
|
|
|
# ========== FastAPI 应用 ==========
|
|
app = FastAPI(title="小红书数据采集器", version="1.0.0")
|
|
|
|
# 添加 CORS 中间件
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"], # 生产环境应该指定具体域名
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
expose_headers=["*"]
|
|
)
|
|
|
|
# 挂载静态文件和模板
|
|
app.mount("/static", StaticFiles(directory="static"), name="static")
|
|
templates = Jinja2Templates(directory="templates")
|
|
|
|
|
|
# ========== 路由 ==========
|
|
|
|
@app.get("/", response_class=HTMLResponse)
|
|
async def index(request: Request):
|
|
"""首页"""
|
|
return templates.TemplateResponse("index.html", {"request": request})
|
|
|
|
|
|
|
|
|
|
@app.get("/api/status", response_model=StatusResponse)
|
|
async def get_status():
|
|
"""获取当前状态"""
|
|
return StatusResponse(
|
|
is_running=app_state.is_running,
|
|
mode=app_state.mode,
|
|
max_threads=app_state.max_threads,
|
|
total_users=app_state.total_users,
|
|
completed_users=app_state.completed_users,
|
|
failed_users=app_state.failed_users,
|
|
active_threads=app_state.active_threads,
|
|
users_status=app_state.users_status
|
|
)
|
|
|
|
|
|
@app.post("/api/start")
|
|
async def start_scraping(request: StartRequest):
|
|
"""启动采集"""
|
|
if app_state.is_running:
|
|
raise HTTPException(status_code=400, detail="采集任务正在运行中")
|
|
|
|
# 读取用户 ID 列表
|
|
users_file = os.path.join(config.DATA_DIR, "users.txt")
|
|
if not os.path.exists(users_file):
|
|
raise HTTPException(status_code=400, detail="用户文件不存在，请先创建 data/users.txt")
|
|
|
|
with open(users_file, 'r', encoding='utf-8') as f:
|
|
user_ids = [line.strip() for line in f if line.strip()]
|
|
|
|
if not user_ids:
|
|
raise HTTPException(status_code=400, detail="用户 ID 列表为空")
|
|
|
|
# 验证模式
|
|
if request.mode not in ["single", "multi"]:
|
|
raise HTTPException(status_code=400, detail="无效的模式，必须是 single 或 multi")
|
|
|
|
# 验证线程数
|
|
if request.mode == "multi":
|
|
if request.max_threads < 1 or request.max_threads > config.MAX_CONCURRENT_THREADS:
|
|
raise HTTPException(status_code=400, detail=f"线程数必须在 1-{config.MAX_CONCURRENT_THREADS} 之间")
|
|
|
|
# 重置状态
|
|
app_state.is_running = True
|
|
app_state.mode = request.mode
|
|
app_state.max_threads = request.max_threads if request.mode == "multi" else 1
|
|
app_state.total_users = len(user_ids)
|
|
app_state.completed_users = 0
|
|
app_state.failed_users = 0
|
|
app_state.active_threads = 0
|
|
app_state.users_status = {}
|
|
app_state.thread_status = {}
|
|
|
|
# 初始化用户状态
|
|
for user_id in user_ids:
|
|
app_state.users_status[user_id] = {
|
|
"status": "pending",
|
|
"note_count": 0,
|
|
"thread_id": None
|
|
}
|
|
|
|
# 创建事件队列
|
|
app_state.event_queue = asyncio.Queue()
|
|
|
|
# 创建采集器
|
|
app_state.scraper = XiaoHongShuScraper(app_state.event_queue, log_manager)
|
|
|
|
# 启动采集任务
|
|
async def run_scraper():
|
|
try:
|
|
await app_state.scraper.run(
|
|
user_ids=user_ids,
|
|
mode=request.mode,
|
|
max_threads=request.max_threads
|
|
)
|
|
finally:
|
|
app_state.is_running = False
|
|
await app_state.event_queue.put({
|
|
"type": "finished"
|
|
})
|
|
|
|
app_state.scraper_task = asyncio.create_task(run_scraper())
|
|
|
|
return {
|
|
"success": True,
|
|
"message": f"采集任务已启动，模式: {request.mode}, 用户数: {len(user_ids)}",
|
|
"total_users": len(user_ids),
|
|
"mode": request.mode,
|
|
"max_threads": app_state.max_threads
|
|
}
|
|
|
|
|
|
@app.post("/api/stop")
|
|
async def stop_scraping():
|
|
"""停止采集"""
|
|
if not app_state.is_running:
|
|
raise HTTPException(status_code=400, detail="没有正在运行的采集任务")
|
|
|
|
if app_state.scraper:
|
|
app_state.scraper.stop()
|
|
|
|
# 取消任务
|
|
if app_state.scraper_task:
|
|
app_state.scraper_task.cancel()
|
|
try:
|
|
await app_state.scraper_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
app_state.is_running = False
|
|
|
|
return {
|
|
"success": True,
|
|
"message": "采集任务已停止"
|
|
}
|
|
|
|
|
|
@app.post("/api/logs/clear")
|
|
async def clear_logs():
|
|
"""清空日志"""
|
|
log_manager.clear()
|
|
return {
|
|
"success": True,
|
|
"message": "日志已清空"
|
|
}
|
|
|
|
|
|
@app.get("/api/logs")
|
|
async def get_logs(level: Optional[str] = None, limit: Optional[int] = None):
|
|
"""获取日志"""
|
|
logs = log_manager.get_logs(level=level, limit=limit)
|
|
return {
|
|
"logs": logs,
|
|
"stats": log_manager.get_stats()
|
|
}
|
|
|
|
|
|
@app.get("/api/logs/search")
|
|
async def search_logs(keyword: str, limit: Optional[int] = None):
|
|
"""搜索日志"""
|
|
logs = log_manager.search_logs(keyword, limit)
|
|
return {
|
|
"logs": logs,
|
|
"count": len(logs)
|
|
}
|
|
|
|
|
|
@app.get("/api/logs/download")
|
|
async def download_logs():
|
|
"""下载日志文件"""
|
|
text = log_manager.export_to_text()
|
|
|
|
return StreamingResponse(
|
|
iter([text]),
|
|
media_type="text/plain",
|
|
headers={
|
|
"Content-Disposition": f"attachment; filename=xiaohongshu_logs_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
|
|
}
|
|
)
|
|
|
|
|
|
@app.get("/api/users")
|
|
async def get_users():
|
|
"""获取用户列表"""
|
|
users_file = os.path.join(config.DATA_DIR, "users.txt")
|
|
|
|
if not os.path.exists(users_file):
|
|
return {"users": []}
|
|
|
|
with open(users_file, 'r', encoding='utf-8') as f:
|
|
users = [{"id": line.strip()} for line in f if line.strip()]
|
|
|
|
# 添加采集状态
|
|
for user in users:
|
|
user_id = user["id"]
|
|
if user_id in app_state.users_status:
|
|
user.update(app_state.users_status[user_id])
|
|
else:
|
|
user["status"] = "pending"
|
|
user["note_count"] = 0
|
|
|
|
return {"users": users}
|
|
|
|
|
|
@app.post("/api/users")
|
|
async def add_user(request: dict):
|
|
"""添加用户"""
|
|
user_id = request.get("user_id", "").strip()
|
|
|
|
if not user_id:
|
|
raise HTTPException(status_code=400, detail="用户ID不能为空")
|
|
|
|
# 验证用户ID格式
|
|
if not user_id or len(user_id) < 10:
|
|
raise HTTPException(status_code=400, detail="用户ID格式不正确")
|
|
|
|
users_file = os.path.join(config.DATA_DIR, "users.txt")
|
|
|
|
# 检查是否已存在
|
|
if os.path.exists(users_file):
|
|
with open(users_file, 'r', encoding='utf-8') as f:
|
|
existing_users = [line.strip() for line in f if line.strip()]
|
|
if user_id in existing_users:
|
|
raise HTTPException(status_code=400, detail="用户ID已存在")
|
|
|
|
# 添加到文件
|
|
with open(users_file, 'a', encoding='utf-8') as f:
|
|
f.write(f"{user_id}\n")
|
|
|
|
return {
|
|
"success": True,
|
|
"message": f"用户 {user_id} 已添加",
|
|
"user_id": user_id
|
|
}
|
|
|
|
|
|
@app.delete("/api/users/{user_id}")
|
|
async def remove_user(user_id: str):
|
|
"""删除用户"""
|
|
users_file = os.path.join(config.DATA_DIR, "users.txt")
|
|
|
|
if not os.path.exists(users_file):
|
|
raise HTTPException(status_code=404, detail="用户文件不存在")
|
|
|
|
# 读取所有用户
|
|
with open(users_file, 'r', encoding='utf-8') as f:
|
|
users = [line.strip() for line in f if line.strip()]
|
|
|
|
# 检查用户是否存在
|
|
if user_id not in users:
|
|
raise HTTPException(status_code=404, detail="用户不存在")
|
|
|
|
# 删除用户
|
|
users.remove(user_id)
|
|
|
|
# 写回文件
|
|
with open(users_file, 'w', encoding='utf-8') as f:
|
|
for user in users:
|
|
f.write(f"{user}\n")
|
|
|
|
# 删除采集数据文件
|
|
data_file = os.path.join(config.DATA_DIR, f"{user_id}.json")
|
|
if os.path.exists(data_file):
|
|
os.remove(data_file)
|
|
|
|
return {
|
|
"success": True,
|
|
"message": f"用户 {user_id} 已删除"
|
|
}
|
|
|
|
|
|
@app.delete("/api/users")
|
|
async def clear_all_users():
|
|
"""清空所有用户"""
|
|
users_file = os.path.join(config.DATA_DIR, "users.txt")
|
|
|
|
if not os.path.exists(users_file):
|
|
raise HTTPException(status_code=404, detail="用户文件不存在")
|
|
|
|
# 清空文件
|
|
with open(users_file, 'w', encoding='utf-8') as f:
|
|
f.write("")
|
|
|
|
# 删除所有采集数据
|
|
data_dir = config.DATA_DIR
|
|
for filename in os.listdir(data_dir):
|
|
if filename.endswith('.json') and filename != 'auth_state.json':
|
|
file_path = os.path.join(data_dir, filename)
|
|
os.remove(file_path)
|
|
|
|
# 重置状态
|
|
app_state.users_status = {}
|
|
|
|
return {
|
|
"success": True,
|
|
"message": "所有用户已清空"
|
|
}
|
|
|
|
|
|
@app.post("/api/users/import/txt")
|
|
async def import_users_from_txt():
|
|
"""从users.txt导入用户（模拟）"""
|
|
users_file = os.path.join(config.DATA_DIR, "users.txt")
|
|
|
|
if not os.path.exists(users_file):
|
|
raise HTTPException(status_code=404, detail="users.txt文件不存在")
|
|
|
|
with open(users_file, 'r', encoding='utf-8') as f:
|
|
user_ids = [line.strip() for line in f if line.strip()]
|
|
|
|
return {
|
|
"success": True,
|
|
"count": len(user_ids),
|
|
"users": user_ids
|
|
}
|
|
|
|
|
|
@app.get("/api/users/export/txt")
|
|
async def export_users_to_txt():
|
|
"""导出用户列表到txt文件"""
|
|
users_file = os.path.join(config.DATA_DIR, "users.txt")
|
|
|
|
if not os.path.exists(users_file):
|
|
raise HTTPException(status_code=404, detail="用户文件不存在")
|
|
|
|
with open(users_file, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
|
|
return StreamingResponse(
|
|
iter([content]),
|
|
media_type="text/plain",
|
|
headers={
|
|
"Content-Disposition": f"attachment; filename=users_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
|
|
}
|
|
)
|
|
|
|
|
|
# ========== 登录相关 API ==========
|
|
|
|
class LoginStatusResponse(BaseModel):
|
|
logged_in: bool
|
|
login_time: Optional[str] = None
|
|
expire_time: Optional[str] = None
|
|
|
|
|
|
@app.get("/api/login/status", response_model=LoginStatusResponse)
|
|
async def get_login_status():
|
|
"""获取登录状态"""
|
|
auth_file = os.path.join(config.DATA_DIR, "auth_state.json")
|
|
|
|
if not os.path.exists(auth_file):
|
|
return LoginStatusResponse(logged_in=False)
|
|
|
|
try:
|
|
with open(auth_file, 'r', encoding='utf-8') as f:
|
|
auth_data = json.load(f)
|
|
|
|
# 检查是否过期
|
|
login_time = auth_data.get("login_time")
|
|
expire_time = auth_data.get("expire_time")
|
|
current_time = datetime.now().isoformat()
|
|
|
|
if expire_time and current_time > expire_time:
|
|
# 登录已过期，清除状态
|
|
os.remove(auth_file)
|
|
return LoginStatusResponse(logged_in=False)
|
|
|
|
return LoginStatusResponse(
|
|
logged_in=auth_data.get("logged_in", False),
|
|
login_time=login_time,
|
|
expire_time=expire_time
|
|
)
|
|
except Exception:
|
|
return LoginStatusResponse(logged_in=False)
|
|
|
|
|
|
@app.post("/api/login/qrcode")
|
|
async def generate_qrcode():
|
|
"""重定向到小红书官方登录页面"""
|
|
# 返回小红书登录页面 URL
|
|
login_url = "https://www.xiaohongshu.com/login"
|
|
|
|
return {
|
|
"login_url": login_url,
|
|
"message": "请在打开的小红书页面完成扫码登录",
|
|
"instruction": "登录成功后，返回此页面点击'验证登录'按钮"
|
|
}
|
|
|
|
|
|
@app.get("/api/check-login-status")
|
|
async def check_login_status():
|
|
"""检查小红书登录状态（直接验证）"""
|
|
try:
|
|
# 如果浏览器可用，直接使用浏览器验证
|
|
if hasattr(app_state, 'browser') and app_state.browser:
|
|
try:
|
|
page = await app_state.browser.new_page()
|
|
await page.goto("https://www.xiaohongshu.com/", timeout=10000)
|
|
await asyncio.sleep(2)
|
|
|
|
# 获取页面cookies
|
|
cookies = await page.context.cookies()
|
|
page.close()
|
|
|
|
# 检查是否有小红书相关的session cookie
|
|
web_session = None
|
|
for cookie in cookies:
|
|
if 'web_session' in cookie.get('name', '').lower():
|
|
web_session = cookie
|
|
break
|
|
|
|
if web_session and web_session.get('value'):
|
|
return {
|
|
"success": True,
|
|
"message": "已登录小红书",
|
|
"logged_in": True
|
|
}
|
|
else:
|
|
return {
|
|
"success": False,
|
|
"message": "未找到登录状态，请先登录"
|
|
}
|
|
|
|
except Exception as e:
|
|
pass
|
|
|
|
# 如果浏览器不可用，尝试访问登录页面检查
|
|
try:
|
|
# 使用 requests 作为备选方案
|
|
import asyncio
|
|
import aiohttp
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.get("https://www.xiaohongshu.com/") as response:
|
|
if response.status == 200:
|
|
# 检查响应中是否有登录相关内容
|
|
content = await response.text()
|
|
|
|
# 如果页面包含登录标识，说明未登录
|
|
if any(keyword in content for keyword in ['登录', 'login', 'signin']):
|
|
return {
|
|
"success": False,
|
|
"message": "未找到登录状态，请先登录"
|
|
}
|
|
|
|
# 尝试获取cookies
|
|
cookies = session.cookie_jar
|
|
if cookies:
|
|
return {
|
|
"success": True,
|
|
"message": "已登录小红书",
|
|
"logged_in": True
|
|
}
|
|
|
|
except Exception as e:
|
|
# 如果 aiohttp 不可用，跳过此检查
|
|
pass
|
|
|
|
return {
|
|
"success": False,
|
|
"message": "未找到登录状态，请先登录"
|
|
}
|
|
|
|
except Exception as e:
|
|
return {
|
|
"success": False,
|
|
"message": f"检查登录状态失败: {str(e)}"
|
|
}
|
|
|
|
|
|
@app.post("/api/login/verify")
|
|
async def verify_login():
|
|
"""验证登录状态"""
|
|
try:
|
|
# 优先检查浏览器的实际登录状态
|
|
if hasattr(app_state, 'browser') and app_state.browser:
|
|
try:
|
|
# 创建新页面检查登录状态
|
|
page = await app_state.browser.new_page()
|
|
|
|
# 访问小红书主页
|
|
await page.goto("https://www.xiaohongshu.com/", timeout=10000)
|
|
|
|
# 等待页面加载
|
|
await asyncio.sleep(2)
|
|
|
|
# 检查页面标题
|
|
title = await page.title()
|
|
|
|
# 如果页面包含"登录"字样，说明未登录
|
|
if "登录" in title or "login" in title.lower():
|
|
await page.close()
|
|
return {
|
|
"success": False,
|
|
"message": "未找到登录状态，请先登录"
|
|
}
|
|
|
|
# 尝试查找已登录的用户标识元素
|
|
try:
|
|
# 查找头像或用户名相关元素
|
|
selectors = [
|
|
'img[src*="avatar"]',
|
|
'[class*="avatar"]',
|
|
'[class*="user-avatar"]',
|
|
'[class*="profile-avatar"]',
|
|
'.user-name',
|
|
'[class*="username"]'
|
|
]
|
|
|
|
for selector in selectors:
|
|
elements = await page.query_selector_all(selector)
|
|
if elements:
|
|
await page.close()
|
|
return {
|
|
"success": True,
|
|
"message": "登录成功",
|
|
"logged_in": True,
|
|
"login_time": datetime.now().isoformat(),
|
|
"expire_time": (datetime.now() + timedelta(hours=24)).isoformat(),
|
|
"expire_in_hours": 24
|
|
}
|
|
except:
|
|
pass
|
|
|
|
await page.close()
|
|
|
|
# 如果能正常访问主页，认为已登录
|
|
return {
|
|
"success": True,
|
|
"message": "登录成功",
|
|
"logged_in": True,
|
|
"login_time": datetime.now().isoformat(),
|
|
"expire_time": (datetime.now() + timedelta(hours=24)).isoformat(),
|
|
"expire_in_hours": 24
|
|
}
|
|
|
|
except Exception as e:
|
|
# 如果浏览器访问失败，回退到本地文件检查
|
|
pass
|
|
|
|
# 检查本地登录状态文件（兼容处理）
|
|
auth_file = os.path.join(config.DATA_DIR, "auth_state.json")
|
|
|
|
if os.path.exists(auth_file):
|
|
with open(auth_file, 'r', encoding='utf-8') as f:
|
|
auth_data = json.load(f)
|
|
|
|
# 检查登录是否过期
|
|
login_time = auth_data.get("login_time")
|
|
expire_time = auth_data.get("expire_time")
|
|
current_time = datetime.now().isoformat()
|
|
|
|
if expire_time and current_time > expire_time:
|
|
# 登录已过期，清除状态
|
|
os.remove(auth_file)
|
|
return {
|
|
"success": False,
|
|
"message": "登录已过期，请重新登录"
|
|
}
|
|
|
|
return {
|
|
"success": True,
|
|
"message": "登录成功",
|
|
"logged_in": True,
|
|
"login_time": login_time,
|
|
"expire_time": expire_time,
|
|
"expire_in_hours": 24
|
|
}
|
|
|
|
return {
|
|
"success": False,
|
|
"message": "未找到登录状态，请先登录"
|
|
}
|
|
|
|
except Exception as e:
|
|
return {
|
|
"success": False,
|
|
"message": f"验证失败: {str(e)}"
|
|
}
|
|
|
|
|
|
@app.post("/api/logout")
|
|
async def logout():
|
|
"""退出登录"""
|
|
auth_file = os.path.join(config.DATA_DIR, "auth_state.json")
|
|
|
|
try:
|
|
if os.path.exists(auth_file):
|
|
os.remove(auth_file)
|
|
|
|
return {
|
|
"success": True,
|
|
"message": "已退出登录"
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"退出登录失败: {str(e)}")
|
|
|
|
|
|
@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="用户数据不存在")
|
|
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
data = json.load(f)
|
|
|
|
return data
|
|
|
|
|
|
@app.get("/api/events")
|
|
async def event_stream():
|
|
"""SSE 事件流"""
|
|
async def event_generator():
|
|
while True:
|
|
try:
|
|
# 等待事件，超时 1 秒发送心跳
|
|
event = await asyncio.wait_for(
|
|
app_state.event_queue.get(),
|
|
timeout=1.0
|
|
)
|
|
|
|
# 更新状态
|
|
if event["type"] == "log":
|
|
# 日志不需要更新状态
|
|
pass
|
|
elif event["type"] == "progress":
|
|
app_state.completed_users = event.get("completed", app_state.completed_users)
|
|
app_state.failed_users = event.get("failed", app_state.failed_users)
|
|
app_state.active_threads = event.get("active_threads", app_state.active_threads)
|
|
elif event["type"] == "user_status":
|
|
if event["user_id"] in app_state.users_status:
|
|
app_state.users_status[event["user_id"]]["status"] = event["status"]
|
|
app_state.users_status[event["user_id"]]["note_count"] = event.get("note_count", 0)
|
|
app_state.users_status[event["user_id"]]["thread_id"] = event.get("thread_id")
|
|
elif event["type"] == "thread_status":
|
|
app_state.thread_status[event["thread_id"]] = {
|
|
"status": event["status"],
|
|
"current_user": event.get("current_user")
|
|
}
|
|
# 更新活跃线程数
|
|
app_state.active_threads = sum(
|
|
1 for s in app_state.thread_status.values() if s.get("status") == "busy"
|
|
)
|
|
|
|
# 发送事件
|
|
yield f"data: {json.dumps(event)}\n\n"
|
|
|
|
except asyncio.TimeoutError:
|
|
# 发送心跳
|
|
yield ": heartbeat\n\n"
|
|
|
|
except Exception as e:
|
|
# 错误事件
|
|
yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
|
|
|
|
return StreamingResponse(
|
|
event_generator(),
|
|
media_type="text/event-stream",
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"Access-Control-Allow-Origin": "*"
|
|
}
|
|
)
|
|
|
|
|
|
# ========== Git 相关 API ==========
|
|
|
|
@app.get("/api/git/status")
|
|
async def git_status():
|
|
"""获取 Git 状态"""
|
|
if not os.path.exists("/app/.git"):
|
|
return {
|
|
"enabled": False,
|
|
"message": "不是 Git 仓库"
|
|
}
|
|
|
|
try:
|
|
# 获取当前分支
|
|
result = subprocess.run(
|
|
["git", "branch", "--show-current"],
|
|
capture_output=True,
|
|
text=True,
|
|
cwd="/app"
|
|
)
|
|
current_branch = result.stdout.strip()
|
|
|
|
# 获取最新提交
|
|
result = subprocess.run(
|
|
["git", "log", "-1", "--format=%h %s %cr"],
|
|
capture_output=True,
|
|
text=True,
|
|
cwd="/app"
|
|
)
|
|
latest_commit = result.stdout.strip()
|
|
|
|
# 检查是否有未提交的更改
|
|
result = subprocess.run(
|
|
["git", "status", "--porcelain"],
|
|
capture_output=True,
|
|
text=True,
|
|
cwd="/app"
|
|
)
|
|
has_changes = len(result.stdout.strip()) > 0
|
|
|
|
return {
|
|
"enabled": True,
|
|
"branch": current_branch,
|
|
"latest_commit": latest_commit,
|
|
"has_changes": has_changes,
|
|
"auto_pull_enabled": config.ENABLE_GIT_PULL,
|
|
"auto_pull_interval": config.AUTO_PULL_INTERVAL
|
|
}
|
|
|
|
except Exception as e:
|
|
return {
|
|
"enabled": False,
|
|
"error": str(e)
|
|
}
|
|
|
|
|
|
@app.post("/api/git/pull")
|
|
async def git_pull():
|
|
"""手动拉取最新代码"""
|
|
if not os.path.exists("/app/.git"):
|
|
raise HTTPException(status_code=400, detail="不是 Git 仓库")
|
|
|
|
try:
|
|
# 获取当前分支
|
|
result = subprocess.run(
|
|
["git", "branch", "--show-current"],
|
|
capture_output=True,
|
|
text=True,
|
|
cwd="/app"
|
|
)
|
|
current_branch = result.stdout.strip()
|
|
|
|
# 拉取最新代码
|
|
result = subprocess.run(
|
|
["git", "fetch", "origin", current_branch],
|
|
capture_output=True,
|
|
text=True,
|
|
cwd="/app"
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
raise HTTPException(status_code=500, detail=f"拉取失败: {result.stderr}")
|
|
|
|
# 检查是否有更新
|
|
result = subprocess.run(
|
|
["git", "rev-list", "HEAD..origin/" + current_branch, "--count"],
|
|
capture_output=True,
|
|
text=True,
|
|
cwd="/app"
|
|
)
|
|
commits_behind = int(result.stdout.strip())
|
|
|
|
if commits_behind == 0:
|
|
return {
|
|
"success": True,
|
|
"message": "已是最新代码",
|
|
"needs_restart": False
|
|
}
|
|
|
|
# 重置到最新代码
|
|
result = subprocess.run(
|
|
["git", "reset", "--hard", "origin/" + current_branch],
|
|
capture_output=True,
|
|
text=True,
|
|
cwd="/app"
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
raise HTTPException(status_code=500, detail=f"重置失败: {result.stderr}")
|
|
|
|
# 检查 Python 文件是否有变化
|
|
result = subprocess.run(
|
|
["git", "diff", "--name-only", "HEAD@{1}", "HEAD"],
|
|
capture_output=True,
|
|
text=True,
|
|
cwd="/app"
|
|
)
|
|
changed_files = result.stdout.strip().split('\n')
|
|
python_changed = any(f.endswith('.py') for f in changed_files if f)
|
|
|
|
return {
|
|
"success": True,
|
|
"message": f"已拉取 {commits_behind} 个新提交",
|
|
"changed_files": changed_files,
|
|
"needs_restart": python_changed,
|
|
"restart_hint": "Python 代码有变化，建议重启容器以应用更新"
|
|
}
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"拉取失败: {str(e)}")
|
|
|
|
|
|
@app.post("/api/git/restart")
|
|
async def restart_container():
|
|
"""
|
|
请求重启容器
|
|
注意：这需要在容器外配置重启策略
|
|
"""
|
|
return {
|
|
"success": True,
|
|
"message": "重启请求已接收",
|
|
"methods": [
|
|
"docker restart xiaohongshu-scraper",
|
|
"docker-compose restart"
|
|
]
|
|
}
|
|
|
|
|
|
# ========== 启动事件 ==========
|
|
@app.on_event("startup")
|
|
async def startup_event():
|
|
"""启动时初始化"""
|
|
# 确保数据目录存在
|
|
os.makedirs(config.DATA_DIR, exist_ok=True)
|
|
os.makedirs(config.SCREENSHOT_DIR, exist_ok=True)
|
|
|
|
# 初始化事件队列
|
|
app_state.event_queue = asyncio.Queue()
|
|
|
|
print("=" * 60)
|
|
print("小红书数据采集器 Web 服务启动")
|
|
print(f"访问地址: http://{config.WEB_HOST}:{config.WEB_PORT}")
|
|
print("=" * 60)
|
|
|
|
|
|
@app.on_event("shutdown")
|
|
async def shutdown_event():
|
|
"""关闭时清理"""
|
|
if app_state.scraper:
|
|
app_state.scraper.stop()
|
|
|
|
if app_state.scraper_task:
|
|
app_state.scraper_task.cancel()
|
|
|
|
print("Web 服务已关闭")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(
|
|
"app.web:app",
|
|
host=config.WEB_HOST,
|
|
port=config.WEB_PORT,
|
|
reload=True
|
|
)
|