""" 注册任务 API 路由 """ import asyncio import logging import uuid import random from datetime import datetime from typing import List, Optional, Dict from fastapi import APIRouter, HTTPException, Query, BackgroundTasks from pydantic import BaseModel, Field from ...database import crud from ...database.session import get_db from ...database.models import RegistrationTask from ...core.register import RegistrationEngine, RegistrationResult from ...services import EmailServiceFactory, EmailServiceType from ...config.settings import get_settings logger = logging.getLogger(__name__) router = APIRouter() # 任务存储(简单的内存存储,生产环境应使用 Redis) running_tasks: dict = {} # 批量任务存储 batch_tasks: Dict[str, dict] = {} # ============== Pydantic Models ============== class RegistrationTaskCreate(BaseModel): """创建注册任务请求""" email_service_type: str = "tempmail" proxy: Optional[str] = None email_service_config: Optional[dict] = None class BatchRegistrationRequest(BaseModel): """批量注册请求""" count: int = 1 # 注册数量 email_service_type: str = "tempmail" proxy: Optional[str] = None email_service_config: Optional[dict] = None interval_min: int = 5 # 最小间隔秒数 interval_max: int = 30 # 最大间隔秒数 class BatchRegistrationResponse(BaseModel): """批量注册响应""" batch_id: str count: int tasks: List[RegistrationTaskResponse] class RegistrationTaskResponse(BaseModel): """注册任务响应""" id: int task_uuid: str status: str email_service_id: Optional[int] = None proxy: Optional[str] = None logs: Optional[str] = None result: Optional[dict] = None error_message: Optional[str] = None created_at: Optional[str] = None started_at: Optional[str] = None completed_at: Optional[str] = None class Config: from_attributes = True class TaskListResponse(BaseModel): """任务列表响应""" total: int tasks: List[RegistrationTaskResponse] # ============== Helper Functions ============== def task_to_response(task: RegistrationTask) -> RegistrationTaskResponse: """转换任务模型为响应""" return RegistrationTaskResponse( id=task.id, task_uuid=task.task_uuid, status=task.status, email_service_id=task.email_service_id, proxy=task.proxy, logs=task.logs, result=task.result, error_message=task.error_message, created_at=task.created_at.isoformat() if task.created_at else None, started_at=task.started_at.isoformat() if task.started_at else None, completed_at=task.completed_at.isoformat() if task.completed_at else None, ) async def run_registration_task(task_uuid: str, email_service_type: str, proxy: Optional[str], email_service_config: Optional[dict]): """异步执行注册任务""" with get_db() as db: try: # 更新任务状态为运行中 task = crud.update_registration_task( db, task_uuid, status="running", started_at=datetime.utcnow() ) if not task: logger.error(f"任务不存在: {task_uuid}") return # 创建邮箱服务 service_type = EmailServiceType(email_service_type) settings = get_settings() if service_type == EmailServiceType.TEMPMAIL: config = { "base_url": settings.tempmail_base_url, "timeout": settings.tempmail_timeout, "max_retries": settings.tempmail_max_retries, "proxy_url": proxy, } elif service_type == EmailServiceType.CUSTOM_DOMAIN: config = { "base_url": settings.custom_domain_base_url, "api_key": settings.custom_domain_api_key.get_secret_value() if settings.custom_domain_api_key else "", "proxy_url": proxy, } else: config = email_service_config or {} email_service = EmailServiceFactory.create(service_type, config) # 创建注册引擎 def log_callback(msg): with get_db() as db_inner: crud.append_task_log(db_inner, task_uuid, msg) engine = RegistrationEngine( email_service=email_service, proxy_url=proxy, callback_logger=log_callback, task_uuid=task_uuid ) # 执行注册 result = engine.run() if result.success: # 保存到数据库 engine.save_to_database(result) # 更新任务状态 crud.update_registration_task( db, task_uuid, status="completed", completed_at=datetime.utcnow(), result=result.to_dict() ) logger.info(f"注册任务完成: {task_uuid}, 邮箱: {result.email}") else: # 更新任务状态为失败 crud.update_registration_task( db, task_uuid, status="failed", completed_at=datetime.utcnow(), error_message=result.error_message ) logger.warning(f"注册任务失败: {task_uuid}, 原因: {result.error_message}") except Exception as e: logger.error(f"注册任务异常: {task_uuid}, 错误: {e}") try: with get_db() as db: crud.update_registration_task( db, task_uuid, status="failed", completed_at=datetime.utcnow(), error_message=str(e) ) except: pass async def run_batch_registration( batch_id: str, task_uuids: List[str], email_service_type: str, proxy: Optional[str], email_service_config: Optional[dict], interval_min: int, interval_max: int ): """异步执行批量注册任务""" batch_tasks[batch_id] = { "total": len(task_uuids), "completed": 0, "success": 0, "failed": 0, "cancelled": False, "task_uuids": task_uuids, "current_index": 0 } try: for i, task_uuid in enumerate(task_uuids): # 检查是否已取消 if batch_tasks[batch_id]["cancelled"]: # 取消剩余任务 with get_db() as db: for remaining_uuid in task_uuids[i:]: crud.update_registration_task(db, remaining_uuid, status="cancelled") logger.info(f"批量任务 {batch_id} 已取消") break batch_tasks[batch_id]["current_index"] = i # 运行单个注册任务 await run_registration_task( task_uuid, email_service_type, proxy, email_service_config ) # 更新统计 with get_db() as db: task = crud.get_registration_task(db, task_uuid) if task: batch_tasks[batch_id]["completed"] += 1 if task.status == "completed": batch_tasks[batch_id]["success"] += 1 elif task.status == "failed": batch_tasks[batch_id]["failed"] += 1 # 如果不是最后一个任务,等待随机间隔 if i < len(task_uuids) - 1 and not batch_tasks[batch_id]["cancelled"]: wait_time = random.randint(interval_min, interval_max) logger.info(f"批量任务 {batch_id}: 等待 {wait_time} 秒后继续下一个任务") await asyncio.sleep(wait_time) logger.info(f"批量任务 {batch_id} 完成: 成功 {batch_tasks[batch_id]['success']}, 失败 {batch_tasks[batch_id]['failed']}") except Exception as e: logger.error(f"批量任务 {batch_id} 异常: {e}") finally: batch_tasks[batch_id]["finished"] = True # ============== API Endpoints ============== @router.post("/start", response_model=RegistrationTaskResponse) async def start_registration( request: RegistrationTaskCreate, background_tasks: BackgroundTasks ): """ 启动注册任务 - email_service_type: 邮箱服务类型 (tempmail, outlook, custom_domain) - proxy: 代理地址 - email_service_config: 邮箱服务配置(outlook 需要提供账户信息) """ # 验证邮箱服务类型 try: EmailServiceType(request.email_service_type) except ValueError: raise HTTPException( status_code=400, detail=f"无效的邮箱服务类型: {request.email_service_type}" ) # 创建任务 task_uuid = str(uuid.uuid4()) with get_db() as db: task = crud.create_registration_task( db, task_uuid=task_uuid, proxy=request.proxy ) # 在后台运行注册任务 background_tasks.add_task( run_registration_task, task_uuid, request.email_service_type, request.proxy, request.email_service_config ) return task_to_response(task) @router.post("/batch", response_model=BatchRegistrationResponse) async def start_batch_registration( request: BatchRegistrationRequest, background_tasks: BackgroundTasks ): """ 启动批量注册任务 - count: 注册数量 (1-100) - email_service_type: 邮箱服务类型 - proxy: 代理地址 - interval_min: 最小间隔秒数 - interval_max: 最大间隔秒数 """ # 验证参数 if request.count < 1 or request.count > 100: raise HTTPException(status_code=400, detail="注册数量必须在 1-100 之间") try: EmailServiceType(request.email_service_type) except ValueError: raise HTTPException( status_code=400, detail=f"无效的邮箱服务类型: {request.email_service_type}" ) if request.interval_min < 0 or request.interval_max < request.interval_min: raise HTTPException(status_code=400, detail="间隔时间参数无效") # 创建批量任务 batch_id = str(uuid.uuid4()) task_uuids = [] with get_db() as db: for _ in range(request.count): task_uuid = str(uuid.uuid4()) task = crud.create_registration_task( db, task_uuid=task_uuid, proxy=request.proxy ) task_uuids.append(task_uuid) # 获取所有任务 with get_db() as db: tasks = [crud.get_registration_task(db, uuid) for uuid in task_uuids] # 在后台运行批量注册 background_tasks.add_task( run_batch_registration, batch_id, task_uuids, request.email_service_type, request.proxy, request.email_service_config, request.interval_min, request.interval_max ) return BatchRegistrationResponse( batch_id=batch_id, count=request.count, tasks=[task_to_response(t) for t in tasks if t] ) @router.get("/batch/{batch_id}") async def get_batch_status(batch_id: str): """获取批量任务状态""" if batch_id not in batch_tasks: raise HTTPException(status_code=404, detail="批量任务不存在") batch = batch_tasks[batch_id] return { "batch_id": batch_id, "total": batch["total"], "completed": batch["completed"], "success": batch["success"], "failed": batch["failed"], "current_index": batch["current_index"], "cancelled": batch["cancelled"], "finished": batch.get("finished", False), "progress": f"{batch['completed']}/{batch['total']}" } @router.post("/batch/{batch_id}/cancel") async def cancel_batch(batch_id: str): """取消批量任务""" if batch_id not in batch_tasks: raise HTTPException(status_code=404, detail="批量任务不存在") batch = batch_tasks[batch_id] if batch.get("finished"): raise HTTPException(status_code=400, detail="批量任务已完成") batch["cancelled"] = True return {"success": True, "message": "批量任务取消请求已提交"} @router.get("/tasks", response_model=TaskListResponse) async def list_tasks( page: int = Query(1, ge=1), page_size: int = Query(20, ge=1, le=100), status: Optional[str] = Query(None), ): """获取任务列表""" with get_db() as db: query = db.query(RegistrationTask) if status: query = query.filter(RegistrationTask.status == status) total = query.count() offset = (page - 1) * page_size tasks = query.order_by(RegistrationTask.created_at.desc()).offset(offset).limit(page_size).all() return TaskListResponse( total=total, tasks=[task_to_response(t) for t in tasks] ) @router.get("/tasks/{task_uuid}", response_model=RegistrationTaskResponse) async def get_task(task_uuid: str): """获取任务详情""" with get_db() as db: task = crud.get_registration_task(db, task_uuid) if not task: raise HTTPException(status_code=404, detail="任务不存在") return task_to_response(task) @router.get("/tasks/{task_uuid}/logs") async def get_task_logs(task_uuid: str): """获取任务日志""" with get_db() as db: task = crud.get_registration_task(db, task_uuid) if not task: raise HTTPException(status_code=404, detail="任务不存在") logs = task.logs or "" return { "task_uuid": task_uuid, "status": task.status, "logs": logs.split("\n") if logs else [] } @router.post("/tasks/{task_uuid}/cancel") async def cancel_task(task_uuid: str): """取消任务""" with get_db() as db: task = crud.get_registration_task(db, task_uuid) if not task: raise HTTPException(status_code=404, detail="任务不存在") if task.status not in ["pending", "running"]: raise HTTPException(status_code=400, detail="任务已完成或已取消") task = crud.update_registration_task(db, task_uuid, status="cancelled") return {"success": True, "message": "任务已取消"} @router.delete("/tasks/{task_uuid}") async def delete_task(task_uuid: str): """删除任务""" with get_db() as db: task = crud.get_registration_task(db, task_uuid) if not task: raise HTTPException(status_code=404, detail="任务不存在") if task.status == "running": raise HTTPException(status_code=400, detail="无法删除运行中的任务") crud.delete_registration_task(db, task_uuid) return {"success": True, "message": "任务已删除"} @router.get("/stats") async def get_registration_stats(): """获取注册统计信息""" with get_db() as db: from sqlalchemy import func # 按状态统计 status_stats = db.query( RegistrationTask.status, func.count(RegistrationTask.id) ).group_by(RegistrationTask.status).all() # 今日注册数 today = datetime.utcnow().date() today_count = db.query(func.count(RegistrationTask.id)).filter( func.date(RegistrationTask.created_at) == today ).scalar() return { "by_status": {status: count for status, count in status_stats}, "today_count": today_count }