From 17f13550992a56838dcb4e65d07ff4827f037d7b Mon Sep 17 00:00:00 2001
From: yinpeng <2291314224@qq.com>
Date: Tue, 11 Feb 2025 21:32:21 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=BC=BA=E5=BA=94=E7=94=A8?=
 =?UTF-8?q?=E6=97=A5=E5=BF=97=E8=AE=B0=E5=BD=95=E5=B9=B6=E4=BC=98=E5=8C=96?=
 =?UTF-8?q?=E9=94=99=E8=AF=AF=E5=A4=84=E7=90=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/main.py | 48 ++++++++++++++++++++++++++++++++----------------
 1 file changed, 32 insertions(+), 16 deletions(-)

diff --git a/app/main.py b/app/main.py
index c48f4d6..dc22b4a 100644
--- a/app/main.py
+++ b/app/main.py
@@ -25,8 +25,13 @@ key_manager = None
 @app.on_event("startup")
 async def startup_event():
     global key_manager
-    key_manager = await get_key_manager_instance(settings.API_KEYS)
-
+    logger.info("Application starting up...")
+    try:
+        key_manager = await get_key_manager_instance(settings.API_KEYS)
+        logger.info("KeyManager initialized successfully")
+    except Exception as e:
+        logger.error(f"Failed to initialize KeyManager: {str(e)}")
+        raise
 
 # 添加中间件来处理未经身份验证的请求
 @app.middleware("http")
@@ -41,7 +46,9 @@ async def auth_middleware(request: Request, call_next):
         not request.url.path.startswith("/hf")):
         auth_token = request.cookies.get("auth_token")
         if not auth_token or not verify_auth_token(auth_token):
+            logger.warning(f"Unauthorized access attempt to {request.url.path}")
             return RedirectResponse(url="/")
+        logger.debug("Request authenticated successfully")
     response = await call_next(request)
     return response
 
@@ -76,31 +83,40 @@ async def authenticate(request: Request):
         form = await request.form()
         auth_token = form.get("auth_token")
         if not auth_token:
+            logger.warning("Authentication attempt with empty token")
             return RedirectResponse(url="/", status_code=302)
         
         if verify_auth_token(auth_token):
+            logger.info("Successful authentication")
             response = RedirectResponse(url="/keys", status_code=302)
             response.set_cookie(key="auth_token", value=auth_token, httponly=True, max_age=3600)
             return response
+        logger.warning("Failed authentication attempt with invalid token")
         return RedirectResponse(url="/", status_code=302)
     except Exception as e:
         logger.error(f"Authentication error: {str(e)}")
         return RedirectResponse(url="/", status_code=302)
 
-
 @app.get("/keys", response_class=HTMLResponse)
 async def keys_page(request: Request):
-    auth_token = request.cookies.get("auth_token")
-    if not auth_token or not verify_auth_token(auth_token):
-        return RedirectResponse(url="/", status_code=302)
-    keys_status = await key_manager.get_keys_by_status()
-    total = len(keys_status["valid_keys"]) + len(keys_status["invalid_keys"])
-    return templates.TemplateResponse("keys_status.html", {
-        "request": request,
-        "valid_keys": keys_status["valid_keys"],
-        "invalid_keys": keys_status["invalid_keys"],
-        "total": total
-    })
+    try:
+        auth_token = request.cookies.get("auth_token")
+        if not auth_token or not verify_auth_token(auth_token):
+            logger.warning("Unauthorized access attempt to keys page")
+            return RedirectResponse(url="/", status_code=302)
+        
+        keys_status = await key_manager.get_keys_by_status()
+        total = len(keys_status["valid_keys"]) + len(keys_status["invalid_keys"])
+        logger.info(f"Keys status retrieved successfully. Total keys: {total}")
+        return templates.TemplateResponse("keys_status.html", {
+            "request": request,
+            "valid_keys": keys_status["valid_keys"],
+            "invalid_keys": keys_status["invalid_keys"],
+            "total": total
+        })
+    except Exception as e:
+        logger.error(f"Error retrieving keys status: {str(e)}")
+        raise
 
 
 @app.get("/health")
@@ -108,7 +124,7 @@ async def health_check(request: Request):
     logger.info("Health check endpoint called")
     return {"status": "healthy"}
     
-
+    
 if __name__ == "__main__":
-
+    logger.info("Starting application server...")
     uvicorn.run(app, host="0.0.0.0", port=8001)