mirror of
https://github.com/cnlimiter/codex-register.git
synced 2026-05-12 02:21:07 +08:00
501 lines
16 KiB
Python
501 lines
16 KiB
Python
"""
|
||
注册任务 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
|
||
}
|