Update app.py

app.py

@@ -20,32 +20,58 @@ logger = logging.getLogger(__name__)
 TOKEN = os.environ.get("TOKEN")
 DOWNLOAD_DIR = "downloads"
 SPACE_URL = "https://barnamemaker-pdfmeger.hf.space"
+MERGED_FILE_TTL = 3600        # مدت نگهداری فایل‌های ادغام‌شده (ثانیه)
+
 os.makedirs(DOWNLOAD_DIR, exist_ok=True)
-user_files = {}
+
+# user_id -> list of file paths
+user_files: dict[int, list[str]] = {}
+# قفل برای جلوگیری از race condition
+user_locks: dict[int, threading.Lock] = {}
 
 app_web = FastAPI()
 
+def get_user_lock(user_id: int) -> threading.Lock:
+    if user_id not in user_locks:
+        user_locks[user_id] = threading.Lock()
+    return user_locks[user_id]
+
+# ───────────────────────── Web server ─────────────────────────
+
 @app_web.get("/")
 async def root():
     return {"message": "Bot is running!"}
 
 @app_web.get("/download/{filename}")
 async def get_file(filename: str):
-    file_path = f"{DOWNLOAD_DIR}/{filename}"
+    # جلوگیری از path traversal
+    if ".." in filename or "/" in filename:
+        return {"error": "Invalid filename"}
+    file_path = os.path.join(DOWNLOAD_DIR, filename)
     if os.path.exists(file_path):
-        return FileResponse(file_path)
+        return FileResponse(file_path, media_type="application/pdf")
     return {"error": "File not found"}
 
 def run_web():
     uvicorn.run(app_web, host="0.0.0.0", port=7860)
 
+# ───────────────────────── Cleanup ─────────────────────────
+
 def cleanup_files():
+    """هر ساعت فایل‌های قدیمی‌تر از TTL را پاک می‌کند."""
     while True:
-        time.sleep(3600)
+        time.sleep(600)  # هر ۱۰ دقیقه چک کن
+        now = time.time()
         for f in os.listdir(DOWNLOAD_DIR):
             path = os.path.join(DOWNLOAD_DIR, f)
-            if os.path.exists(path) and (time.time() - os.path.getmtime(path) > 3600):
-                os.remove(path)
+            try:
+                if os.path.isfile(path) and (now - os.path.getmtime(path) > MERGED_FILE_TTL):
+                    os.remove(path)
+                    logger.info(f"Cleaned up old file: {f}")
+            except Exception as e:
+                logger.warning(f"Could not remove {f}: {e}")
+
+# ───────────────────────── UI helpers ─────────────────────────
 
 def get_menu():
     keyboard = [
@@ -54,62 +80,191 @@ def get_menu():
     ]
     return ReplyKeyboardMarkup(keyboard, resize_keyboard=True, one_time_keyboard=False)
 
+HELP_TEXT = (
+    "📖 راهنمای استفاده:\n\n"
+    "۱. فایل‌های PDF خود را به هر تعدادی که می‌خواهید یکی‌یکی بفرستید.\n"
+    "۲. روی «📄 ادغام فایل‌ها» کلیک کنید.\n"
+    "۳. لینک دانلود فایل ادغام‌شده را دریافت کنید.\n\n"
+    "برای پاک کردن صف: «🗑 پاک کردن صف»\n"
+    "فایل‌ها به ترتیب ارسال ادغام می‌شوند."
+)
+
+# ───────────────────────── Handlers ─────────────────────────
+
 async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
-    logger.info(f"User {update.effective_user.id} started the bot.")
-    text = "سلام! فایل‌های خود را بفرستید و سپس روی «📄 ادغام فایل‌ها» کلیک کنید."
-    await update.message.reply_text(text, reply_markup=get_menu())
+    user_id = update.effective_user.id
+    logger.info(f"User {user_id} started the bot.")
+    await update.message.reply_text(
+        "سلام! 👋\nفایل‌های PDF خود را بفرستید و سپس «📄 ادغام فایل‌ها» را بزنید.\n\n"
+        "برای راهنما: ❓ راهنما",
+        reply_markup=get_menu()
+    )
 
 async def handle_document(update: Update, context: ContextTypes.DEFAULT_TYPE):
     user_id = update.effective_user.id
