app.py

from fastapi import FastAPI, Request, HTTPException, Depends, status
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from fastapi.responses import JSONResponse, StreamingResponse
import httpx
import yaml
import json
from typing import Dict, List, Optional
import urllib.parse
from pathlib import Path
import os
import secrets

app = FastAPI(title="Dynamic Multi-Site Reverse Proxy")

# 配置基本认证
security = HTTPBasic()

# 从环境变量获取认证用户名和密码
ADMIN_USERNAME = os.getenv("ADMIN_USERNAME", "admin")
ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD", "default_password")  # 请务必修改!

def verify_admin(credentials: HTTPBasicCredentials = Depends(security)):
    """验证管理员身份"""
    correct_username = secrets.compare_digest(credentials.username, ADMIN_USERNAME)
    correct_password = secrets.compare_digest(credentials.password, ADMIN_PASSWORD)
    
    if not (correct_username and correct_password):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="无效的认证信息",
            headers={"WWW-Authenticate": "Basic"},
        )
    return credentials.username

# 内存中的站点配置
SITE_CONFIG = {}
DOMAIN_MAPPING = {}

async def load_initial_config():
    """从环境变量或本地文件加载初始配置"""
    global SITE_CONFIG, DOMAIN_MAPPING
    
    # 方式1: 从环境变量读取配置 (JSON字符串)
    config_json = os.getenv("PROXY_CONFIG")
    if config_json:
        try:
            config = json.loads(config_json)
            SITE_CONFIG = config.get("sites", {})
            DOMAIN_MAPPING = config.get("domains", {})
            print("从环境变量加载配置成功")
            return
        except json.JSONDecodeError as e:
            print(f"环境变量配置JSON解析失败: {e}")
    
    # 方式2: 从本地config.json文件加载
    config_file = Path("config.json")
    if config_file.exists():
        try:
            with open(config_file, 'r') as f:
                config = json.load(f)
            SITE_CONFIG = config.get("sites", {})
            DOMAIN_MAPPING = config.get("domains", {})
            print("从本地config.json加载配置成功")
        except Exception as e:
            print(f"加载本地配置文件失败: {e}")
    
    # 方式3: 如果没有配置，使用示例配置
    if not SITE_CONFIG:
        SITE_CONFIG = {
            "site1": {
                "target": "http://106.15.4.153:8085",
                "description": "主站点",
                "enabled": True
            }
        }
        DOMAIN_MAPPING = {
            "www.example.com": "site1"
        }
        print("使用内置示例配置")

@app.on_event("startup")
async def startup_event():
    """应用启动时加载配置"""
    await load_initial_config()
    
    # 可选: 从远程URL加载配置
    remote_config_url = os.getenv("REMOTE_CONFIG_URL")
    if remote_config_url:
        try:
            await load_remote_config(remote_config_url)
        except Exception as e:
            print(f"加载远程配置失败: {e}")

async def load_remote_config(config_url: str):
    """从远程URL加载配置"""
    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(config_url, timeout=10.0)
            if response.status_code == 200:
                remote_config = response.json()
                SITE_CONFIG.update(remote_config.get("sites", {}))
                DOMAIN_MAPPING.update(remote_config.get("domains", {}))
                print("从远程URL加载配置成功")
                return True
    except Exception as e:
        print(f"加载远程配置失败: {e}")
    return False

# 管理API端点
@app.get("/_config")
async def get_config(username: str = Depends(verify_admin)):
    """获取当前配置（需要认证）"""
    return {
        "sites": SITE_CONFIG,
        "domains": DOMAIN_MAPPING
    }

@app.post("/_config/update")
async def update_config(
    site_id: str, 
    target: str, 
    description: str = "",
    username: str = Depends(verify_admin)
):
    """动态更新站点配置"""
    SITE_CONFIG[site_id] = {
        "target": target,
        "description": description,
        "enabled": True
    }
    return {"message": "配置更新成功", "site_id": site_id}

