app.py

# ⚙️ 后端逻辑/核心服务端.py (Hugging Face Spaces app.py)
from fastapi import FastAPI, File, UploadFile, Form, Depends, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response, JSONResponse
from sqlalchemy.orm import Session
from pydantic import BaseModel
from huggingface_hub import hf_hub_download, HfApi
import hashlib
import urllib.parse
import urllib.request
import urllib.error  # 【新增】：用于捕获 HTTPError 以判断私有库状态
import os
import 数据库连接 as db

from router_users import router as users_router
from router_items import router as items_router
from router_comments import router as comments_router
from router_messages import router as messages_router
from router_wallet import router as wallet_router
from router_proxy import router as proxy_router # 导入代理路由

from database_sql import init_sql_db, get_db
from models_sql import Ownership

app = FastAPI(title="ComfyUI Ranking Community API")

@app.on_event("startup")
def on_startup():
    init_sql_db()
    print("关系型数据库加载完毕，金融表同步完成。")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(users_router)
app.include_router(items_router)
app.include_router(comments_router)
app.include_router(messages_router)
app.include_router(wallet_router)
app.include_router(proxy_router) 

@app.get("/")
def read_root(): 
    return {"status": "ok", "message": "API System Protected & Running"}

# 【安全优化】：允许的文件后缀白名单，防挂马
ALLOWED_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".json", ".zip"}

@app.post("/api/upload")
async def upload_file(file: UploadFile = File(...), file_type: str = Form(...)):
    # 验证后缀名
    _, ext = os.path.splitext(file.filename)
    if ext.lower() not in ALLOWED_EXTENSIONS:
        return JSONResponse(status_code=400, content={"error": f"安全拦截：不支持上传 {ext} 格式的文件"})

    # 限制单次读取文件大小，防止撑爆内存
    content = await file.read()
    if len(content) > 10 * 1024 * 1024: # 10MB 限制
        return JSONResponse(status_code=400, content={"error": "文件超过 10MB 大小限制"})

    file_hash = hashlib.md5(content).hexdigest()[:10]
    
    new_filename = f"{file_hash}{ext.lower()}"
    safe_filename = urllib.parse.quote(file.filename)
    safe_url_filename = f"{file_hash}_{safe_filename}"
    
    dir_mapping = {"avatar": "avatars", "cover": "covers", "tool": "tools", "app": "apps"}
    target_dir = dir_mapping.get(file_type, "others")
    full_path_in_repo = f"{target_dir}/{new_filename}"
    
    # 交给底层带锁与异步线程的 db 处理
    db.save_file(full_path_in_repo, content)
    
    url = f"https://huggingface.co/datasets/{db.DATASET_REPO_ID}/resolve/main/{target_dir}/{safe_url_filename}"
    return {"status": "success", "url": url, "display_name": file.filename, "hashed_name": new_filename}

class ValidateRequest(BaseModel):
    item_id: str

