Update app.py

app.py

@@ -3,6 +3,7 @@ import json
 import random
 import asyncio
 import re
+import fcntl  # برای قفل کردن فایل و هماهنگی بین Worker ها
 from datetime import datetime
 from concurrent.futures import ThreadPoolExecutor
 from fastapi import FastAPI, HTTPException, Request
@@ -15,40 +16,31 @@ load_dotenv()
 
 app = FastAPI()
 
-# --- حافظه موقت برای محدودیت روزانه ---
-# ساختار: { "identifier_string": { "date": "YYYY-MM-DD", "count": int } }
-DAILY_LIMITS = {}
+# --- تنظیمات محدودیت ---
 MAX_FREE_DAILY_REQUESTS = 5
+USAGE_DB_FILE = "usage_data.json"  # فایل ذخیره سوابق
 
-# --- بخش اصلاح شده برای بارگذاری تمیز کلیدها ---
+# --- بخش بارگذاری کلیدها ---
 all_keys_str = os.getenv("ALL_GEMINI_API_KEYS", "")
-
-# استفاده از Regex برای جدا کردن کلیدها بر اساس کاما، خط جدید یا فاصله
 raw_keys = re.split(r'[,\s\n]+', all_keys_str)
 
 GEMINI_KEYS = []
 for k in raw_keys:
-    # حذف کوتیشن‌های احتمالی (' یا ") و کاراکترهای کنترلی مثل \r
     clean_key = k.strip().replace('"', '').replace("'", "").replace("\r", "").replace("\n", "")
-    
-    # فیلتر کردن کلیدهای خراب (کلیدهای جمینای معمولاً طولانی هستند)
     if len(clean_key) > 30:
         GEMINI_KEYS.append(clean_key)
 
 if GEMINI_KEYS:
-    print(f"تعداد {len(GEMINI_KEYS)} کلید سالم و تمیز بارگذاری شد.")
-    # نمایش ۵ کاراکتر اول یکی از کلیدها جهت اطمینان در لاگ
-    print(f"Sample loaded key start: {GEMINI_KEYS[0][:5]}...")
+    print(f"Loaded {len(GEMINI_KEYS)} keys.")
 else:
-    print("هشدار جدی: هیچ کلید معتبری یافت نشد. لطفاً متغیر ALL_GEMINI_API_KEYS را بررسی کنید.")
+    print("Warning: No valid keys found.")
 
-# تنظیم ThreadPool برای درخواست‌های موازی
 executor = ThreadPoolExecutor(max_workers=10)
 
 class IdeaRequest(BaseModel):
     idea: str
-    fingerprint: str  # شناسه مرورگر کاربر
-    is_premium: bool = False # وضعیت اشتراک
+    fingerprint: str
+    is_premium: bool = False
 
 @app.get("/", response_class=HTMLResponse)
 async def read_root():
@@ -59,80 +51,100 @@ async def read_root():
         return "<h1>Error: index.html not found</h1>"
 
 def sync_generate_content(prompt):
-    """
-    این تابع همگام است و تلاش می‌کند با کلیدهای مختلف محتوا تولید کند.
-    """
     max_retries = 1000
-    
-    # اگر کلیدی موجود نیست، فوراً خارج شو
     if not GEMINI_KEYS:
-        print("No keys available to use.")
         return None
 
     for _ in range(max_retries):
         try:
             current_key = random.choice(GEMINI_KEYS)
-            
             genai.configure(api_key=current_key)
-            model = genai.GenerativeModel('gemini-2.5-flash') # استفاده از مدل جدیدتر و سریعتر
-            
+            model = genai.GenerativeModel('gemini-2.5-flash')
             response = model.generate_content(prompt)
-            
             clean_json = response.text.replace("```json", "").replace("```", "").strip()
-            data = json.loads(clean_json)
-            
-            return data
-
-        except Exception as e:
-            # خطا را لاگ نمی‌کنیم تا لاگ‌ها شلوغ نشود، فقط ریترای می‌کنیم
+            return json.loads(clean_json)
+        except Exception:
             continue
-            
     return None
 
-def check_and_update_limit(identifier: str):
+def check_and_increment_usage(ip: str, fingerprint: str):
     """
-    بررسی محدودیت روزانه برای یک شناسه (آی‌پی یا فینگرپرینت).
-    اگر محدودیت رد شده باشد False برمی‌گرداند.
+    این تابع با استفاده از قفل فایل، دسترسی همزمان چندین Worker را مدیریت می‌کند
+    تا محدودیت دقیقاً اعمال شود.
     """
     today = datetime.now().strftime("%Y-%m-%d")
     