@app.delete("/_config/remove/{site_id}")
async def remove_config(site_id: str, username: str = Depends(verify_admin)):
    """删除站点配置"""
    if site_id in SITE_CONFIG:
        del SITE_CONFIG[site_id]
        # 同时移除相关的域名映射
        global DOMAIN_MAPPING
        DOMAIN_MAPPING = {k: v for k, v in DOMAIN_MAPPING.items() if v != site_id}
        return {"message": f"站点 {site_id} 已删除"}
    raise HTTPException(status_code=404, detail="站点不存在")

@app.post("/_config/domain")
async def add_domain_mapping(
    domain: str, 
    site_id: str,
    username: str = Depends(verify_admin)
):
    """添加域名映射"""
    if site_id not in SITE_CONFIG:
        raise HTTPException(status_code=404, detail="站点不存在")
    
    DOMAIN_MAPPING[domain] = site_id
    return {"message": f"域名 {domain} 已映射到站点 {site_id}"}

# 代理路由
@app.api_route("/{site_id}/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "PATCH", "HEAD"])
@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "PATCH", "HEAD"])
async def proxy_request(
    request: Request, 
    site_id: Optional[str] = None, 
    path: str = ""
):
    """多站点反向代理"""
    try:
        # 确定目标站点
        host = request.headers.get("host", "").split(":")[0]
        if host in DOMAIN_MAPPING:
            site_id = DOMAIN_MAPPING[host]
        elif not site_id:
            site_id = request.query_params.get("__site")
        
        if not site_id or site_id not in SITE_CONFIG:
            site_id = next(iter(SITE_CONFIG.keys()), None)  # 获取第一个站点作为默认
        
        if not site_id:
            raise HTTPException(status_code=404, detail="没有可用的代理站点")
        
        site_config = SITE_CONFIG.get(site_id)
        if not site_config or not site_config.get("enabled", True):
            raise HTTPException(status_code=404, detail=f"站点 {site_id} 未找到或已禁用")
        
        # 构建目标URL并转发请求
        target_base = site_config["target"].rstrip("/")
        target_url = f"{target_base}/{path}" if path else target_base
        
        # 处理查询参数
        query_params = dict(request.query_params)
        query_params.pop("__site", None)
        if query_params:
            query_str = "&".join([f"{k}={urllib.parse.quote(v)}" for k, v in query_params.items()])
            target_url += f"?{query_str}"
        
        # 转发请求
        async with httpx.AsyncClient(timeout=60.0) as client:
            headers = dict(request.headers)
            headers.pop("host", None)
            headers.pop("content-length", None)
            
            response = await client.request(
                method=request.method,
                url=target_url,
                headers=headers,
                content=await request.body(),
                timeout=60.0,
                follow_redirects=False
            )
        
        # 返回响应
        response_headers = dict(response.headers)
        response_headers.pop("content-length", None)
        
        return StreamingResponse(
            content=response.aiter_bytes(),
            status_code=response.status_code,
            headers=response_headers,
            media_type=response.headers.get("content-type")
        )
            
    except httpx.RequestError as e:
        raise HTTPException(status_code=502, detail=f"代理错误: {str(e)}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"内部错误: {str(e)}")

@app.get("/")
async def root():
    """显示使用说明"""
    return {
        "message": "多站点动态反向代理",
        "usage": {
            "通过域名访问": "不同的域名自动路由到不同站点",
            "通过路径访问": "/site1/path 或 /site2/path",
            "通过参数访问": "/any/path?__site=site1",
            "管理配置": {
                "查看配置": "GET /_config (需要认证)",
                "更新配置": "POST /_config/update?site_id=xxx&target=xxx (需要认证)",
                "删除配置": "DELETE /_config/remove/site_id (需要认证)",
                "添加域名映射": "POST /_config/domain?domain=xxx&site_id=xxx (需要认证)"
            },
            "当前可用站点": list(SITE_CONFIG.keys())
        }
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)