delete client ip

main.py

@@ -20,8 +20,18 @@ from starlette.middleware.base import BaseHTTPMiddleware
 import uvicorn
 from collections import defaultdict
 from PIL import Image
+from Crypto.PublicKey import RSA
+from Crypto.Cipher import PKCS1_OAEP
+from datetime import datetime
+from pymongo.mongo_client import MongoClient
+from pymongo.server_api import ServerApi
 
 
+# mongodb uri
+URI = os.getenv("URI")
+# Create a new client and connect to the server
+client = MongoClient(URI, server_api=ServerApi('1'))
+
 app = FastAPI()
 
 # 设置最大连接数
@@ -84,6 +94,7 @@ class OptimizationRequest(BaseModel):
     suggestions: List[str]
     request_model_id: str = 'gpt-image-1'
     openai_key: str = ''
+    user_key: str = ''
 
 class CaseStudyRequest(BaseModel):
     user_input: str
@@ -202,41 +213,114 @@ async def analyze_feedback(request: AnalysisRequest):
         logging.error(f'Error: {e}, traceback: {traceback.format_exc()}')
         raise HTTPException(status_code=500, detail=f'Error: {e}, traceback: {traceback.format_exc()}')
 
+def decrypt_user_key(encrypted_key: str) -> str:
+    try:
+        # 使用与前端相同的私钥
+        private_key = 'wangyuejunhui'
+        
+        # 创建RSA解密对象
+        cipher = PKCS1_OAEP.new(RSA.import_key(private_key))
+        
+        # 解码base64并解密
+        encrypted_data = base64.b64decode(encrypted_key)
+        decrypted_data = cipher.decrypt(encrypted_data)
+        
+        # 返回解密后的字符串
+        return decrypted_data.decode('utf-8')
+    except Exception as e:
+        logging.error(f'Error decrypting user key: {e}')
+        return None
+
+def search_user_key(user_key: str) -> dict:
+    try:
+        # Create a new client and connect to the server
+        client = MongoClient(URI, server_api=ServerApi('1'))
+        
+        # Get the user_info database and user_keys collection
+        user_info_db = client.user_info
+        user_keys_collection = user_info_db['user_keys']
+        
+        # Search for the user key
+        search_res = user_keys_collection.aggregate([
+            {
+                '$search': {
+                    'index': "user_key",
+                    'text': {
+                        'query': user_key,
+                        'path': {
+                            'wildcard': "*"
+                        }
+                    }
+                }
+            },
+            {
+                '$limit': 3
+            },
+            {
+                '$project': {
+                    '_id': 0,
+                    'user_key': 1,
+                    'user_name': 1,
+                    'credit': 1,
+                    'expiration': 1
+                }
+            }
+        ])
+        
+        # Get the first result
+        result = next(search_res, None)
+        
+        if result:
+            return result
+        else:
+            return None
+            
+    except Exception as e:
+        logging.error(f'Error searching user key: {e}')
+        return None
+
 @app.post("/api/optimize-design")
-async def optimize_design(request: OptimizationRequest, client_ip: str = None):
+async def optimize_design(request: OptimizationRequest):
     try:
-        # 查看客户是否提供了openai_key，获取客户端IP（如果未提供，使用默认值）
-        if not request.openai_key and client_ip not in white_list:
-            print(f'user_ip {client_ip} not in whitelist')
+        # 检查用户密钥
+        user_key = None
+        if request.user_key:
+            user_key = decrypt_user_key(request.user_key)
+            if not user_key:
+                raise HTTPException(
+                    status_code=400,
+                    detail="无效的用户密钥"
+                )
+            user_info = eval(user_key)
+            
+            credit_data = {
+                "user_key": request.user_key,
+                "user_name": user_info['user_id'],
+                "credit": user_info['credit'],
+                "expiration": datetime.fromisoformat(user_info['expiration'])
+            }
+            
+            query_res = search_user_key(request.user_key)
+            user_info_db = client.user_info
+            user_keys_collection = user_info_db['user_keys']
+            if not query_res:
+                user_keys_collection.insert_many([credit_data])
+            else:
+                credit_data = query_res
+            
+            if credit_data['credit'] < 1:
+                raise HTTPException(
+                    status_code=503,
+                    detail="当前user-key额度已用尽。"
+                )
+        
+        # 查看客户是否提供了openai_key
+        if not request.openai_key and not user_key:
             raise HTTPException(
                 status_code=503,
-                detail="当前用户无生图权限，请点击'想使用自己的OpenAI API Key？'输入您的OpenAI API Key或联系@wangyue161添加白名单后重试。"
+                detail="当前用户无生图权限，请点击'想使用自己的OpenAI API Key？'输入您的OpenAI API Key或联系@wangyue161并添加白名单user-key后重试。"
             )
