import asyncio import hashlib import json from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any from uuid import UUID import psycopg from fastapi import Depends, FastAPI, HTTPException, Request, Response, status import httpx from fastapi.responses import FileResponse, JSONResponse from psycopg import Connection from psycopg.types.json import Json from .config import settings from .db import get_connection from .schemas import ( AuthResponse, LoginRequest, MembershipRequest, PasswordChangeRequest, QuotaAdjustmentRequest, RegisterRequest, StatusRequest, UserPublic, VerificationConfirmRequest, VerificationSendRequest, TtsTaskRequest, TtsTaskPublic, TtsVoicePublic, normalize_email, ) from .security import ( check_origin, clear_session_cookie, code_digest, current_user, hash_password, new_csrf_token, new_token, password_hasher, rate_limiter, require_admin, require_csrf, set_csrf_cookie, set_session_cookie, token_digest, utc_now, verify_password, ) app = FastAPI(title="Kaotings Business API", version="0.2.0", docs_url=None if settings.app_env == "production" else "/docs") def error(code: str, message: str, http_status: int) -> HTTPException: return HTTPException(status_code=http_status, detail={"code": code, "message": message}) def period_bounds(now: datetime | None = None) -> tuple[datetime, datetime]: current = now or utc_now() start = current.replace(day=1, hour=0, minute=0, second=0, microsecond=0) if start.month == 12: end = start.replace(year=start.year + 1, month=1) else: end = start.replace(month=start.month + 1) return start, end def effective_plan(connection: Connection, user: dict) -> str: if user["role"] == "admin": return user["plan"] grant = connection.execute( """ SELECT plan FROM membership_grants WHERE user_id = %s AND revoked_at IS NULL AND starts_at <= now() AND (expires_at IS NULL OR expires_at > now()) ORDER BY starts_at DESC LIMIT 1 """, (user["id"],), ).fetchone() return grant["plan"] if grant else "free" def user_public(connection: Connection, user: dict) -> UserPublic: return UserPublic( id=user["id"], email=user["email"], phone=user["phone"], role=user["role"], plan=effective_plan(connection, user), status=user["status"], email_verified=user["email_verified"], phone_verified=user["phone_verified"], created_at=user["created_at"], last_login_at=user["last_login_at"], ) def ensure_quota(connection: Connection, user_id: UUID, plan: str) -> dict: start, end = period_bounds() policy = connection.execute("SELECT period_limit FROM plan_policies WHERE code = %s", (plan,)).fetchone() if not policy: raise error("POLICY_MISSING", "权益策略未配置", 500) connection.execute( """ INSERT INTO quota_accounts(user_id, period_start, period_end, limit_snapshot) VALUES (%s, %s, %s, %s) ON CONFLICT (user_id, period_start, period_end) DO NOTHING """, (user_id, start, end, policy["period_limit"]), ) row = connection.execute( "SELECT * FROM quota_accounts WHERE user_id = %s AND period_start = %s AND period_end = %s", (user_id, start, end), ).fetchone() return row def usage_view(quota: dict) -> dict[str, Any]: available = max(0, quota["limit_snapshot"] + quota["adjustment"] - quota["used"] - quota["reserved"]) return { "period_start": quota["period_start"], "period_end": quota["period_end"], "limit": quota["limit_snapshot"], "adjustment": quota["adjustment"], "used": quota["used"], "reserved": quota["reserved"], "available": available, } def task_view(connection: Connection, task: dict) -> dict[str, Any]: audio = connection.execute("SELECT expires_at, status FROM audio_files WHERE task_id = %s", (task["id"],)).fetchone() return { "id": task["id"], "status": task["status"], "text_length": task["text_length"], "voice_id": task["provider_voice_id"], "parameters": task["parameters"], "error_code": task["error_code"], "audio_available": bool(audio and audio["status"] == "available"), "audio_expires_at": audio["expires_at"] if audio else None, "created_at": task["created_at"], "started_at": task["started_at"], "finished_at": task["finished_at"], } def claim_next_task() -> dict | None: with psycopg.connect(settings.database_url, row_factory=psycopg.rows.dict_row) as connection: task = connection.execute("SELECT * FROM tts_tasks WHERE status = 'queued' ORDER BY created_at FOR UPDATE SKIP LOCKED LIMIT 1").fetchone() if not task: return None updated = connection.execute("UPDATE tts_tasks SET status = 'running', started_at = now(), updated_at = now(), attempt_count = attempt_count + 1 WHERE id = %s RETURNING *", (task["id"],)).fetchone() connection.commit() return updated def finish_task_failure(task: dict, code: str) -> None: with psycopg.connect(settings.database_url, row_factory=psycopg.rows.dict_row) as connection: connection.execute("UPDATE tts_tasks SET status = 'failed', error_code = %s, finished_at = now(), updated_at = now() WHERE id = %s AND status = 'running'", (code, task["id"])) connection.execute("UPDATE quota_accounts SET reserved = GREATEST(0, reserved - %s), version = version + 1 WHERE id = %s", (task["reserved_amount"], task["quota_account_id"])) connection.execute("INSERT INTO usage_records(user_id, quota_account_id, task_id, type, amount) VALUES (%s, %s, %s, 'release', %s)", (task["user_id"], task["quota_account_id"], task["id"], task["reserved_amount"])) connection.commit() def finish_task_success(task: dict, audio: bytes, mime_type: str) -> None: storage_dir = Path(settings.audio_storage_dir) storage_dir.mkdir(parents=True, exist_ok=True) extension = "mp3" if mime_type == "audio/mpeg" else "wav" storage_key = f"{task['user_id']}/{task['id']}.{extension}" target = storage_dir / storage_key target.parent.mkdir(parents=True, exist_ok=True) target.write_bytes(audio) checksum = hashlib.sha256(audio).hexdigest() with psycopg.connect(settings.database_url, row_factory=psycopg.rows.dict_row) as connection: connection.execute("INSERT INTO audio_files(task_id, owner_id, storage_key, mime_type, size_bytes, checksum) VALUES (%s, %s, %s, %s, %s, %s)", (task["id"], task["user_id"], storage_key, mime_type, len(audio), checksum)) connection.execute("UPDATE tts_tasks SET status = 'succeeded', finished_at = now(), updated_at = now() WHERE id = %s AND status = 'running'", (task["id"],)) connection.execute("UPDATE quota_accounts SET reserved = GREATEST(0, reserved - %s), used = used + %s, version = version + 1 WHERE id = %s", (task["reserved_amount"], task["reserved_amount"], task["quota_account_id"])) connection.execute("INSERT INTO usage_records(user_id, quota_account_id, task_id, type, amount) VALUES (%s, %s, %s, 'consume', %s)", (task["user_id"], task["quota_account_id"], task["id"], task["reserved_amount"])) connection.commit() async def process_task(task: dict) -> None: if not settings.tts_upstream_url: await asyncio.to_thread(finish_task_failure, task, "UPSTREAM_NOT_CONFIGURED") return parameters = task["parameters"] response_format = parameters.get("format", "wav") payload = {"model": parameters.get("model", "qwen3-tts"), "input": task["text"], "voice": task["provider_voice_id"], "response_format": response_format, "speed": parameters.get("speed", 1.0)} headers = {"Content-Type": "application/json"} if settings.tts_api_key: headers["Authorization"] = f"Bearer {settings.tts_api_key}" try: async with httpx.AsyncClient(timeout=settings.tts_timeout_seconds) as client: response = await client.post(f"{settings.tts_upstream_url}/v1/audio/speech", headers=headers, json=payload) if response.status_code == 401: raise RuntimeError("UPSTREAM_UNAUTHORIZED") if response.status_code >= 400: raise RuntimeError(f"UPSTREAM_HTTP_{response.status_code}") audio = response.content if not audio or len(audio) > 50 * 1024 * 1024: raise RuntimeError("UPSTREAM_AUDIO_INVALID") mime = "audio/mpeg" if response_format == "mp3" else "audio/wav" await asyncio.to_thread(finish_task_success, task, audio, mime) except httpx.TimeoutException: await asyncio.to_thread(finish_task_failure, task, "UPSTREAM_TIMEOUT") except Exception as exc: await asyncio.to_thread(finish_task_failure, task, str(exc)[:80]) async def task_worker() -> None: while True: task = await asyncio.to_thread(claim_next_task) if task: await process_task(task) else: await asyncio.sleep(1) worker_task: asyncio.Task | None = None @app.on_event("startup") async def start_worker(): global worker_task worker_task = asyncio.create_task(task_worker()) @app.on_event("shutdown") async def stop_worker(): if worker_task: worker_task.cancel() @app.exception_handler(HTTPException) async def http_error_handler(_: Request, exc: HTTPException): detail = exc.detail if isinstance(exc.detail, dict) else {"code": "REQUEST_FAILED", "message": str(exc.detail)} return JSONResponse(status_code=exc.status_code, content={"error": detail}) @app.get("/healthz") def health(connection: Connection = Depends(get_connection)): connection.execute("SELECT 1") return {"status": "ok", "database": "ok", "service": "api"} @app.get("/api/v1/auth/csrf") def csrf(response: Response, request: Request): check_origin(request) token = request.cookies.get("kaotings_csrf") or new_csrf_token() set_csrf_cookie(response, token) return {"csrf_token": token} @app.post("/api/v1/auth/register", response_model=AuthResponse, status_code=201) def register(payload: RegisterRequest, request: Request, connection: Connection = Depends(get_connection)): check_origin(request) key = f"register:{request.client.host if request.client else 'unknown'}" if not rate_limiter.allow(key, 5, 3600): raise error("RATE_LIMITED", "请求过于频繁", 429) email = normalize_email(str(payload.email)) try: user = connection.execute( """ INSERT INTO users(email, password_hash) VALUES (%s, %s) RETURNING * """, (email, hash_password(payload.password)), ).fetchone() ensure_quota(connection, user["id"], "free") connection.commit() except psycopg.errors.UniqueViolation: connection.rollback() raise error("REGISTRATION_FAILED", "注册信息不可用", 409) return {"user": user_public(connection, user)} @app.post("/api/v1/auth/login", response_model=AuthResponse) def login(payload: LoginRequest, request: Request, response: Response, connection: Connection = Depends(get_connection)): check_origin(request) key = f"login:{request.client.host if request.client else 'unknown'}:{normalize_email(str(payload.email))}" if not rate_limiter.allow(key, 10, 300): raise error("RATE_LIMITED", "请求过于频繁", 429) user = connection.execute("SELECT * FROM users WHERE email = %s", (normalize_email(str(payload.email)),)).fetchone() if not user or not verify_password(payload.password, user["password_hash"]): raise error("LOGIN_FAILED", "邮箱或密码错误", 401) if user["status"] != "active": raise error("ACCOUNT_DISABLED", "账户不可用", 403) token = new_token() connection.execute("INSERT INTO sessions(user_id, token_hash, expires_at) VALUES (%s, %s, %s)", (user["id"], token_digest(token), utc_now() + timedelta(seconds=settings.session_ttl_seconds))) user = connection.execute("UPDATE users SET last_login_at = now(), updated_at = now() WHERE id = %s RETURNING *", (user["id"],)).fetchone() ensure_quota(connection, user["id"], effective_plan(connection, user)) connection.commit() set_session_cookie(response, token) return {"user": user_public(connection, user)} @app.post("/api/v1/auth/logout", status_code=204, dependencies=[Depends(require_csrf)]) def logout(response: Response, request: Request, connection: Connection = Depends(get_connection)): token = request.cookies.get(settings.session_cookie_name) if token: connection.execute("UPDATE sessions SET revoked_at = now() WHERE token_hash = %s", (token_digest(token),)) connection.commit() clear_session_cookie(response) return Response(status_code=204) @app.get("/api/v1/auth/me", response_model=UserPublic) def me(user: dict = Depends(current_user), connection: Connection = Depends(get_connection)): return user_public(connection, user) @app.post("/api/v1/auth/password/change", dependencies=[Depends(require_csrf)]) def change_password(payload: PasswordChangeRequest, request: Request, connection: Connection = Depends(get_connection), user: dict = Depends(current_user)): if not verify_password(payload.current_password, user["password_hash"]): raise error("PASSWORD_INVALID", "当前密码错误", 400) connection.execute("UPDATE users SET password_hash = %s, updated_at = now() WHERE id = %s", (hash_password(payload.new_password), user["id"])) connection.execute("UPDATE sessions SET revoked_at = now() WHERE user_id = %s AND id <> %s", (user["id"], user["session_id"])) connection.commit() return {"status": "ok"} @app.get("/api/v1/account/usage") def account_usage(user: dict = Depends(current_user), connection: Connection = Depends(get_connection)): quota = ensure_quota(connection, user["id"], effective_plan(connection, user)) connection.commit() return {"plan": effective_plan(connection, user), **usage_view(quota)} @app.post("/api/v1/auth/verification/send") def verification_send(payload: VerificationSendRequest, request: Request, user: dict = Depends(current_user)): check_origin(request) enabled = settings.email_verification_enabled if payload.channel == "email" else settings.phone_verification_enabled if not enabled: raise error("VERIFICATION_NOT_ENABLED", "验证渠道尚未启用", 503) raise error("VERIFICATION_PROVIDER_UNAVAILABLE", "验证渠道尚未配置", 503) @app.post("/api/v1/auth/verification/confirm", dependencies=[Depends(require_csrf)]) def verification_confirm(payload: VerificationConfirmRequest, request: Request, connection: Connection = Depends(get_connection), user: dict = Depends(current_user)): check_origin(request) challenge = connection.execute("SELECT * FROM verification_challenges WHERE id = %s AND user_id = %s", (payload.challenge_id, user["id"])).fetchone() if not challenge or challenge["consumed_at"] or challenge["expires_at"] <= utc_now() or challenge["attempt_count"] >= 5: raise error("VERIFICATION_INVALID", "验证挑战无效或已过期", 400) if not __import__("hmac").compare_digest(challenge["code_digest"], code_digest(payload.code)): connection.execute("UPDATE verification_challenges SET attempt_count = attempt_count + 1 WHERE id = %s", (payload.challenge_id,)) connection.commit() raise error("VERIFICATION_INVALID", "验证码错误", 400) connection.execute("UPDATE verification_challenges SET consumed_at = now() WHERE id = %s", (payload.challenge_id,)) if challenge["channel"] == "email": connection.execute("UPDATE users SET email_verified = true, updated_at = now() WHERE id = %s", (user["id"],)) else: connection.execute("UPDATE users SET phone_verified = true, updated_at = now() WHERE id = %s", (user["id"],)) connection.commit() return {"status": "ok"} @app.get("/api/v1/tts/voices", response_model=list[TtsVoicePublic]) def tts_voices(connection: Connection = Depends(get_connection)): return connection.execute("SELECT id, provider_voice_id, name, language, description, supported_parameters FROM tts_voices WHERE enabled = true ORDER BY name").fetchall() @app.post("/api/v1/tts/tasks", status_code=202) def create_tts_task(payload: TtsTaskRequest, request: Request, connection: Connection = Depends(get_connection), user: dict = Depends(current_user), _: None = Depends(require_csrf)): idempotency_key = request.headers.get("idempotency-key", "").strip() if len(idempotency_key) < 8 or len(idempotency_key) > 120: raise error("IDEMPOTENCY_REQUIRED", "需要有效的 Idempotency-Key", 400) plan = effective_plan(connection, user) policy = connection.execute("SELECT * FROM plan_policies WHERE code = %s", (plan,)).fetchone() text_length = len(payload.text) if not policy or text_length > policy["max_text_length"]: raise error("TEXT_TOO_LONG", "文本超过当前计划限制", 422) parameters = dict(payload.parameters) response_format = str(parameters.get("format", "wav")).lower() speed = float(parameters.get("speed", 1.0)) if response_format not in {"wav", "mp3"} or not 0.5 <= speed <= 2: raise error("INVALID_PARAMETERS", "音频格式或语速不可用", 422) voice = connection.execute("SELECT * FROM tts_voices WHERE provider_voice_id = %s AND enabled = true", (payload.voice_id,)).fetchone() if not voice or plan not in (voice["allowed_plans"] or ["free", "vip"]): raise error("VOICE_NOT_ALLOWED", "音色不可用", 422) request_hash = hashlib.sha256(json.dumps({"text": payload.text, "voice": payload.voice_id, "parameters": parameters}, ensure_ascii=False, sort_keys=True).encode()).hexdigest() existing = connection.execute("SELECT * FROM tts_tasks WHERE user_id = %s AND idempotency_key = %s", (user["id"], idempotency_key)).fetchone() if existing: if existing["request_hash"] != request_hash: raise error("IDEMPOTENCY_CONFLICT", "幂等键已用于其他请求", 409) return task_view(connection, existing) quota = ensure_quota(connection, user["id"], plan) quota = connection.execute("SELECT * FROM quota_accounts WHERE id = %s FOR UPDATE", (quota["id"],)).fetchone() available = quota["limit_snapshot"] + quota["adjustment"] - quota["used"] - quota["reserved"] if available < text_length: raise error("QUOTA_EXCEEDED", "额度不足", 409) task = connection.execute( """ INSERT INTO tts_tasks(user_id, text, text_length, voice_id, provider_voice_id, parameters, idempotency_key, request_hash, policy_version, quota_account_id, reserved_amount) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING * """, (user["id"], payload.text, text_length, voice["id"], payload.voice_id, Json(parameters), idempotency_key, request_hash, policy["version"], quota["id"], text_length), ).fetchone() connection.execute("UPDATE quota_accounts SET reserved = reserved + %s, version = version + 1 WHERE id = %s", (text_length, quota["id"])) connection.execute("INSERT INTO usage_records(user_id, quota_account_id, task_id, type, amount, idempotency_key) VALUES (%s, %s, %s, 'reserve', %s, %s)", (user["id"], quota["id"], task["id"], text_length, f"reserve:{task['id']}")) connection.commit() return task_view(connection, task) @app.get("/api/v1/tts/tasks", response_model=list[TtsTaskPublic]) def list_tts_tasks(user: dict = Depends(current_user), connection: Connection = Depends(get_connection)): tasks = connection.execute("SELECT * FROM tts_tasks WHERE user_id = %s ORDER BY created_at DESC LIMIT 100", (user["id"],)).fetchall() return [task_view(connection, task) for task in tasks] def owned_task(task_id: UUID, user: dict, connection: Connection) -> dict: task = connection.execute("SELECT * FROM tts_tasks WHERE id = %s", (task_id,)).fetchone() if not task or (task["user_id"] != user["id"] and user["role"] != "admin"): raise error("NOT_FOUND", "任务不存在", 404) return task @app.get("/api/v1/tts/tasks/{task_id}", response_model=TtsTaskPublic) def get_tts_task(task_id: UUID, user: dict = Depends(current_user), connection: Connection = Depends(get_connection)): return task_view(connection, owned_task(task_id, user, connection)) def audio_response(task_id: UUID, user: dict, connection: Connection, download: bool): task = owned_task(task_id, user, connection) audio = connection.execute("SELECT * FROM audio_files WHERE task_id = %s AND status = 'available'", (task_id,)).fetchone() if task["status"] != "succeeded" or not audio: raise error("AUDIO_NOT_AVAILABLE", "音频尚不可用", 404) path = Path(settings.audio_storage_dir) / audio["storage_key"] if not path.is_file(): raise error("AUDIO_NOT_AVAILABLE", "音频文件不可用", 404) filename = f"kaotings-{task_id}.{'mp3' if audio['mime_type'] == 'audio/mpeg' else 'wav'}" if download else None return FileResponse(path, media_type=audio["mime_type"], filename=filename) @app.get("/api/v1/tts/tasks/{task_id}/audio") def play_tts_audio(task_id: UUID, user: dict = Depends(current_user), connection: Connection = Depends(get_connection)): return audio_response(task_id, user, connection, False) @app.get("/api/v1/tts/tasks/{task_id}/download") def download_tts_audio(task_id: UUID, user: dict = Depends(current_user), connection: Connection = Depends(get_connection)): return audio_response(task_id, user, connection, True) @app.get("/api/v1/admin/users") def admin_users(user: dict = Depends(require_admin), connection: Connection = Depends(get_connection)): rows = connection.execute("SELECT id, email, phone, role, plan, status, email_verified, phone_verified, created_at, last_login_at FROM users ORDER BY created_at DESC LIMIT 100").fetchall() return {"items": rows} @app.patch("/api/v1/admin/users/{user_id}/status", dependencies=[Depends(require_csrf)]) def admin_status(user_id: UUID, payload: StatusRequest, connection: Connection = Depends(get_connection), actor: dict = Depends(require_admin)): target = connection.execute("SELECT * FROM users WHERE id = %s", (user_id,)).fetchone() if not target: raise error("NOT_FOUND", "用户不存在", 404) if target["role"] == "admin" and payload.status == "disabled": count = connection.execute("SELECT count(*) FROM users WHERE role = 'admin' AND status = 'active'").fetchone()["count"] if count <= 1: raise error("LAST_ADMIN_PROTECTED", "不能禁用最后一个可用管理员", 409) updated = connection.execute("UPDATE users SET status = %s, updated_at = now() WHERE id = %s RETURNING id, status", (payload.status, user_id)).fetchone() connection.execute("INSERT INTO admin_audit_logs(actor_id, action, target_type, target_id, before_value, after_value, reason) VALUES (%s, 'user_status', 'user', %s, %s, %s, %s)", (actor["id"], user_id, Json({"status": target["status"]}), Json({"status": updated["status"]}), payload.reason)) connection.commit() return updated @app.put("/api/v1/admin/users/{user_id}/membership", dependencies=[Depends(require_csrf)]) def admin_membership(user_id: UUID, payload: MembershipRequest, connection: Connection = Depends(get_connection), actor: dict = Depends(require_admin)): target = connection.execute("SELECT * FROM users WHERE id = %s", (user_id,)).fetchone() if not target: raise error("NOT_FOUND", "用户不存在", 404) starts_at = payload.starts_at or utc_now() connection.execute("UPDATE membership_grants SET revoked_at = now() WHERE user_id = %s AND revoked_at IS NULL", (user_id,)) grant = connection.execute("INSERT INTO membership_grants(user_id, plan, starts_at, expires_at, created_by, reason) VALUES (%s, %s, %s, %s, %s, %s) RETURNING id, plan, starts_at, expires_at", (user_id, payload.plan, starts_at, payload.expires_at, actor["id"], payload.reason)).fetchone() connection.execute("UPDATE users SET plan = %s, updated_at = now() WHERE id = %s", (payload.plan, user_id)) connection.execute("INSERT INTO admin_audit_logs(actor_id, action, target_type, target_id, before_value, after_value, reason) VALUES (%s, 'membership', 'user', %s, %s, %s, %s)", (actor["id"], user_id, Json({"plan": target["plan"]}), Json({"plan": payload.plan, "expires_at": payload.expires_at.isoformat() if payload.expires_at else None}), payload.reason)) connection.commit() return grant @app.post("/api/v1/admin/users/{user_id}/quota-adjustments", dependencies=[Depends(require_csrf)]) def admin_quota(user_id: UUID, payload: QuotaAdjustmentRequest, connection: Connection = Depends(get_connection), actor: dict = Depends(require_admin)): target = connection.execute("SELECT * FROM users WHERE id = %s", (user_id,)).fetchone() if not target: raise error("NOT_FOUND", "用户不存在", 404) quota = ensure_quota(connection, user_id, effective_plan(connection, target)) existing = connection.execute("SELECT id FROM usage_records WHERE user_id = %s AND idempotency_key = %s", (user_id, payload.idempotency_key)).fetchone() if existing: raise error("IDEMPOTENCY_CONFLICT", "该调整已提交", 409) connection.execute("UPDATE quota_accounts SET adjustment = adjustment + %s, version = version + 1 WHERE id = %s", (payload.amount, quota["id"])) connection.execute("INSERT INTO usage_records(user_id, quota_account_id, type, amount, idempotency_key) VALUES (%s, %s, 'adjust', %s, %s)", (user_id, quota["id"], payload.amount, payload.idempotency_key)) connection.execute("INSERT INTO admin_audit_logs(actor_id, action, target_type, target_id, after_value, reason) VALUES (%s, 'quota_adjustment', 'user', %s, %s, %s)", (actor["id"], user_id, Json({"amount": payload.amount}), payload.reason)) connection.commit() return {"status": "ok", "amount": payload.amount}