-    logger.info(f"User {user_id} sent a document.")
-    file = await update.message.document.get_file()
-    file_id = f"{user_id}_{uuid.uuid4().hex[:6]}.pdf"
-    file_path = f"{DOWNLOAD_DIR}/{file_id}"
-    await file.download_to_drive(file_path)
-    if user_id not in user_files: user_files[user_id] = []
-    user_files[user_id].append(file_path)
-    await update.message.reply_text(f"✅ فایل {len(user_files[user_id])} دریافت شد.", reply_markup=get_menu())
+    doc = update.message.document
+
+    # بررسی نوع فایل
+    if doc.mime_type != "application/pdf":
+        await update.message.reply_text("⚠️ فقط فایل PDF قبول می‌شود.", reply_markup=get_menu())
+        return
+
+    lock = get_user_lock(user_id)
+    with lock:
+        if user_id not in user_files:
+            user_files[user_id] = []
+
+        logger.info(f"User {user_id} sent a document: {doc.file_name}")
+        file = await doc.get_file()
+        file_id = f"{user_id}_{uuid.uuid4().hex[:8]}.pdf"
+        file_path = os.path.join(DOWNLOAD_DIR, file_id)
+
+        try:
+            await file.download_to_drive(file_path)
+        except Exception as e:
+            logger.error(f"Download error for user {user_id}: {e}")
+            await update.message.reply_text("❌ خطا در دریافت فایل. دوباره امتحان کنید.", reply_markup=get_menu())
+            return
+
+        # اعتبارسنجی PDF (جلوگیری از فایل‌های خراب)
+        try:
+            with fitz.open(file_path) as test_doc:
+                if test_doc.page_count == 0:
+                    raise ValueError("Empty PDF")
+        except Exception:
+            os.remove(file_path)
+            await update.message.reply_text("⚠️ فایل PDF معتبر نیست یا خراب است.", reply_markup=get_menu())
+            return
+
+        user_files[user_id].append(file_path)
+        count = len(user_files[user_id])
+
+    await update.message.reply_text(
+        f"✅ فایل {count} دریافت شد: {doc.file_name}",
+        reply_markup=get_menu()
+    )
 
 async def merge_pdfs(update: Update, context: ContextTypes.DEFAULT_TYPE):
     user_id = update.effective_user.id
-    files = user_files.get(user_id, [])
+    lock = get_user_lock(user_id)
+
+    with lock:
+        files = list(user_files.get(user_id, []))
+
     if not files:
-        await update.message.reply_text("⚠️ صف شما خالی است!", reply_markup=get_menu())
+        await update.message.reply_text("⚠️ صف شما خالی است! ابتدا فایل‌های PDF بفرستید.", reply_markup=get_menu())
         return
-    await update.message.reply_text("⏳ در حال ادغام...")
-    output_filename = f"merged_{user_id}_{uuid.uuid4().hex[:6]}.pdf"
-    output_path = f"{DOWNLOAD_DIR}/{output_filename}"
+
+    if len(files) == 1:
+        await update.message.reply_text("⚠️ حداقل ۲ فایل برای ادغام لازم است.", reply_markup=get_menu())
+        return
+
+    await update.message.reply_text(f"⏳ در حال ادغام {len(files)} فایل...", reply_markup=get_menu())
+
+    output_filename = f"merged_{user_id}_{uuid.uuid4().hex[:8]}.pdf"
+    output_path = os.path.join(DOWNLOAD_DIR, output_filename)
+
     try:
         result = fitz.open()
-        for pdf in files:
-            with fitz.open(pdf) as doc:
+        for pdf_path in files:
+            if not os.path.exists(pdf_path):
+                logger.warning(f"Missing file during merge: {pdf_path}")
+                continue
+            with fitz.open(pdf_path) as doc:
                 result.insert_pdf(doc)
+
+        if result.page_count == 0:
+            result.close()
+            await update.message.reply_text("❌ هیچ صفحه‌ای برای ادغام یافت نشد.", reply_markup=get_menu())
+            return
+
         result.save(output_path)