-
-        # user_rate_limit = white_list[client_ip]
-        #
-        # current_time = time.time()
-        #
-        # # 检查是否需要重置计数器（超过24小时）
-        # if current_time - optimize_design_timestamps[client_ip] > 3600*24:
-        #     optimize_design_requests[client_ip] = 0
-        #     optimize_design_timestamps[client_ip] = current_time
-        #
-        # # 如果是首次请求，记录时间戳
-        # if optimize_design_requests[client_ip] == 0:
-        #     optimize_design_timestamps[client_ip] = current_time
-        #
-        # # 检查是否超过限制
-        # print(f'user_ip {client_ip}, total requests in routine: {optimize_design_requests[client_ip]}, rate limit: {user_rate_limit}')
-        # if optimize_design_requests[client_ip] >= user_rate_limit:
-        #     raise HTTPException(
-        #         status_code=503,
-        #         detail="用户当日改图接口访问已达上限，请24小时后重试"
-        #     )
-        #
-        # # 增加请求计数
-        # optimize_design_requests[client_ip] += 1
-
+    
         # 提取设计类型
         design_type = f"设计类型：{request.text.split()[0]}\n" if len(request.text.split()) > 1 else ""
         
@@ -254,16 +338,18 @@ async def optimize_design(request: OptimizationRequest, client_ip: str = None):
             openai_key=request.openai_key
         )
         
+        # Update credit count by decrementing it by 1
+        user_keys_collection.update_one(
+            {"user_key": request.user_key},
+            {"$inc": {"credit": -1}}
+        )
+        
         return JSONResponse(response)
     except HTTPException as he:
         raise he
     except Exception as e:
         logging.error(f'Error: {e}, traceback: {traceback.format_exc()}')
         raise HTTPException(status_code=500, detail=f'Error: {e}, traceback: {traceback.format_exc()}')
-    finally:
-        # 如果发生异常，减少请求计数
-        if 'he' in locals() and isinstance(he, HTTPException):
-            optimize_design_requests[client_ip] -= 1
 
 @app.post("/api/optimize-text")
 async def optimize_text(request: TextOptimizationRequest):

requirements.txt

@@ -7,4 +7,6 @@ python-dotenv==1.0.0
 pydantic==2.5.2
 requests==2.31.0
 openai
-pillow
+pillow
+Crypto
+pymongo

static/script.js

@@ -714,6 +714,7 @@ document.addEventListener('DOMContentLoaded', function() {
                     suggestions: suggestions,
                     request_model_id: 'gpt-image-1',
                     openai_key: sessionStorage.getItem('userOpenaiKey') || '',
+                    user_key: sessionStorage.getItem('userKey') || ''
                 })
             });
 
@@ -785,7 +786,8 @@ document.addEventListener('DOMContentLoaded', function() {
                     image_data: uploadedImage,
                     suggestions: suggestions,
                     request_model_id: 'dall-e-2',
-                    openai_key: sessionStorage.getItem('userOpenaiKey') || ''
+                    openai_key: sessionStorage.getItem('userOpenaiKey') || '',
+                    user_key: sessionStorage.getItem('userKey') || ''
                 })
             });
             if (response.status === 503) {
@@ -1002,4 +1004,52 @@ document.addEventListener('DOMContentLoaded', function() {
             }
         });
     }
+
+    // User Key Dialog functionality
+    const userKeyLink = document.getElementById('userKeyLink');
+    const userKeyDialog = document.getElementById('userKeyDialog');
+    const userKeyInput = document.getElementById('userKeyInput');
+    const saveUserKey = document.getElementById('saveUserKey');
+    const cancelUserKey = document.getElementById('cancelUserKey');
+
+    if (userKeyLink && userKeyDialog) {
+        // Show dialog when clicking the link
+        userKeyLink.addEventListener('click', function(e) {
+            e.preventDefault();
+            userKeyDialog.style.display = 'flex';
+            // Load existing key if available
+            const existingKey = sessionStorage.getItem('userKey');
+            if (existingKey) {
+                userKeyInput.value = existingKey;
+            }
+        });
+
+        // Save User Key
+        saveUserKey.addEventListener('click', function() {
+            const key = userKeyInput.value.trim();
+            if (key) {
+                sessionStorage.setItem('userKey', key);
+                userKeyDialog.style.display = 'none';
+                // Show success message
+                const toast = document.getElementById('toast');
+                toast.textContent = '用户密钥已保存';
+                toast.style.display = 'block';
+                setTimeout(() => {
+                    toast.style.display = 'none';
+                }, 2000);
+            }
+        });
+
+        // Cancel button
+        cancelUserKey.addEventListener('click', function() {
+            userKeyDialog.style.display = 'none';
+        });
+
+        // Close dialog when clicking outside
+        userKeyDialog.addEventListener('click', function(e) {
+            if (e.target === userKeyDialog) {
+                userKeyDialog.style.display = 'none';
+            }
+        });
+    }
 });