@app.post("/api/validate_resource")
async def validate_resource(req: ValidateRequest):
    items_db = db.load_data("items.json", default_data=[])
    item = next((i for i in items_db if i["id"] == req.item_id), None)
    if not item:
        return JSONResponse(content={"error": "该资源已被原作者删除"}, status_code=404)
        
    link = item.get("link", "")
    itype = item.get("type", "")
    
    if itype.startswith("tool"):
        headers = {'User-Agent': 'Mozilla/5.0'}
        # 【核心升级】：提取该资源绑定的私有库密匙，如果没有则用全局兜底
        github_token = item.get("github_token") or os.environ.get("GITHUB_PAT")
        
        if github_token and link.startswith("https://github.com/"):
            # 针对私有库：调用 GitHub API 带着 Token 进行身份核验探测
            repo_parts = link.rstrip("/").split("/")
            if len(repo_parts) >= 2:
                owner, repo = repo_parts[-2], repo_parts[-1]
                api_link = f"https://api.github.com/repos/{owner}/{repo}"
                headers["Authorization"] = f"Bearer {github_token}"
                headers["Accept"] = "application/vnd.github.v3+json"
                try:
                    req_obj = urllib.request.Request(api_link, method="GET", headers=headers)
                    with urllib.request.urlopen(req_obj, timeout=5) as response:
                        if response.status >= 400:
                            return JSONResponse(content={"error": "私有仓库访问失败，可能密匙已失效"}, status_code=400)
                except urllib.error.HTTPError as e:
                    # 如果抛出 HTTPError (如 401 Unauthorized 或 404 Not Found)，说明 Token 假了或库被删了
                    return JSONResponse(content={"error": f"该私有库的访问密匙已失效或仓库已被原作者删除 (HTTP {e.code})"}, status_code=400)
                except Exception:
                    return JSONResponse(content={"error": "无法连接到 GitHub 验证仓库有效性"}, status_code=400)
                
                # 走到这里说明私有库和密匙都是 100% 有效的，放行！
                return {"status": "success"}

        # 针对普通公开库的无感探测
        try:
            req_obj = urllib.request.Request(link, method="HEAD", headers=headers)
            with urllib.request.urlopen(req_obj, timeout=5) as response:
                if response.status >= 400:
                    return JSONResponse(content={"error": "原作者的 Git 仓库已失效或设为私有"}, status_code=400)
        except Exception:
            return JSONResponse(content={"error": "原作者的 Git 仓库无法访问，链接已失效"}, status_code=400)
            
    elif itype.startswith("app"):
        if "resolve/main/" in link:
            repo_path = urllib.parse.unquote(link.split("resolve/main/")[-1])
            hf_token = os.environ.get("HF_TOKEN")
            try:
                api = HfApi()
                exists = api.file_exists(repo_id=db.DATASET_REPO_ID, filename=repo_path, repo_type="dataset", token=hf_token)
                if not exists:
                    return JSONResponse(content={"error": "该工作流的 JSON 文件已在云端损坏或丢失"}, status_code=400)
            except Exception:
                pass 
                
    return {"status": "success"}

class ProxyDownloadRequest(BaseModel):
    url: str
    item_id: str
    account: str

@app.post("/api/proxy_download")
async def proxy_download(req_data: ProxyDownloadRequest, sql_db: Session = Depends(get_db)):
    target_url = req_data.url
    if not target_url or "resolve/main/" not in target_url:
        return JSONResponse(content={"error": "无效的 Hugging Face 下载链接"}, status_code=400)
        
    items_db = db.load_data("items.json", default_data=[])
    item = next((i for i in items_db if i["id"] == req_data.item_id), None)
    
    if not item: return JSONResponse(content={"error": "资源不存在或已被删除"}, status_code=404)
        
    price = int(item.get("price", 0))
    author = item.get("author")
    
    if price > 0 and req_data.account != author:
        owned = sql_db.query(Ownership).filter(Ownership.account == req_data.account, Ownership.item_id == req_data.item_id).first()
        if not owned:
            return JSONResponse(content={"error": "🚨 非法下载：云端数据库未找到您的购买凭证！"}, status_code=403)

    hf_token = os.environ.get("HF_TOKEN")
    if not hf_token: return JSONResponse(content={"error": "云端环境变量未配置 HF_TOKEN"}, status_code=401)
        
    try:
        repo_path_encoded = target_url.split("resolve/main/")[-1]
        repo_path = urllib.parse.unquote(repo_path_encoded)
        
        cached_file_path = hf_hub_download(
            repo_id=db.DATASET_REPO_ID,
            repo_type="dataset",
            filename=repo_path,
            token=hf_token
        )
        
        with open(cached_file_path, "rb") as f:
            content = f.read()
            
        return Response(content=content, media_type="application/json")
        
    except Exception as e:
        return JSONResponse(content={"error": "云端代理读取失败，可能是源文件损坏"}, status_code=500)