+        result.close()
+
+        # پاک کردن فایل‌های موقت کاربر
+        with lock:
+            for f in user_files.get(user_id, []):
+                try:
+                    if os.path.exists(f):
+                        os.remove(f)
+                except Exception as e:
+                    logger.warning(f"Could not delete temp file {f}: {e}")
+            user_files[user_id] = []
+
         link = f"{SPACE_URL}/download/{output_filename}"
-        await update.message.reply_text(f"🎉 فایل آماده است:\n📥 {link}", reply_markup=get_menu())
-        user_files[user_id] = []
+        size_kb = os.path.getsize(output_path) // 1024
+        await update.message.reply_text(
+            f"🎉 ادغام موفق!\n"
+            f"📄 تعداد فایل‌ها: {len(files)}\n"
+            f"📦 حجم خروجی: {size_kb} KB\n\n"
+            f"📥 لینک دانلود:\n{link}",
+            reply_markup=get_menu()
+        )
+
     except Exception as e:
-        logger.error(f"Error merging PDFs: {e}")
-        await update.message.reply_text("❌ خطایی در پردازش رخ داد.", reply_markup=get_menu())
+        logger.error(f"Error merging PDFs for user {user_id}: {e}")
+        # پاک کردن فایل ناقص
+        if os.path.exists(output_path):
+            try:
+                os.remove(output_path)
+            except Exception:
+                pass
+        await update.message.reply_text(
+            "❌ خطایی در پردازش رخ داد. لطفاً دوباره امتحان کنید.",
+            reply_markup=get_menu()
+        )
 
 async def clear_queue(update: Update, context: ContextTypes.DEFAULT_TYPE):
-    user_files[update.effective_user.id] = []
-    await update.message.reply_text("🗑 صف پاک شد.", reply_markup=get_menu())
+    user_id = update.effective_user.id
+    lock = get_user_lock(user_id)
+
+    with lock:
+        files = list(user_files.get(user_id, []))
+        user_files[user_id] = []
+
+    # پاک کردن فایل‌های دانلودشده از دیسک
+    removed = 0
+    for f in files:
+        try:
+            if os.path.exists(f):
+                os.remove(f)
+                removed += 1
+        except Exception as e:
+            logger.warning(f"Could not delete {f}: {e}")
+
+    await update.message.reply_text(
+        f"🗑 صف پاک شد. ({removed} فایل حذف شد)",
+        reply_markup=get_menu()
+    )
 
 async def handle_text(update: Update, context: ContextTypes.DEFAULT_TYPE):
     text = update.message.text
-    if text == "📄 ادغام فایل‌ها": await merge_pdfs(update, context)
-    elif text == "🗑 پاک کردن صف": await clear_queue(update, context)
-    else: await update.message.reply_text("لطفاً از منو استفاده کنید.", reply_markup=get_menu())
+    if text == "📄 ادغام فایل‌ها":
+        await merge_pdfs(update, context)
+    elif text == "🗑 پاک کردن صف":
+        await clear_queue(update, context)
+    elif text == "❓ راهنما":
+        await update.message.reply_text(HELP_TEXT, reply_markup=get_menu())
+    else:
+        count = len(user_files.get(update.effective_user.id, []))
+        await update.message.reply_text(
+            f"لطفاً از منو استفاده کنید.\n📂 فایل‌های در صف: {count}",
+            reply_markup=get_menu()
+        )
+
+# ───────────────────────── Main ─────────────────────────
 
 if __name__ == "__main__":
+    if not TOKEN:
+        raise RuntimeError("متغیر محیطی TOKEN تنظیم نشده است!")
+
     threading.Thread(target=run_web, daemon=True).start()
     threading.Thread(target=cleanup_files, daemon=True).start()
-    
+
     app = Application.builder().token(TOKEN).build()
     app.add_handler(CommandHandler("start", start))
     app.add_handler(MessageHandler(filters.Document.PDF, handle_document))
     app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_text))
-    
+
     logger.info("Bot is polling...")
     app.run_polling(drop_pending_updates=True)