Update app.py

app.py

@@ -1,6 +1,7 @@
 from fastapi import FastAPI, HTTPException, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, HTMLResponse
+from fastapi.staticfiles import StaticFiles
 import httpx
 import os
 import json
@@ -27,7 +28,8 @@ class Config:
     OPENAI_API_BASE = "https://api.x.ai/v1"
     KEYS_URL = os.getenv("KEYS_URL", "")
     WHITELIST_IPS = os.getenv("WHITELIST_IPS", "").split(",")
-    CHUNK_SIZE = 1  # 控制流式传输的块大小
+    ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD", "")
+    CHUNK_SIZE = 1
 
 # 全局变量
 keys = []
@@ -53,7 +55,6 @@ def init_keys():
 
 # 获取真实IP地址
 def get_client_ip(request: Request) -> str:
-    # 尝试从各种头部获取IP
     forwarded_for = request.headers.get("x-forwarded-for")
     if forwarded_for:
         return forwarded_for.split(",")[0].strip()
@@ -84,6 +85,92 @@ def get_next_key():
         raise HTTPException(status_code=500, detail="No API keys available")
     return next(key_cycle)
 
+# 添加静态文件支持
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+# 管理页面路由
+@app.get("/admin", response_class=HTMLResponse)
+async def admin_page():
+    try:
+        with open("static/admin.html", "r", encoding="utf-8") as f:
+            return f.read()
+    except Exception as e:
+        print(f"Error loading admin page: {e}")
+        raise HTTPException(status_code=500, detail="Error loading admin page")
+
+# 验证管理密码
+@app.post("/api/admin/login")
+async def admin_login(password: str):
+    if password == Config.ADMIN_PASSWORD:
+        return {"status": "success"}
+    raise HTTPException(status_code=401, detail="Invalid password")
+
+# 获取所有keys（需要密码）
+@app.get("/api/keys")
+async def get_keys(password: str):
+    if password != Config.ADMIN_PASSWORD:
+        raise HTTPException(status_code=401, detail="Invalid password")
+    return {"keys": list(keys)}
+
+# 检查key是否有效
+async def check_key_valid(key: str) -> bool:
+    headers = {
+        "Authorization": f"Bearer {key}",
+        "Content-Type": "application/json"
+    }
+    
+    async with httpx.AsyncClient() as client:
+        try:
+            response = await client.get(
+                f"{Config.OPENAI_API_BASE}/models",
+                headers=headers
+            )
+            return response.status_code == 200
+        except:
+            return False
+
+# 添加新key（需要密码）
+@app.post("/api/keys")
+async def add_key(key: str, password: str):
+    if password != Config.ADMIN_PASSWORD:
+        raise HTTPException(status_code=401, detail="Invalid password")
+        
+    if key not in keys:
+        # 检查key是否有效
+        is_valid = await check_key_valid(key)
+        if not is_valid:
+            raise HTTPException(status_code=400, detail="Invalid key")
+            
+        keys.append(key)
+        global key_cycle
+        key_cycle = cycle(keys)
+        return {"status": "success", "message": "Key added successfully"}
+    
+    raise HTTPException(status_code=400, detail="Key already exists")
+
+# 删除key（需要密码）
+@app.delete("/api/keys/{key}")
+async def delete_key(key: str, password: str):
+    if password != Config.ADMIN_PASSWORD:
+        raise HTTPException(status_code=401, detail="Invalid password")
+        
+    if key in keys:
+        keys.remove(key)
+        global key_cycle
+        key_cycle = cycle(keys)
+        return {"status": "success", "message": "Key deleted successfully"}
+    
+    raise HTTPException(status_code=404, detail="Key not found")
+
+# 检查key状态（需要密码）
+@app.get("/api/keys/check/{key}")
+async def check_key(key: str, password: str):
+    if password != Config.ADMIN_PASSWORD:
+        raise HTTPException(status_code=401, detail="Invalid password")
+        
+    is_valid = await check_key_valid(key)
+    return {"valid": is_valid}
+
 # 模型列表路由
 @app.get("/api/v1/models")
 async def list_models():
@@ -209,7 +296,7 @@ async def proxy(path: str, request: Request):
     # 获取并处理headers
     headers = dict(request.headers)
     headers.pop("host", None)
-    headers["Authorization"] = f"Bearer {get_next_key()}"
+    headers["Authorization"] = f"Bearer {get_next_key()}",
     headers["Content-Type"] = "application/json"
     
     # 构建目标URL