diff --git a/conversation_handlers.py b/conversation_handlers.py index 50a469869..1675f6f0a 100644 --- a/conversation_handlers.py +++ b/conversation_handlers.py @@ -870,6 +870,54 @@ async def start_zip_create_flow(update: Update, context: ContextTypes.DEFAULT_TY pass return ConversationHandler.END + +def _cleanup_zip_state(context: ContextTypes.DEFAULT_TYPE) -> None: + """מנקה את דגלי מצב יצירת ה‑ZIP (אחרי סיום/ביטול/בחירת שם).""" + for key in ('upload_mode', 'zip_create_items', 'awaiting_zip_name'): + context.user_data.pop(key, None) + + +async def finalize_zip_create(update: Update, context: ContextTypes.DEFAULT_TYPE, zip_name: Optional[str] = None) -> None: + """בונה ZIP מהקבצים שנאספו ושולח למשתמש. + + אם zip_name ניתן — משתמשים בו לאחר ניקוי (TextUtils.clean_filename) והבטחת סיומת .zip; + אחרת נופלים לשם ברירת המחדל my-files-.zip. + """ + items = context.user_data.get('zip_create_items') or [] + msg = update.effective_message + if not items: + if msg is not None: + await msg.reply_text("ℹ️ לא נאספו קבצים. שלח/י קבצים ואז נסה שוב.") + _cleanup_zip_state(context) + return + try: + from io import BytesIO as _BytesIO + from utils import build_zip_bytes, TextUtils + # בניית ה-ZIP (סינכרונית/כבדה) מחוץ ל-event loop — עם ניקוי שמות (Zip-Slip) ואכיפת מגבלות + zip_bytes = await asyncio.to_thread(build_zip_bytes, items) + buf = _BytesIO(zip_bytes) + buf.seek(0) + # קביעת שם ה‑ZIP: שם מהמשתמש (מנוקה) או ברירת מחדל לפי חותמת זמן + default_base = f"my-files-{datetime.now(timezone.utc).strftime('%Y%m%d-%H%M%S')}" + base = '' + if zip_name: + try: + base = (TextUtils.clean_filename(zip_name) or '').strip() + except Exception: + base = '' + if base.lower().endswith('.zip'): + base = base[:-4] + safe_name = f"{base or default_base}.zip" + await msg.reply_document(document=buf, filename=safe_name) + await msg.reply_text(f'✅ נוצר ZIP "{safe_name}" עם {len(items)} קבצים ונשלח אליך.') + except Exception as e: + logger.exception(f"finalize_zip_create failed: {e}") + if msg is not None: + await msg.reply_text(f"❌ שגיאה ביצירת ה‑ZIP: {e}") + finally: + _cleanup_zip_state(context) + + async def show_by_repo_menu(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int: """מציג תפריט קבוצות לפי תגיות ריפו ומאפשר בחירה.""" user_id = update.effective_user.id @@ -3739,9 +3787,8 @@ async def handle_callback_query(update: Update, context: ContextTypes.DEFAULT_TY pass return next_state elif data == "zip_create_cancel": - # ביטול מצב יצירת ZIP בלבד - context.user_data.pop('upload_mode', None) - context.user_data.pop('zip_create_items', None) + # ביטול מצב יצירת ZIP (כולל מצב המתנה לשם) — ניקוי מלא של הדגלים כולל awaiting_zip_name + _cleanup_zip_state(context) await query.edit_message_text("🚫 יצירת ה‑ZIP בוטלה.") await query.message.reply_text( "🎮 בחר פעולה מתקדמת:", @@ -3749,32 +3796,31 @@ async def handle_callback_query(update: Update, context: ContextTypes.DEFAULT_TY ) return ConversationHandler.END elif data == "zip_create_finish": - # בניית ZIP מהקבצים שנאספו ושליחה למשתמש - try: - items = context.user_data.get('zip_create_items') or [] - if not items: - await query.edit_message_text("ℹ️ לא נאספו קבצים. שלח/י קבצים ואז נסה שוב.") - return ConversationHandler.END - from io import BytesIO as _BytesIO - import zipfile as _zip - buf = _BytesIO() - with _zip.ZipFile(buf, 'w', compression=_zip.ZIP_DEFLATED) as z: - for it in items: - # it: {"filename": str, "bytes": bytes} - try: - z.writestr(it.get('filename') or 'file', it.get('bytes') or b'') - except Exception: - pass - buf.seek(0) - safe_name = f"my-files-{datetime.now(timezone.utc).strftime('%Y%m%d-%H%M%S')}.zip" - await query.message.reply_document(document=buf, filename=safe_name) - await query.edit_message_text(f"✅ נוצר ZIP עם {len(items)} קבצים ונשלח אליך.") - except Exception as e: - logger.exception(f"zip_create_finish failed: {e}") - await query.edit_message_text(f"❌ שגיאה ביצירת ה‑ZIP: {e}") - finally: - context.user_data.pop('upload_mode', None) - context.user_data.pop('zip_create_items', None) + # לפני יצירת ה‑ZIP — מבקשים מהמשתמש שם (או דילוג לשם אוטומטי) + items = context.user_data.get('zip_create_items') or [] + if not items: + await query.edit_message_text("ℹ️ לא נאספו קבצים. שלח/י קבצים ואז נסה שוב.") + _cleanup_zip_state(context) + return ConversationHandler.END + context.user_data['awaiting_zip_name'] = True + # לא אוספים עוד קבצים בזמן שממתינים לשם + context.user_data.pop('upload_mode', None) + kb = [ + [InlineKeyboardButton("⏭️ דלג (שם אוטומטי)", callback_data="zip_create_skip_name")], + [InlineKeyboardButton("❌ ביטול", callback_data="zip_create_cancel")], + ] + await query.edit_message_text( + "✍️ איך לקרוא ל‑ZIP?\n" + "שלח/י שם (בלי הסיומת .zip), או לחצ/י דלג לשם אוטומטי (או ביטול).\n" + f"📦 {len(items)} קבצים ייכללו.", + reply_markup=InlineKeyboardMarkup(kb), + ) + return ConversationHandler.END + elif data == "zip_create_skip_name": + # יצירת ZIP עם שם ברירת מחדל (דילוג על בחירת שם) + context.user_data.pop('awaiting_zip_name', None) + await query.edit_message_text("⏳ יוצר ZIP…") + await finalize_zip_create(update, context, zip_name=None) return ConversationHandler.END elif data.startswith("replace_") or data == "rename_file" or data == "cancel_save": return await handle_duplicate_callback(update, context) @@ -4768,7 +4814,7 @@ async def show_batch_zips_menu(update: Update, context: ContextTypes.DEFAULT_TYP end = min(start + PAGE_SIZE, total) items = backups[start:end] - lines = [f"📦 קבצי ZIP שמורים — סה""כ: {total}\n📄 עמוד {page} מתוך {total_pages}\n"] + lines = [f'📦 קבצי ZIP שמורים — סה"כ: {total}\n📄 עמוד {page} מתוך {total_pages}\n'] keyboard = [] # חישוב גרסאות vN לפי ריפו repo_to_sorted: Dict[str, list] = {} diff --git a/handlers/documents.py b/handlers/documents.py index cc6aeaae0..1a7cbe80a 100644 --- a/handlers/documents.py +++ b/handlers/documents.py @@ -839,18 +839,40 @@ async def _handle_zip_create(self, update: Update, context: ContextTypes.DEFAULT document.file_name, document.file_size, ) - file = await context.bot.get_file(document.file_id) - buf = BytesIO() - await file.download_to_memory(buf) - raw = buf.getvalue() items = context.user_data.get("zip_create_items") if items is None: items = [] context.user_data["zip_create_items"] = items + # אכיפת מגבלות מוקדמת — לפני ההורדה לזיכרון (הגנה מ-DoS/זיכרון): מספר קבצים + גודל מוצהר + from utils import ZIP_CREATE_MAX_FILES, ZIP_CREATE_MAX_TOTAL_BYTES + limit_mb = ZIP_CREATE_MAX_TOTAL_BYTES // (1024 * 1024) + if len(items) >= ZIP_CREATE_MAX_FILES: + await update.message.reply_text( + f"⚠️ הגעת למקסימום {ZIP_CREATE_MAX_FILES} קבצים ל-ZIP. לחצ/י 'סיום' כדי ליצור." + ) + return + current_total = sum(len(it.get("bytes") or b"") for it in items) + incoming_size = int(getattr(document, "file_size", 0) or 0) + if current_total + incoming_size > ZIP_CREATE_MAX_TOTAL_BYTES: + await update.message.reply_text( + f"⚠️ הקובץ לא נוסף — חריגה מהמגבלה של {limit_mb}MB לכלל ה-ZIP." + ) + return + # הורדה רק לאחר שהמגבלות המוקדמות עברו + file = await context.bot.get_file(document.file_id) + buf = BytesIO() + await file.download_to_memory(buf) + raw = buf.getvalue() + # אימות סופי לפי הגודל בפועל (למקרה של פער מול file_size המוצהר) + if current_total + len(raw) > ZIP_CREATE_MAX_TOTAL_BYTES: + await update.message.reply_text( + f"⚠️ הקובץ לא נוסף — חריגה מהמגבלה של {limit_mb}MB לכלל ה-ZIP." + ) + return safe_name = (document.file_name or f"file_{len(items)+1}").strip() or f"file_{len(items)+1}" items.append({"filename": safe_name, "bytes": raw}) await update.message.reply_text( - f"✅ נוסף: {html_escape(safe_name)} (סה""כ {len(items)} קבצים)", + f'✅ נוסף: {html_escape(safe_name)} (סה"כ {len(items)} קבצים)', parse_mode=ParseMode.HTML, ) except Exception as err: diff --git a/handlers/save_flow.py b/handlers/save_flow.py index f77698cb5..39818a4b0 100644 --- a/handlers/save_flow.py +++ b/handlers/save_flow.py @@ -323,7 +323,7 @@ async def long_collect_timeout_job(context: ContextTypes.DEFAULT_TYPE): chat_id=chat_id, text=( f"⏳ מצב איסוף נסגר לאחר חוסר פעילות.\n" - f"✅ נאספו {len(parts)} חלקים (סה""כ ~{total_bytes // 1024}KB).\n" + f'✅ נאספו {len(parts)} חלקים (סה"כ ~{total_bytes // 1024}KB).\n' f"שלח/י /done לאיחוד לקובץ אחד או /cancel לביטול." ) ) diff --git a/main.py b/main.py index 1e91a1435..7603f5bf5 100644 --- a/main.py +++ b/main.py @@ -4472,6 +4472,16 @@ async def handle_text_message(self, update: Update, context: ContextTypes.DEFAUL return text = message_text + # מצב "בחירת שם ל-ZIP" — הופעל אחרי לחיצה על 'סיום' בזרימת יצירת ZIP. + # נבדק ראשון כדי שהשם לא ייבלע/ייחשב כקוד (הטקסט נופל לכאן במצב הזה). + if context.user_data.pop('awaiting_zip_name', False): + try: + from conversation_handlers import finalize_zip_create + await finalize_zip_create(update, context, zip_name=text) + except Exception as _zip_err: + logger.exception("zip name finalize failed: %s", _zip_err) + return + # מצב חיפוש אינטראקטיבי (מופעל מהכפתור "🔎 חפש קובץ") if context.user_data.get('awaiting_search_text'): query_text = (text or '').strip() diff --git a/tests/handlers/test_documents.py b/tests/handlers/test_documents.py index 323cc7a9b..7b9501e8e 100644 --- a/tests/handlers/test_documents.py +++ b/tests/handlers/test_documents.py @@ -234,6 +234,11 @@ async def test_handle_document_collects_zip_items(handler_env): items = context.user_data.get("zip_create_items") assert items and items[0]["filename"] == "bundle.txt" assert replies.messages, "צפויה הודעה על הוספת הפריט ל-ZIP" + # הגנה מפני באג f-string השבור (סה""כ): המספר חייב להופיע בפועל, לא כטקסט מילולי + reply_text = replies.messages[0][0] + assert "1 קבצים" in reply_text, "ההודעה חייבת להציג את מספר הקבצים בפועל" + assert 'סה"כ' in reply_text, "הטקסט 'סה\"כ' חייב להופיע תקין" + assert "{len(items)}" not in reply_text, "אסור ש-{len(items)} יופיע כטקסט מילולי" @pytest.mark.asyncio diff --git a/tests/test_zip_bundle_utils.py b/tests/test_zip_bundle_utils.py new file mode 100644 index 000000000..10ec7e70b --- /dev/null +++ b/tests/test_zip_bundle_utils.py @@ -0,0 +1,88 @@ +"""טסטים לכלי בניית ה-ZIP המרוכז ב-utils (הגנת Zip-Slip + מגבלות איסוף). + +חשוב אבטחתית: safe_zip_entry_name מנרמל שמות רשומה לשם בסיס בטוח (בלי נתיב מוחלט/ +מקונן/"..") ו-build_zip_bytes אוכף מגבלת מספר קבצים וגודל מצטבר. נבדק ע"י פענוח ה-ZIP. +""" + +import zipfile +from io import BytesIO + +from utils import ( + safe_zip_entry_name, + build_zip_bytes, + ZIP_CREATE_MAX_FILES, + ZIP_CREATE_MAX_TOTAL_BYTES, +) + + +def _names(zip_bytes): + with zipfile.ZipFile(BytesIO(zip_bytes)) as z: + return z.namelist() + + +def test_safe_zip_entry_name_strips_paths(): + # basename בלבד — שום רכיב נתיב לא שורד + assert safe_zip_entry_name("../../etc/passwd") == "passwd" + assert safe_zip_entry_name("/abs/path/x.py") == "x.py" + assert safe_zip_entry_name("a\\b\\c.txt") == "c.txt" + assert safe_zip_entry_name("plain.md") == "plain.md" + # שמות מסוכנים "טהורים" נופלים ל-fallback + assert safe_zip_entry_name("..", fallback="f") == "f" + assert safe_zip_entry_name(".", fallback="f") == "f" + assert safe_zip_entry_name("", fallback="f") == "f" + assert safe_zip_entry_name(None, fallback="f") == "f" + + +def test_build_zip_sanitizes_entry_names(): + items = [ + {"filename": "../../etc/passwd", "bytes": b"x"}, + {"filename": "/abs/evil.sh", "bytes": b"y"}, + {"filename": "ok.txt", "bytes": b"z"}, + ] + names = _names(build_zip_bytes(items)) + # אף רשומה ללא מפריד נתיב או ".." + assert all("/" not in n and "\\" not in n and ".." not in n for n in names) + assert not any(n.startswith("/") for n in names) + assert set(names) == {"passwd", "evil.sh", "ok.txt"} + + +def test_build_zip_enforces_max_files(): + items = [{"filename": f"f{i}.txt", "bytes": b"x"} for i in range(10)] + names = _names(build_zip_bytes(items, max_files=3)) + assert len(names) == 3 + + +def test_build_zip_enforces_max_total_bytes(): + items = [ + {"filename": "a.bin", "bytes": b"0" * 600}, + {"filename": "b.bin", "bytes": b"1" * 600}, # 600+600=1200 > 1000 ⇒ נעצר לפני + {"filename": "c.bin", "bytes": b"2" * 600}, + ] + names = _names(build_zip_bytes(items, max_total_bytes=1000)) + assert names == ["a.bin"] + + +def test_build_zip_skips_missing_names_with_fallback(): + items = [ + {"filename": None, "bytes": b"content"}, # שם חסר ⇒ fallback file_1 + {"bytes": b"nofield"}, # ללא מפתח filename ⇒ fallback file_2 + ] + names = _names(build_zip_bytes(items)) + assert names == ["file_1", "file_2"] + + +def test_build_zip_disambiguates_duplicate_names(): + # מקורות שונים שמנורמלים לאותו שם — לא נדרסים; הבאים מקבלים סיומת ממספרת עם שמירת הסיומת + items = [ + {"filename": "a/x.txt", "bytes": b"1"}, + {"filename": "b/x.txt", "bytes": b"2"}, + {"filename": "c/x.txt", "bytes": b"3"}, + ] + names = _names(build_zip_bytes(items)) + assert names == ["x.txt", "x_2.txt", "x_3.txt"] + assert len(set(names)) == 3 + + +def test_defaults_are_sane(): + assert ZIP_CREATE_MAX_FILES == 50 + assert ZIP_CREATE_MAX_TOTAL_BYTES == 45 * 1024 * 1024 diff --git a/utils.py b/utils.py index f883c241a..d54aa4bfe 100644 --- a/utils.py +++ b/utils.py @@ -1580,3 +1580,67 @@ def _should_keep(ch: str) -> bool: except Exception: # במקרה של שגיאה, החזר את הטקסט המקורי return text + + +# ----- כלי יצירת ZIP מרוכז (בטוח מפני Zip-Slip + חסום-מגבלות, טהור וניתן להרצה ב-thread) ----- + +# מגבלות זרימת "יצירת ZIP" בבוט (מספר קבצים וגודל מצטבר) — הגנה מפני צריכת זיכרון/DoS +ZIP_CREATE_MAX_FILES = 50 +ZIP_CREATE_MAX_TOTAL_BYTES = 45 * 1024 * 1024 # 45MB (מתחת למגבלת שליחת document של טלגרם) + + +def safe_zip_entry_name(name, fallback: str = "file") -> str: + """מחזיר שם רשומת ZIP בטוח: basename בלבד, ללא נתיב מוחלט/מקונן/‏‎".."/".‎"; אחרת fallback. + + מונע Zip-Slip: שמות כמו '../../etc/passwd' או '/abs/x' מנורמלים לשם בסיס בטוח. + """ + raw = str(name or "").replace("\\", "/") + base = os.path.basename(raw) # מסיר כל רכיב נתיב (absolute/nested) + cleaned = TextUtils.clean_filename(base) # מסיר תווים אסורים + נקודות מובילות/סוגרות + if not cleaned or cleaned in (".", ".."): + return fallback + return cleaned + + +def build_zip_bytes(items, *, max_files: int = ZIP_CREATE_MAX_FILES, + max_total_bytes: int = ZIP_CREATE_MAX_TOTAL_BYTES) -> bytes: + """בונה ZIP (סינכרוני, טהור) מרשימת פריטים [{'filename': str, 'bytes': bytes}]. + + - מנקה כל שם רשומה דרך safe_zip_entry_name (הגנת Zip-Slip). + - אוכף מגבלת מספר קבצים וגודל מצטבר (הגנה כפולה מעבר לאיסוף). + - שמות רשומה כפולים מקבלים סיומת ממספרת (x.txt, x_2.txt) לשמירת ייחודיות. + - מדלג בשקט על פריט בודד שנכשל (שומר על ההתנהגות הקיימת). + + מיועד להרצה תחת asyncio.to_thread כדי לא לחסום את לולאת האירועים. + """ + from io import BytesIO + buf = BytesIO() + total = 0 + count = 0 + used = set() + with zipfile.ZipFile(buf, "w", compression=zipfile.ZIP_DEFLATED) as z: + for i, it in enumerate(items or []): + if count >= max_files: + break + try: + data = it.get("bytes") or b"" + if total + len(data) > max_total_bytes: + break + entry = safe_zip_entry_name(it.get("filename"), fallback=f"file_{i + 1}") + # מניעת שמות כפולים: משמרים את הראשון, ולבאים מוסיפים סיומת ממספרת (עם שמירת הסיומת) + if entry in used: + stem, ext = os.path.splitext(entry) + n = 2 + entry = f"{stem}_{n}{ext}" + while entry in used: + n += 1 + entry = f"{stem}_{n}{ext}" + used.add(entry) + z.writestr(entry, data) + total += len(data) + count += 1 + except Exception: + # שמירה על ההתנהגות הקיימת: פריט בעייתי מדולג ולא מפיל את כל ה-ZIP + continue + buf.seek(0) + return buf.getvalue()