import aiohttp from functools import wraps from fastapi import HTTPException import os from typing import Callable, Any, Dict, Optional, Tuple, List from settings import settings from fastapi.responses import StreamingResponse import json import requests import io import re from pathlib import Path from urllib.parse import urlparse from PIL import Image async def verify_api_credentials_and_balance(api_key: str, api_secret: str, api_name: str) -> Dict[str, Any]: """验证API密钥和团队余额""" try: async with aiohttp.ClientSession() as session: async with session.post( f"{settings.jingrow_api_url}/api/action/jcloud.api.account.verify_api_credentials_and_balance", headers={"Authorization": f"token {settings.jingrow_api_key}:{settings.jingrow_api_secret}"}, json={"api_key": api_key, "api_secret": api_secret, "api_name": api_name} ) as response: if response.status != 200: raise HTTPException(status_code=500, detail="验证服务暂时不可用") result = await response.json() if "message" in result and isinstance(result["message"], dict): result = result["message"] if not result.get("success"): raise HTTPException(status_code=401, detail=result.get("message", "验证失败")) return result except HTTPException: raise except Exception as e: raise HTTPException(status_code=500, detail=f"验证服务暂时不可用: {str(e)}") async def deduct_jingrow_api_usage_fee(api_key: str, api_secret: str, api_name: str, usage_count: int = 1) -> Dict[str, Any]: """从Jingrow平台扣除API使用费""" try: async with aiohttp.ClientSession() as session: async with session.post( f"{settings.jingrow_api_url}/api/action/jcloud.api.account.deduct_api_usage_fee", headers={"Authorization": f"token {settings.jingrow_api_key}:{settings.jingrow_api_secret}"}, json={ "api_key": api_key, "api_secret": api_secret, "api_name": api_name, "usage_count": usage_count } ) as response: if response.status != 200: raise HTTPException(status_code=500, detail="扣费服务暂时不可用") result = await response.json() if "message" in result and isinstance(result["message"], dict): result = result["message"] return result except HTTPException: raise except Exception as e: return {"success": False, "message": f"扣费服务暂时不可用: {str(e)}"} def get_token_from_request(request) -> str: """从请求中获取访问令牌""" if not request: raise HTTPException(status_code=400, detail="无法获取请求信息") auth_header = request.headers.get("Authorization", "") if not auth_header or not auth_header.startswith("token "): raise HTTPException(status_code=401, detail="无效的Authorization头格式") token = auth_header[6:] if ":" not in token: raise HTTPException(status_code=401, detail="无效的令牌格式") return token def jingrow_api_verify_and_billing(api_name: str): """Jingrow API 验证装饰器(带余额检查和扣费)""" def decorator(func: Callable) -> Callable: @wraps(func) async def wrapper(*args, **kwargs): try: request = kwargs.get('request') if not request: raise HTTPException(status_code=400, detail="无法获取请求信息") token = get_token_from_request(request) api_key, api_secret = token.split(":", 1) verify_result = await verify_api_credentials_and_balance(api_key, api_secret, api_name) if not verify_result.get("success"): raise HTTPException(status_code=401, detail=verify_result.get("message", "验证失败")) result = await func(*args, **kwargs) usage_count = 1 try: body_data = await request.json() if isinstance(body_data, dict): for key in ["items", "urls", "images", "files"]: if key in body_data and isinstance(body_data[key], list): usage_count = len(body_data[key]) break except Exception: pass if isinstance(result, StreamingResponse): original_generator = result.body_iterator success_count = 0 async def wrapped_generator(): nonlocal success_count async for chunk in original_generator: try: data = json.loads(chunk) if isinstance(data, dict) and data.get("status") == "success": success_count += 1 except: pass yield chunk if success_count > 0: await deduct_jingrow_api_usage_fee(api_key, api_secret, api_name, success_count) return StreamingResponse( wrapped_generator(), media_type=result.media_type, headers=result.headers ) if isinstance(result, dict) and result.get("success") is True: actual_usage_count = result.get("success_count", usage_count) await deduct_jingrow_api_usage_fee(api_key, api_secret, api_name, actual_usage_count) return result await deduct_jingrow_api_usage_fee(api_key, api_secret, api_name, usage_count) return result except HTTPException: raise except Exception as e: raise HTTPException(status_code=500, detail=f"API验证过程发生异常: {str(e)}") return wrapper return decorator def is_valid_image_url(url: str) -> bool: if not url or not isinstance(url, str): return False try: parsed = urlparse(url) if not parsed.scheme or not parsed.netloc: return False # 检查文件扩展名 path = parsed.path.lower() valid_extensions = ['.jpg', '.jpeg', '.png', '.webp', '.gif'] return any(path.endswith(ext) for ext in valid_extensions) except: return False def get_new_image_url(image_url: str) -> str: try: # 使用settings中的upload_url upload_url = settings.upload_url if not upload_url: raise HTTPException(status_code=500, detail="未配置上传URL") # 下载图片 response = requests.get(image_url, verify=False, timeout=30) if response.status_code != 200: raise HTTPException(status_code=400, detail=f"无法下载图片: HTTP {response.status_code}") image_data = response.content # 解析文件名和扩展名 parsed_url = urlparse(image_url) file_name = Path(parsed_url.path).name file_name = sanitize_filename(file_name) file_ext = Path(file_name).suffix.lower() # 如果图片是webp格式,转换为png格式 if file_ext == '.webp': image = Image.open(io.BytesIO(image_data)) png_buffer = io.BytesIO() image.save(png_buffer, format='PNG') image_data = png_buffer.getvalue() file_name = file_name.replace('.webp', '.png') # 准备文件上传 files = {"file": (file_name, image_data)} # 上传图片 upload_response = requests.post(upload_url, files=files, verify=False, timeout=30) if upload_response.status_code != 200: error_message = f"图片URL转换失败: 状态码 {upload_response.status_code}, 响应: {upload_response.text[:200]}" raise HTTPException(status_code=500, detail=error_message) result = upload_response.json() new_url = result.get("url") if not new_url: raise HTTPException(status_code=500, detail="上传成功但未返回URL") return new_url except HTTPException: raise except Exception as e: raise HTTPException(status_code=500, detail=f"图片URL转换异常: {str(e)}") def sanitize_filename(filename: str) -> str: # 移除路径分隔符和空字符 filename = re.sub(r'[\\/:*?"<>|\x00-\x1f]', '', filename) # 移除首尾空白字符 filename = filename.strip() # 如果文件名为空,使用默认名称 if not filename: filename = "untitled" # 限制文件名长度 if len(filename) > 255: name, ext = os.path.splitext(filename) filename = name[:255-len(ext)] + ext return filename