router_proxy.py

# router_proxy.py
from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from sqlalchemy.orm import Session
from pydantic import BaseModel
import httpx
import os
import 数据库连接 as json_db
from database_sql import get_db
from models_sql import Ownership

router = APIRouter()

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

@router.post("/api/proxy_github_zip")
async def proxy_github_zip(req_data: ProxyGithubZipRequest, db: Session = Depends(get_db)):
    """云端代理：校验所有权后，拉取 GitHub 私有库 ZIP 流透传给本地"""
    items_db = json_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)
    
    # 1. 核心鉴权：验证用户是否购买过该工具
    price = int(item.get("price", 0))
    if price > 0 and req_data.account != item.get("author"):
        owned = 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)

    # 2. 解析 GitHub 仓库信息
    repo_url = item.get("link", "").rstrip("/")
    if not repo_url.startswith("https://github.com/"):
        return JSONResponse(content={"error": "无效的仓库地址，目前仅支持 GitHub 私有库代理"}, status_code=400)

    repo_parts = repo_url.split("/")
    if len(repo_parts) < 2: return JSONResponse(content={"error": "无效的仓库地址格式"}, status_code=400)
    owner = repo_parts[-2]
    repo = repo_parts[-1].replace(".git", "")
    
    # GitHub 官方提供的打包下载 API
    github_zip_api = f"https://api.github.com/repos/{owner}/{repo}/zipball"
    
    # 【核心修改】：优先读取该资源在数据库中绑定的专属创作者 Token
    creator_token = item.get("github_token")
    # 如果没填，尝试使用官方全局兜底的 PAT
    fallback_token = os.environ.get("GITHUB_PAT") 
    
    active_token = creator_token if creator_token else fallback_token

    headers = {
        "Accept": "application/vnd.github.v3+json",
        "User-Agent": "ComfyUI-Ranking-SaaS"
    }
    if active_token:
        headers["Authorization"] = f"Bearer {active_token}"

    # 3. 异步请求 GitHub API 并以流形式透传回客户端 (防内存打爆)
    async def stream_generator():
        async with httpx.AsyncClient(follow_redirects=True) as client:
            async with client.stream("GET", github_zip_api, headers=headers) as response:
                if response.status_code != 200:
                    yield b"GITHUB_DOWNLOAD_FAILED"
                    return
                async for chunk in response.aiter_bytes():
                    yield chunk

    return StreamingResponse(stream_generator(), media_type="application/zip")

# ==========================================
# 新增：工作流/应用 (App) JSON 代理下载接口
# ==========================================
class ProxyDownloadRequest(BaseModel):
    url: str
    item_id: str
    account: str

@router.post("/api/proxy_download")
async def proxy_download(req_data: ProxyDownloadRequest, db: Session = Depends(get_db)):
    """云端代理：校验所有权后，拉取真实的工作流 JSON 文件并透传给本地"""
    items_db = json_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)
    
    # 1. 核心鉴权：验证用户是否购买/获取过该工作流
    price = int(item.get("price", 0))
    if price > 0 and req_data.account != item.get("author"):
        owned = 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)

    target_url = req_data.url
    
    # 2. 异步拉取真实 JSON 数据并透传回客户端
    try:
        # 开启 follow_redirects=True 防止源地址有 302 重定向跳转
        async with httpx.AsyncClient(follow_redirects=True) as client:
            resp = await client.get(target_url, timeout=30.0)
            
            if resp.status_code != 200:
                return JSONResponse(content={"error": f"源文件拉取失败，HTTP状态码: {resp.status_code}"}, status_code=resp.status_code)
            
            # 成功则直接将 JSON 二进制流原样返回给本地 ComfyUI 引擎
            from fastapi.responses import Response
            return Response(content=resp.content, media_type="application/json")
            
    except Exception as e:
        return JSONResponse(content={"error": f"代理下载时发生网络异常: {str(e)}"}, status_code=500)