v2 token key in db

.gitignore

@@ -1,2 +1,3 @@
 __pycache__
-test.db
+test.db
+database.db

routes/v2.py

@@ -1,9 +1,72 @@
 import httpx
-from fastapi import APIRouter, HTTPException, Request
+from fastapi import APIRouter, HTTPException, Request, Depends
 import json
+import random
+import sqlite3
+from datetime import datetime
 
 router = APIRouter()
 
+# SQLite database connection
+def get_db_connection():
+    conn = sqlite3.connect('database.db')
+    conn.row_factory = sqlite3.Row
+    return conn
+
+# Initialize database
+def init_db():
+    conn = get_db_connection()
+    conn.execute('''
+        CREATE TABLE IF NOT EXISTS authorized_tokens (
+            token TEXT PRIMARY KEY,
+            expiration TEXT
+        )
+    ''')
+    conn.execute('''
+        CREATE TABLE IF NOT EXISTS api_keys (
+            api_key TEXT PRIMARY KEY,
+            expiration TEXT
+        )
+    ''')
+    conn.commit()
+    conn.close()
+
+init_db()
+
+# Import data to database
+@router.post("/import/authorized_tokens")
+async def import_authorized_tokens(tokens: dict):
+    conn = get_db_connection()
+    for token, expiration in tokens.items():
+        conn.execute('INSERT OR REPLACE INTO authorized_tokens (token, expiration) VALUES (?, ?)', (token, expiration))
+    conn.commit()
+    conn.close()
+    return {"message": "Authorized tokens imported successfully"}
+
+@router.post("/import/api_keys")
+async def import_api_keys(keys: dict):
+    conn = get_db_connection()
+    for api_key, expiration in keys.items():
+        conn.execute('INSERT OR REPLACE INTO api_keys (api_key, expiration) VALUES (?, ?)', (api_key, expiration))
+    conn.commit()
+    conn.close()
+    return {"message": "API keys imported successfully"}
+
+# Export data from database
+@router.get("/export/authorized_tokens")
+async def export_authorized_tokens():
+    conn = get_db_connection()
+    tokens = conn.execute('SELECT * FROM authorized_tokens').fetchall()
+    conn.close()
+    return {token['token']: token['expiration'] for token in tokens}
+
+@router.get("/export/api_keys")
+async def export_api_keys():
+    conn = get_db_connection()
+    keys = conn.execute('SELECT * FROM api_keys').fetchall()
+    conn.close()
+    return {key['api_key']: key['expiration'] for key in keys}
+
 @router.get("/")
 async def v2():
     return {"message": "Hello World from v2"}
@@ -26,6 +89,22 @@ async def proxy_request(protocol:str, url: str):
 
 @router.post("/{protocol}/{url:path}")
 async def proxy_post(protocol:str, url: str, request: Request):
+    # 获取请求头中的Authorization字段
+    auth_header = request.headers.get("Authorization")
+    if not auth_header:
+        raise HTTPException(status_code=401, detail="Unauthorized")
+    
+    # 去掉Bearer前缀
+    token = auth_header.split(" ")[1] if "Bearer " in auth_header else auth_header
+    
+    # 检查token是否过期
+    current_time = datetime.utcnow().isoformat()
+    conn = get_db_connection()
+    token_row = conn.execute('SELECT * FROM authorized_tokens WHERE token = ?', (token,)).fetchone()
+    conn.close()
+    if not token_row or token_row['expiration'] < current_time:
+        raise HTTPException(status_code=401, detail="Unauthorized")
+
     body = await request.body()
     headers = dict(request.headers)  # 将请求头转换为字典
     
@@ -39,6 +118,17 @@ async def proxy_post(protocol:str, url: str, request: Request):
 
     dest_url = f"{protocol}://{url}"
     
+    # 随机选择一个有效的API Key
+    conn = get_db_connection()
+    keys = conn.execute('SELECT * FROM api_keys WHERE expiration >= ?', (current_time,)).fetchall()
+    conn.close()
+    if not keys:
+        raise HTTPException(status_code=503, detail="No valid API keys available")
+    selected_api_key = random.choice(keys)['api_key']
+    print(selected_api_key)
+    print(f"Bearer {selected_api_key}")
+    headers['Authorization'] = f"Bearer {selected_api_key}"
+    
     async with httpx.AsyncClient() as client:
         try:
             # 向目标 URL 发送 POST 请求