-    if identifier not in DAILY_LIMITS:
-        DAILY_LIMITS[identifier] = {"date": today, "count": 1}
-        return True
-    
-    record = DAILY_LIMITS[identifier]
-    
-    # اگر تاریخ رکورد قدیمی است، ریست کن
-    if record["date"] != today:
-        DAILY_LIMITS[identifier] = {"date": today, "count": 1}
-        return True
-    
-    # بررسی تعداد
-    if record["count"] >= MAX_FREE_DAILY_REQUESTS:
-        return False
-    
-    # افزایش شمارنده
-    DAILY_LIMITS[identifier]["count"] += 1
-    return True
+    # اطمینان از وجود فایل
+    if not os.path.exists(USAGE_DB_FILE):
+        with open(USAGE_DB_FILE, 'w') as f:
+            json.dump({}, f)
+
+    # باز کردن فایل با قابلیت خواندن �� نوشتن
+    with open(USAGE_DB_FILE, 'r+') as f:
+        try:
+            # اعمال قفل انحصاری (فقط یک Worker در لحظه می‌تواند فایل را بخواند/بنویسد)
+            fcntl.flock(f, fcntl.LOCK_EX)
+            
+            try:
+                data = json.load(f)
+            except json.JSONDecodeError:
+                data = {}
+
+            # آماده‌سازی کلیدها
+            ip_key = f"ip:{ip}"
+            fp_key = f"fp:{fingerprint}"
+            
+            # خواندن وضعیت فعلی IP
+            ip_record = data.get(ip_key, {"date": today, "count": 0})
+            if ip_record["date"] != today: # ریست کردن روز جدید
+                ip_record = {"date": today, "count": 0}
+            
+            # خواندن وضعیت فعلی Fingerprint
+            fp_record = data.get(fp_key, {"date": today, "count": 0})
+            if fp_record["date"] != today: # ریست کردن روز جدید
+                fp_record = {"date": today, "count": 0}
+
+            # بررسی محدودیت (اگر هر کدام از سقف رد شده باشند، مسدود می‌شود)
+            if ip_record["count"] >= MAX_FREE_DAILY_REQUESTS:
+                return False
+            if fingerprint and fp_record["count"] >= MAX_FREE_DAILY_REQUESTS:
+                return False
+
+            # افزایش شمارنده برای هر دو (برای جلوگیری از دور زدن)
+            ip_record["count"] += 1
+            data[ip_key] = ip_record
+            
+            if fingerprint:
+                fp_record["count"] += 1
+                data[fp_key] = fp_record
+
+            # بازنویسی فایل
+            f.seek(0)
+            f.truncate()
+            json.dump(data, f)
+            
+            return True
+
+        finally:
+            # آزاد کردن قفل
+            fcntl.flock(f, fcntl.LOCK_UN)
 
 @app.post("/api/refine")
 async def refine_text(request: IdeaRequest, req: Request):
     if not GEMINI_KEYS:
         raise HTTPException(status_code=500, detail="کلید API تنظیم نشده است.")
 
-    # --- بررسی محدودیت روزانه ---
+    # --- بررسی دقیق محدودیت روزانه ---
     if not request.is_premium:
         client_ip = req.client.host
-        client_fingerprint = request.fingerprint
-
-        # بررسی آی‌پی
-        if not check_and_update_limit(client_ip):
-             return JSONResponse(content={"error": "LIMIT_REACHED"}, status_code=429)
+        client_fp = request.fingerprint
+        
+        # فراخوانی تابع امن با قفل فایل
+        allowed = check_and_increment_usage(client_ip, client_fp)
         
-        # بررسی فینگرپرینت (اگر متفاوت از آی‌پی باشد)
-        if client_fingerprint and client_fingerprint != client_ip:
-             if not check_and_update_limit(client_fingerprint):
-                 return JSONResponse(content={"error": "LIMIT_REACHED"}, status_code=429)
+        if not allowed:
+            # کد 429 به معنی Too Many Requests
+            return JSONResponse(content={"error": "LIMIT_REACHED"}, status_code=429)
     # ---------------------------
 
     prompt = f"""
@@ -175,4 +187,4 @@ async def refine_text(request: IdeaRequest, req: Request):
     if result:
         return JSONResponse(content=result)
     else:
-        return JSONResponse(content={"error": "Server busy (All keys limit reached or invalid). Please try again later."}, status_code=503)
+        return JSONResponse(content={"error": "Server busy. Please try again later."}, status_code=503)