update user_key dialogue

main.py

@@ -21,7 +21,7 @@ import uvicorn
 from collections import defaultdict
 from PIL import Image
 from Crypto.PublicKey import RSA
-from Crypto.Cipher import PKCS1_OAEP
+from Crypto.Cipher import PKCS1_v1_5
 from datetime import datetime
 from pymongo.mongo_client import MongoClient
 from pymongo.server_api import ServerApi
@@ -213,17 +213,28 @@ 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:
+
+def decrypt_user_key(encrypted_key: str) -> dict|None:
     try:
-        # 创建RSA解密对象
-        cipher = PKCS1_OAEP.new(RSA.import_key(SECRET_KEY))
-        
-        # 解码base64并解密
-        encrypted_data = base64.b64decode(encrypted_key)
-        decrypted_data = cipher.decrypt(encrypted_data)
-        
-        # 返回解密后的字符串
-        return decrypted_data.decode('utf-8')
+        key = RSA.import_key(SECRET_KEY)
+        cipher = PKCS1_v1_5.new(key)
+
+        # Decode the base64 encrypted text
+        encrypted_bytes = base64.b64decode(encrypted_key)
+
+        # Decrypt the message
+        decrypted_bytes = cipher.decrypt(encrypted_bytes, None)
+
+        # Convert bytes to string
+        decrypted_text = decrypted_bytes.decode('utf-8')
+
+        # Try to parse as JSON for pretty printing
+        try:
+            data = json.loads(decrypted_text)
+            return data
+        except Exception as e:
+            logging.error(f'Error: {e}, traceback: {traceback.format_exc()}')
+            return None
     except Exception as e:
         logging.error(f'Error decrypting user key: {e}')
         return None
@@ -282,15 +293,8 @@ async def optimize_design(request: OptimizationRequest):
         # 检查用户密钥
         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="无效的用户密钥"
-                )
-            try:
-                user_info = eval(user_key)
-            except:
+            user_info = decrypt_user_key(request.user_key)
+            if not user_info:
                 raise HTTPException(
                     status_code=400,
                     detail="无效的用户密钥"
@@ -340,12 +344,15 @@ async def optimize_design(request: OptimizationRequest):
             request_model_id=request.request_model_id,
             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}}
-        )
+
+        if user_key:
+            # Update credit count by decrementing it by 1
+            user_info_db = client.user_info
+            user_keys_collection = user_info_db['user_keys']
+            user_keys_collection.update_one(
+                {"user_key": request.user_key},
+                {"$inc": {"credit": -1}}
+            )
         
         return JSONResponse(response)
     except HTTPException as he: