diff --git a/config.py b/config.py index fcefa545..6ca88de8 100644 --- a/config.py +++ b/config.py @@ -1,78 +1,112 @@ -# Copyright (c) 2025 devgagan : https://github.com/devgaganin. -# Licensed under the GNU General Public License v3.0. -# See LICENSE file in the repository root for full license text. - -import os -from dotenv import load_dotenv -load_dotenv() - -# ════════════════════════════════════════════════════════════════════════════════ -# ░ CONFIGURATION SETTINGS -# ════════════════════════════════════════════════════════════════════════════════ - -# VPS --- FILL COOKIES 🍪 in """ ... """ -INST_COOKIES = """ -# write up here insta cookies -""" - -YTUB_COOKIES = """ -# write here yt cookies -""" - -# ─── BOT / DATABASE CONFIG ────────────────────────────────────────────────────── -API_ID = os.getenv("API_ID", "") -API_HASH = os.getenv("API_HASH", "") -BOT_TOKEN = os.getenv("BOT_TOKEN", "") -MONGO_DB = os.getenv("MONGO_DB", "") -DB_NAME = os.getenv("DB_NAME", "telegram_downloader") - -# ─── OWNER / CONTROL SETTINGS ─────────────────────────────────────────────────── -OWNER_ID = list(map(int, os.getenv("OWNER_ID", "").split())) # space-separated list -STRING = os.getenv("STRING", None) # optional session string -LOG_GROUP = int(os.getenv("LOG_GROUP", "-1001234456")) -FORCE_SUB = int(os.getenv("FORCE_SUB", "-10012345567")) - -# ─── SECURITY KEYS ────────────────────────────────────────────────────────────── -MASTER_KEY = os.getenv("MASTER_KEY", "gK8HzLfT9QpViJcYeB5wRa3DmN7P2xUq") # session encryption -IV_KEY = os.getenv("IV_KEY", "s7Yx5CpVmE3F") # decryption key - -# ─── COOKIES HANDLING ─────────────────────────────────────────────────────────── -YT_COOKIES = os.getenv("YT_COOKIES", YTUB_COOKIES) -INSTA_COOKIES = os.getenv("INSTA_COOKIES", INST_COOKIES) - -# ─── USAGE LIMITS ─────────────────────────────────────────────────────────────── -FREEMIUM_LIMIT = int(os.getenv("FREEMIUM_LIMIT", "0")) -PREMIUM_LIMIT = int(os.getenv("PREMIUM_LIMIT", "500")) - -# ─── UI / LINKS ───────────────────────────────────────────────────────────────── -JOIN_LINK = os.getenv("JOIN_LINK", "https://t.me/team_spy_pro") -ADMIN_CONTACT = os.getenv("ADMIN_CONTACT", "https://t.me/username_of_admin") - -# ════════════════════════════════════════════════════════════════════════════════ -# ░ PREMIUM PLANS CONFIGURATION -# ════════════════════════════════════════════════════════════════════════════════ - -P0 = { - "d": { - "s": int(os.getenv("PLAN_D_S", 1)), - "du": int(os.getenv("PLAN_D_DU", 1)), - "u": os.getenv("PLAN_D_U", "days"), - "l": os.getenv("PLAN_D_L", "Daily"), - }, - "w": { - "s": int(os.getenv("PLAN_W_S", 3)), - "du": int(os.getenv("PLAN_W_DU", 1)), - "u": os.getenv("PLAN_W_U", "weeks"), - "l": os.getenv("PLAN_W_L", "Weekly"), - }, - "m": { - "s": int(os.getenv("PLAN_M_S", 5)), - "du": int(os.getenv("PLAN_M_DU", 1)), - "u": os.getenv("PLAN_M_U", "month"), - "l": os.getenv("PLAN_M_L", "Monthly"), - }, -} - -# ════════════════════════════════════════════════════════════════════════════════ -# ░ DEVGAGAN -# ════════════════════════════════════════════════════════════════════════════════ +import os +from dotenv import load_dotenv + +load_dotenv() + + +def _getenv(*names, default=""): + for name in names: + value = os.getenv(name) + if value is not None and str(value).strip() != "": + return str(value).strip() + return default + + +def _getint(*names, default=0): + value = _getenv(*names, default=str(default)) + try: + return int(value) + except Exception: + return default + + +def _get_owner_ids(): + raw = _getenv("OWNER_ID", "ADMIN_ID", default="") + if not raw: + return [] + if "," in raw: + parts = [part.strip() for part in raw.split(",")] + else: + parts = [part.strip() for part in raw.split()] + ids = [] + for part in parts: + if not part: + continue + try: + ids.append(int(part)) + except Exception: + pass + return ids + + +INST_COOKIES = """ +# write up here insta cookies +""" + +YTUB_COOKIES = """ +# write here yt cookies +""" + + +# Core bot/database config +API_ID = _getenv("API_ID", "TG_API_ID", default="") +API_HASH = _getenv("API_HASH", "TG_API_HASH", default="") +BOT_TOKEN = _getenv("BOT_TOKEN", "TG_BOT_TOKEN", default="") +MONGO_DB = _getenv("MONGO_DB", default="") +DB_NAME = _getenv("DB_NAME", default="telegram_downloader") +STORAGE_CHANNEL_ID = _getint("TG_STORAGE_CHANNEL", "LOG_GROUP", default=0) +BACKUP_CHANNEL_ID = _getint("TG_BACKUP_CHANNEL", default=0) +TEMP_DOWNLOAD_DIR = _getenv( + "TEMP_DIR", + default="/dev/shm/telegram_vault_tmp" if os.path.isdir("/dev/shm") else "/tmp/telegram_vault_tmp" +) + + +# Owner / control settings +OWNER_ID = _get_owner_ids() +STRING = _getenv("STRING", default=None) +LOG_GROUP = _getint("LOG_GROUP", "TG_STORAGE_CHANNEL", "TG_BACKUP_CHANNEL", default=0) +FORCE_SUB = _getint("FORCE_SUB", "CHANNEL_ID", default=0) + + +# Security keys +MASTER_KEY = _getenv("MASTER_KEY", "ENCRYPTION_KEY", default="gK8HzLfT9QpViJcYeB5wRa3DmN7P2xUq") +IV_KEY = _getenv("IV_KEY", "DB_PASSWORD", default="s7Yx5CpVmE3F") + + +# Cookies +YT_COOKIES = _getenv("YT_COOKIES", default=YTUB_COOKIES) +INSTA_COOKIES = _getenv("INSTA_COOKIES", default=INST_COOKIES) + + +# Limits +FREEMIUM_LIMIT = _getint("FREEMIUM_LIMIT", default=0) +PREMIUM_LIMIT = _getint("PREMIUM_LIMIT", default=500) + + +# UI / links +JOIN_LINK = _getenv("JOIN_LINK", default="https://t.me/team_spy_pro") +ADMIN_CONTACT = _getenv("ADMIN_CONTACT", default="https://t.me/username_of_admin") + + +# Premium plan config +P0 = { + "d": { + "s": _getint("PLAN_D_S", default=1), + "du": _getint("PLAN_D_DU", default=1), + "u": _getenv("PLAN_D_U", default="days"), + "l": _getenv("PLAN_D_L", default="Daily"), + }, + "w": { + "s": _getint("PLAN_W_S", default=3), + "du": _getint("PLAN_W_DU", default=1), + "u": _getenv("PLAN_W_U", default="weeks"), + "l": _getenv("PLAN_W_L", default="Weekly"), + }, + "m": { + "s": _getint("PLAN_M_S", default=5), + "du": _getint("PLAN_M_DU", default=1), + "u": _getenv("PLAN_M_U", default="month"), + "l": _getenv("PLAN_M_L", default="Monthly"), + }, +} diff --git a/main.py b/main.py index edac6aac..c4e218b8 100644 --- a/main.py +++ b/main.py @@ -2,15 +2,17 @@ # Licensed under the GNU General Public License v3.0. # See LICENSE file in the repository root for full license text. -import asyncio -from shared_client import start_client -import importlib -import os -import sys - -async def load_and_run_plugins(): - await start_client() - plugin_dir = "plugins" +import asyncio +from shared_client import start_client +import importlib +import os +import sys +from utils.func import ensure_vault_indexes + +async def load_and_run_plugins(): + await start_client() + await ensure_vault_indexes() + plugin_dir = "plugins" plugins = [f[:-3] for f in os.listdir(plugin_dir) if f.endswith(".py") and f != "__init__.py"] for plugin in plugins: diff --git a/plugins/batch.py b/plugins/batch.py index 59fc23fb..d07d6902 100644 --- a/plugins/batch.py +++ b/plugins/batch.py @@ -6,10 +6,11 @@ from pyrogram import Client, filters from pyrogram.types import Message from pyrogram.errors import UserNotParticipant -from config import API_ID, API_HASH, LOG_GROUP, STRING, FORCE_SUB, FREEMIUM_LIMIT, PREMIUM_LIMIT -from utils.func import get_user_data, screenshot, thumbnail, get_video_metadata -from utils.func import get_user_data_key, process_text_with_rules, is_premium_user, E -from shared_client import app as X +from config import API_ID, API_HASH, LOG_GROUP, STRING, FORCE_SUB, FREEMIUM_LIMIT, PREMIUM_LIMIT, STORAGE_CHANNEL_ID, OWNER_ID +from utils.func import get_user_data, screenshot, thumbnail, get_video_metadata +from utils.func import get_user_data_key, process_text_with_rules, is_premium_user, E +from utils.func import create_vault_collection, add_vault_file, cache_source_file, get_cached_source_file, get_vault_file_by_hash, compute_file_hash +from shared_client import app as X, userbot as Y from plugins.settings import rename_file from plugins.start import subscribe as sub from utils.custom_filters import login_in_progress @@ -17,8 +18,7 @@ from typing import Dict, Any, Optional -Y = None if not STRING else __import__('shared_client').userbot -Z, P, UB, UC, emp = {}, {}, {}, {}, {} +Z, P, UB, UC, emp = {}, {}, {}, {}, {} ACTIVE_USERS = {} ACTIVE_USERS_FILE = "active_users.json" @@ -75,7 +75,23 @@ async def remove_active_batch(user_id: int): def get_batch_info(user_id: int) -> Optional[Dict[str, Any]]: return ACTIVE_USERS.get(str(user_id)) -ACTIVE_USERS = load_active_users() +ACTIVE_USERS = load_active_users() + + +def parse_source_input(text): + raw = text.strip() + i, d, lt = E(raw) + if i and d: + return str(i), int(d), lt, 1 + + m = re.match(r'^(-?\d+)\s+(\d+)(?:\s+(\d+))?$', raw) + if not m: + return None + + chat_id = m.group(1) + msg_id = int(m.group(2)) + count = int(m.group(3)) if m.group(3) else 1 + return chat_id, msg_id, 'private', count async def upd_dlg(c): try: @@ -167,12 +183,13 @@ async def get_msg(c, u, i, d, lt): return None -async def get_ubot(uid): - bt = await get_user_data_key(uid, "bot_token", None) - if not bt: return None - if uid in UB: return UB.get(uid) - try: - bot = Client(f"user_{uid}", bot_token=bt, api_id=API_ID, api_hash=API_HASH) +async def get_ubot(uid): + bt = await get_user_data_key(uid, "bot_token", None) + if not bt: + return X if uid in OWNER_ID else None + if uid in UB: return UB.get(uid) + try: + bot = Client(f"user_{uid}", bot_token=bt, api_id=API_ID, api_hash=API_HASH) await bot.start() UB[uid] = bot return bot @@ -180,25 +197,25 @@ async def get_ubot(uid): print(f"Error starting bot for user {uid}: {e}") return None -async def get_uclient(uid): - ud = await get_user_data(uid) - ubot = UB.get(uid) - cl = UC.get(uid) - if cl: return cl - if not ud: return ubot if ubot else None - xxx = ud.get('session_string') - if xxx: - try: - ss = dcs(xxx) - gg = Client(f'{uid}_client', api_id=API_ID, api_hash=API_HASH, device_model="v3saver", session_string=ss) +async def get_uclient(uid): + ud = await get_user_data(uid) + cl = UC.get(uid) + if cl: return cl + if not ud: + return Y + xxx = ud.get('session_string') + if xxx: + try: + ss = dcs(xxx) + gg = Client(f'{uid}_client', api_id=API_ID, api_hash=API_HASH, device_model="v3saver", session_string=ss) await gg.start() - await upd_dlg(gg) - UC[uid] = gg - return gg - except Exception as e: - print(f'User client error: {e}') - return ubot if ubot else Y - return Y + await upd_dlg(gg) + UC[uid] = gg + return gg + except Exception as e: + print(f'User client error: {e}') + return Y + return Y async def prog(c, t, C, h, m, st): global P @@ -215,7 +232,7 @@ async def prog(c, t, C, h, m, st): await C.edit_message_text(h, m, f"__**Pyro Handler...**__\n\n{bar}\n\n⚡**__Completed__**: {c_mb:.2f} MB / {t_mb:.2f} MB\n📊 **__Done__**: {p:.2f}%\n🚀 **__Speed__**: {speed:.2f} MB/s\n⏳ **__ETA__**: {eta}\n\n**__Powered by Team SPY__**") if p >= 100: P.pop(m, None) -async def send_direct(c, m, tcid, ft=None, rtmid=None): +async def send_direct(c, m, tcid, ft=None, rtmid=None): try: if m.video: await c.send_video(tcid, m.video.file_id, caption=ft, duration=m.video.duration, width=m.video.width, height=m.video.height, reply_to_message_id=rtmid) @@ -237,9 +254,229 @@ async def send_direct(c, m, tcid, ft=None, rtmid=None): return True except Exception as e: print(f'Direct send error: {e}') - return False - -async def process_msg(c, u, m, d, lt, uid, i): + return False + +async def archive_and_forward(c, m, local_file, target_chat_id, reply_to_message_id, caption_text, vault_collection, source_chat_id, deliver_now=True, content_hash=None): + file_name = os.path.basename(local_file) + file_ext = os.path.splitext(local_file)[1].lower() + is_video = file_ext in ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm', '.m4v', '.3gp', '.ogv'] + is_audio = file_ext in ['.mp3', '.wav', '.flac', '.aac', '.ogg', '.wma', '.m4a', '.opus', '.aiff', '.ac3'] + is_photo = file_ext in ['.jpg', '.jpeg', '.png', '.webp'] + file_size = os.path.getsize(local_file) + storage_msg = None + + if is_video: + mtd = await get_video_metadata(local_file) + dur, h, w = mtd['duration'], mtd['height'], mtd['width'] + th = await screenshot(local_file, dur, str(target_chat_id)) + storage_msg = await c.send_video( + STORAGE_CHANNEL_ID, + video=local_file, + caption=caption_text or None, + thumb=th, + width=w, + height=h, + duration=dur, + ) + elif is_audio: + storage_msg = await c.send_audio( + STORAGE_CHANNEL_ID, + audio=local_file, + caption=caption_text or None, + ) + elif is_photo: + storage_msg = await c.send_photo( + STORAGE_CHANNEL_ID, + photo=local_file, + caption=caption_text or None, + ) + else: + storage_msg = await c.send_document( + STORAGE_CHANNEL_ID, + document=local_file, + caption=caption_text or None, + ) + + media = storage_msg.video or storage_msg.audio or storage_msg.photo or storage_msg.document or storage_msg.voice + saved = await add_vault_file( + collection_id=vault_collection["_id"] if vault_collection else None, + source_chat_id=source_chat_id, + source_message_id=m.id, + storage_chat_id=STORAGE_CHANNEL_ID, + storage_message_id=storage_msg.id, + file_id=getattr(media, "file_id", None), + file_unique_id=getattr(media, "file_unique_id", None), + file_name=file_name, + mime_type="video/mp4" if is_video else "audio/mpeg" if is_audio else "image/jpeg" if is_photo else "application/octet-stream", + file_size=file_size, + caption=caption_text or "", + storage_mode="telegram_vault", + source_media_group_id=getattr(m, "media_group_id", None), + content_hash=content_hash, + ) + await cache_source_file(source_chat_id, m.id, saved["_id"]) + if deliver_now: + await c.copy_message(target_chat_id, STORAGE_CHANNEL_ID, storage_msg.id, reply_to_message_id=reply_to_message_id) + return storage_msg + + +async def replay_cached(c, cached_file, target_chat_id, reply_to_message_id, vault_collection=None, deliver_now=True): + if vault_collection and cached_file.get("collection_id") != vault_collection["_id"]: + # duplicate the reference into this collection by writing a new lightweight vault file doc + await add_vault_file( + collection_id=vault_collection["_id"], + source_chat_id=cached_file["source_chat_id"], + source_message_id=cached_file["source_message_id"], + storage_chat_id=cached_file["storage_chat_id"], + storage_message_id=cached_file["storage_message_id"], + file_id=cached_file["file_id"], + file_unique_id=cached_file["file_unique_id"], + file_name=cached_file["file_name"], + mime_type=cached_file["mime_type"], + file_size=cached_file["file_size"], + caption=cached_file.get("caption", ""), + storage_mode=cached_file.get("storage_mode", "telegram_vault"), + source_media_group_id=cached_file.get("source_media_group_id"), + content_hash=cached_file.get("content_hash"), + ) + if deliver_now: + await c.copy_message(target_chat_id, cached_file["storage_chat_id"], cached_file["storage_message_id"], reply_to_message_id=reply_to_message_id) + return True + + +async def run_batch_request(c, m, uid, ubot, uc, source_chat, start_msg_id, count, lt): + success = 0 + collection = None + pt = await m.reply_text('Processing batch...') + + if is_user_active(uid): + await pt.edit('Active task exists. Use /stop first.') + return + + await add_active_batch(uid, { + "total": count, + "current": 0, + "success": 0, + "cancel_requested": False, + "progress_message_id": pt.id + }) + + if STORAGE_CHANNEL_ID: + source_name = sanitize(str(source_chat)) + collection_name = f"batch_{source_name}_{int(time.time())}" + collection = await create_vault_collection(uid, collection_name) + + try: + for j in range(count): + if should_cancel(uid): + await pt.edit(f'Cancelled at {j}/{count}. Success: {success}') + break + + await update_batch_progress(uid, j, success) + mid = int(start_msg_id) + j + + try: + msg = await get_msg(ubot, uc, source_chat, mid, lt) + if msg: + res = await process_msg( + ubot, + uc, + msg, + str(m.chat.id), + lt, + uid, + source_chat, + vault_collection=collection, + deliver_now=False, + ) + if 'Done' in res or 'Copied' in res or 'Sent' in res or 'Archived' in res: + success += 1 + except Exception as e: + try: + await pt.edit(f'{j+1}/{count}: Error - {str(e)[:30]}') + except Exception: + pass + + await asyncio.sleep(10) + + suffix = "" + display_success = success + if collection: + from plugins.vault import _dedupe_files + from utils.func import get_vault_collection_files + files = _dedupe_files(await get_vault_collection_files(collection["_id"])) + display_success = len(files) + suffix = f"\n🔑 Collection key: `{collection['access_key']}`" + await m.reply_text(f'Batch Completed ✅ Success: {display_success}/{count}{suffix}') + if collection: + try: + from plugins.vault import _show_collection_page + from utils.func import get_vault_collection_files + files = _dedupe_files(await get_vault_collection_files(collection["_id"])) + if files: + await _show_collection_page(m, collection, files, page=1, edit=False) + except Exception: + pass + finally: + await remove_active_batch(uid) + + +async def run_single_request(c, m, uid, ubot, uc, source_chat, msg_id, lt): + pt = await m.reply_text('Processing...') + collection = None + if STORAGE_CHANNEL_ID: + source_name = sanitize(str(source_chat)) + collection_name = f"single_{source_name}_{int(time.time())}" + collection = await create_vault_collection(uid, collection_name) + + try: + msg = await get_msg(ubot, uc, source_chat, msg_id, lt) + if msg: + res = await process_msg( + ubot, + uc, + msg, + str(m.chat.id), + lt, + uid, + source_chat, + vault_collection=collection, + deliver_now=True, + ) + suffix = "" + if collection: + suffix = f"\n🔑 Collection key: `{collection['access_key']}`" + await pt.edit(f'1/1: {res}{suffix}') + if collection: + try: + from plugins.vault import _show_collection_page + from utils.func import get_vault_collection_files + files = await get_vault_collection_files(collection["_id"]) + if files: + await _show_collection_page(m, collection, files, page=1, edit=False) + except Exception: + pass + else: + await pt.edit('Message not found') + except Exception as e: + await pt.edit(f'Error: {str(e)[:50]}') + + +async def run_parsed_request(m, uid, parsed, force_single=False): + source_chat, start_msg_id, lt, count = parsed + ubot = await get_ubot(uid) + uc = await get_uclient(uid) + if not ubot or not uc: + await m.reply_text('Missing client setup') + return + + if force_single or count <= 1: + await run_single_request(ubot, m, uid, ubot, uc, source_chat, start_msg_id, lt) + return + + await run_batch_request(ubot, m, uid, ubot, uc, source_chat, start_msg_id, count, lt) + +async def process_msg(c, u, m, d, lt, uid, i, vault_collection=None, deliver_now=True): try: cfg_chat = await get_user_data_key(d, 'chat_id', None) tcid = d @@ -252,15 +489,27 @@ async def process_msg(c, u, m, d, lt, uid, i): else: tcid = int(cfg_chat) - if m.media: - orig_text = m.caption.markdown if m.caption else '' + if m.media: + cached_file = await get_cached_source_file(i, m.id) + if cached_file: + await replay_cached( + c, + cached_file, + tcid, + rtmid, + vault_collection=vault_collection, + deliver_now=deliver_now, + ) + return 'Done (cached).' if deliver_now else 'Archived (cached).' + + orig_text = m.caption.markdown if m.caption else '' proc_text = await process_text_with_rules(d, orig_text) user_cap = await get_user_data_key(d, 'caption', '') ft = f'{proc_text}\n\n{user_cap}' if proc_text and user_cap else user_cap if user_cap else proc_text - if lt == 'public' and not emp.get(i, False): - await send_direct(c, m, tcid, ft, rtmid) - return 'Sent directly.' + if lt == 'public' and not emp.get(i, False) and not vault_collection and deliver_now: + await send_direct(c, m, tcid, ft, rtmid) + return 'Sent directly.' st = time.time() p = await c.send_message(d, 'Downloading...') @@ -300,48 +549,107 @@ async def process_msg(c, u, m, d, lt, uid, i): ): f = await rename_file(f, d, p) - fsize = os.path.getsize(f) / (1024 * 1024 * 1024) - th = thumbnail(d) - - if fsize > 2 and Y: - st = time.time() - await c.edit_message_text(d, p.id, 'File is larger than 2GB. Using alternative method...') - await upd_dlg(Y) - mtd = await get_video_metadata(f) - dur, h, w = mtd['duration'], mtd['width'], mtd['height'] - th = await screenshot(f, dur, d) - - send_funcs = {'video': Y.send_video, 'video_note': Y.send_video_note, - 'voice': Y.send_voice, 'audio': Y.send_audio, - 'photo': Y.send_photo, 'document': Y.send_document} - - for mtype, func in send_funcs.items(): - if f.endswith('.mp4'): mtype = 'video' - if getattr(m, mtype, None): - sent = await func(LOG_GROUP, f, thumb=th if mtype == 'video' else None, - duration=dur if mtype == 'video' else None, - height=h if mtype == 'video' else None, - width=w if mtype == 'video' else None, - caption=ft if m.caption and mtype not in ['video_note', 'voice'] else None, - reply_to_message_id=rtmid, progress=prog, progress_args=(c, d, p.id, st)) - break - else: - sent = await Y.send_document(LOG_GROUP, f, thumb=th, caption=ft if m.caption else None, - reply_to_message_id=rtmid, progress=prog, progress_args=(c, d, p.id, st)) - - await c.copy_message(d, LOG_GROUP, sent.id) - os.remove(f) - await c.delete_messages(d, p.id) - - return 'Done (Large file).' + fsize = os.path.getsize(f) / (1024 * 1024 * 1024) + th = thumbnail(d) + content_hash = await compute_file_hash(f) if vault_collection and STORAGE_CHANNEL_ID else None + + if vault_collection and STORAGE_CHANNEL_ID and content_hash: + hash_hit = await get_vault_file_by_hash(content_hash) + if hash_hit: + await replay_cached( + c, + hash_hit, + tcid, + rtmid, + vault_collection=vault_collection, + deliver_now=deliver_now, + ) + await cache_source_file(i, m.id, hash_hit["_id"]) + if os.path.exists(f): + os.remove(f) + await c.delete_messages(d, p.id) + return 'Done (hash cache).' if deliver_now else 'Archived (hash cache).' + + if fsize > 2 and Y: + st = time.time() + await c.edit_message_text(d, p.id, 'File is larger than 2GB. Using alternative method...') + await upd_dlg(Y) + mtd = await get_video_metadata(f) + dur, h, w = mtd['duration'], mtd['width'], mtd['height'] + th = await screenshot(f, dur, d) + upload_chat_id = STORAGE_CHANNEL_ID if (vault_collection and STORAGE_CHANNEL_ID) else LOG_GROUP + + send_funcs = {'video': Y.send_video, 'video_note': Y.send_video_note, + 'voice': Y.send_voice, 'audio': Y.send_audio, + 'photo': Y.send_photo, 'document': Y.send_document} + + for mtype, func in send_funcs.items(): + if f.endswith('.mp4'): mtype = 'video' + if getattr(m, mtype, None): + sent = await func(upload_chat_id, f, thumb=th if mtype == 'video' else None, + duration=dur if mtype == 'video' else None, + height=h if mtype == 'video' else None, + width=w if mtype == 'video' else None, + caption=ft if m.caption and mtype not in ['video_note', 'voice'] else None, + reply_to_message_id=rtmid, progress=prog, progress_args=(c, d, p.id, st)) + break + else: + sent = await Y.send_document(upload_chat_id, f, thumb=th, caption=ft if m.caption else None, + reply_to_message_id=rtmid, progress=prog, progress_args=(c, d, p.id, st)) + + if vault_collection and STORAGE_CHANNEL_ID: + media = sent.video or sent.audio or sent.photo or sent.document or sent.voice + saved = await add_vault_file( + collection_id=vault_collection["_id"], + source_chat_id=i, + source_message_id=m.id, + storage_chat_id=STORAGE_CHANNEL_ID, + storage_message_id=sent.id, + file_id=getattr(media, "file_id", None), + file_unique_id=getattr(media, "file_unique_id", None), + file_name=os.path.basename(f), + mime_type="video/mp4" if sent.video else "audio/mpeg" if sent.audio else "image/jpeg" if sent.photo else "application/octet-stream", + file_size=os.path.getsize(f), + caption=ft or "", + storage_mode="telegram_vault", + source_media_group_id=getattr(m, "media_group_id", None), + content_hash=content_hash, + ) + await cache_source_file(i, m.id, saved["_id"]) + if deliver_now: + await c.copy_message(tcid, STORAGE_CHANNEL_ID, sent.id, reply_to_message_id=rtmid) + else: + await c.copy_message(d, LOG_GROUP, sent.id) + os.remove(f) + await c.delete_messages(d, p.id) + + return 'Done (Large file).' if deliver_now else 'Archived (Large file).' - await c.edit_message_text(d, p.id, 'Uploading...') - st = time.time() - - try: - video_extensions = ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm', '.m4v', '.3gp', '.ogv'] - audio_extensions = ['.mp3', '.wav', '.flac', '.aac', '.ogg', '.wma', '.m4a', '.opus', '.aiff', '.ac3'] - file_ext = os.path.splitext(f)[1].lower() + await c.edit_message_text(d, p.id, 'Uploading...') + st = time.time() + + try: + if STORAGE_CHANNEL_ID and vault_collection: + await archive_and_forward( + c=c, + m=m, + local_file=f, + target_chat_id=tcid, + reply_to_message_id=rtmid, + caption_text=ft if m.caption else None, + vault_collection=vault_collection, + source_chat_id=i, + deliver_now=deliver_now, + content_hash=content_hash, + ) + if os.path.exists(f): + os.remove(f) + await c.delete_messages(d, p.id) + return 'Done.' if deliver_now else 'Archived.' + + video_extensions = ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm', '.m4v', '.3gp', '.ogv'] + audio_extensions = ['.mp3', '.wav', '.flac', '.aac', '.ogg', '.wma', '.m4a', '.opus', '.aiff', '.ac3'] + file_ext = os.path.splitext(f)[1].lower() if m.video or (m.document and file_ext in video_extensions): mtd = await get_video_metadata(f) dur, h, w = mtd['duration'], mtd['width'], mtd['height'] @@ -391,13 +699,13 @@ async def process_msg(c, u, m, d, lt, uid, i): return f'Error: {str(e)[:50]}' @X.on_message(filters.command(['batch', 'single'])) -async def process_cmd(c, m): - uid = m.from_user.id - cmd = m.command[0] - - if FREEMIUM_LIMIT == 0 and not await is_premium_user(uid): - await m.reply_text("This bot does not provide free servies, get subscription from OWNER") - return +async def process_cmd(c, m): + uid = m.from_user.id + cmd = m.command[0] + + if uid not in OWNER_ID and FREEMIUM_LIMIT == 0 and not await is_premium_user(uid): + await m.reply_text("This bot does not provide free servies, get subscription from OWNER") + return if await sub(c, m) == 1: return pro = await m.reply_text('Doing some checks hold on...') @@ -414,79 +722,114 @@ async def process_cmd(c, m): Z[uid] = {'step': 'start' if cmd == 'batch' else 'start_single'} await pro.edit(f'Send {"start link..." if cmd == "batch" else "link you to process"}.') -@X.on_message(filters.command(['cancel', 'stop'])) -async def cancel_cmd(c, m): +@X.on_message(filters.command(['cancel', 'stop'])) +async def cancel_cmd(c, m): uid = m.from_user.id if is_user_active(uid): if await request_batch_cancel(uid): await m.reply_text('Cancellation requested. The current batch will stop after the current download completes.') else: await m.reply_text('Failed to request cancellation. Please try again.') - else: - await m.reply_text('No active batch process found.') - -@X.on_message(filters.text & filters.private & ~login_in_progress & ~filters.command([ - 'start', 'batch', 'cancel', 'login', 'logout', 'stop', 'set', - 'pay', 'redeem', 'gencode', 'single', 'generate', 'keyinfo', 'encrypt', 'decrypt', 'keys', 'setbot', 'rembot'])) -async def text_handler(c, m): - uid = m.from_user.id - if uid not in Z: return - s = Z[uid].get('step') - x = await get_ubot(uid) - if not x: - await message.reply("Add your bot /setbot `token`") - return - - if s == 'start': - L = m.text - i, d, lt = E(L) - if not i or not d: - await m.reply_text('Invalid link format.') - Z.pop(uid, None) + else: + await m.reply_text('No active batch process found.') + + +@X.on_message(filters.regex(r"^(📥 ?开始下载|📥 ?批量下载|❌ ?取消操作|🔙 ?返回主菜单)$") & filters.private) +async def legacy_button_bridge(c, m): + uid = m.from_user.id + if uid not in OWNER_ID: + return + + text = m.text.strip() + if '取消' in text: + Z.pop(uid, None) + await remove_active_batch(uid) + await m.reply_text('Cancelled. Send link / IDs / collection key directly.') + return + + if '返回' in text: + await m.reply_text('Send link / `频道ID 消息ID [数量]` / `file_store...` directly.') + return + + Z[uid] = {'step': 'start'} + await m.reply_text('Send start link or `频道ID 消息ID 数量`.') + +@X.on_message(filters.text & filters.private & ~login_in_progress & ~filters.command([ + 'start', 'batch', 'cancel', 'login', 'logout', 'stop', 'set', + 'pay', 'redeem', 'gencode', 'single', 'generate', 'keyinfo', 'encrypt', 'decrypt', 'keys', 'setbot', 'rembot', + 'mycollections' +]), group=-5) +async def global_source_input_handler(c, m): + uid = m.from_user.id + if uid not in OWNER_ID: + return + + parsed = parse_source_input(m.text) + if not parsed: + return + + state = Z.get(uid, {}).get('step') + force_single = state == 'start_single' + Z.pop(uid, None) + await run_parsed_request(m, uid, parsed, force_single=force_single) + m.stop_propagation() + + +@X.on_message(filters.text & filters.private & ~login_in_progress & ~filters.command([ + 'start', 'batch', 'cancel', 'login', 'logout', 'stop', 'set', + 'pay', 'redeem', 'gencode', 'single', 'generate', 'keyinfo', 'encrypt', 'decrypt', 'keys', 'setbot', 'rembot', + 'mycollections' +])) +async def text_handler(c, m): + uid = m.from_user.id + if uid not in Z: return + s = Z[uid].get('step') + x = await get_ubot(uid) + if not x: + await m.reply("Add your bot /setbot `token`") + return + + if s == 'start': + L = m.text + i, d, lt = E(L) + if not i or not d: + await m.reply_text('Invalid link format.') + Z.pop(uid, None) return Z[uid].update({'step': 'count', 'cid': i, 'sid': d, 'lt': lt}) await m.reply_text('How many messages?') - - elif s == 'start_single': - L = m.text - i, d, lt = E(L) - if not i or not d: - await m.reply_text('Invalid link format.') - Z.pop(uid, None) + + elif s == 'start_single': + L = m.text + i, d, lt = E(L) + if not i or not d: + await m.reply_text('Invalid link format.') + Z.pop(uid, None) return - - Z[uid].update({'step': 'process_single', 'cid': i, 'sid': d, 'lt': lt}) - i, s, lt = Z[uid]['cid'], Z[uid]['sid'], Z[uid]['lt'] - pt = await m.reply_text('Processing...') - - ubot = UB.get(uid) - if not ubot: - await pt.edit('Add bot with /setbot first') - Z.pop(uid, None) - return - - uc = await get_uclient(uid) - if not uc: - await pt.edit('Cannot proceed without user client.') - Z.pop(uid, None) - return - - if is_user_active(uid): - await pt.edit('Active task exists. Use /stop first.') - Z.pop(uid, None) - return - - try: - msg = await get_msg(ubot, uc, i, s, lt) - if msg: - res = await process_msg(ubot, uc, msg, str(m.chat.id), lt, uid, i) - await pt.edit(f'1/1: {res}') - else: - await pt.edit('Message not found') - except Exception as e: - await pt.edit(f'Error: {str(e)[:50]}') - finally: - Z.pop(uid, None) + + Z[uid].update({'step': 'process_single', 'cid': i, 'sid': d, 'lt': lt}) + i, s, lt = Z[uid]['cid'], Z[uid]['sid'], Z[uid]['lt'] + ubot = await get_ubot(uid) + if not ubot: + await m.reply_text('Add bot with /setbot first') + Z.pop(uid, None) + return + + uc = await get_uclient(uid) + if not uc: + await m.reply_text('Cannot proceed without user client.') + Z.pop(uid, None) + return + + if is_user_active(uid): + await m.reply_text('Active task exists. Use /stop first.') + Z.pop(uid, None) + return + + try: + await run_single_request(ubot, m, uid, ubot, uc, i, s, lt) + finally: + Z.pop(uid, None) elif s == 'count': if not m.text.isdigit(): @@ -494,69 +837,42 @@ async def text_handler(c, m): return count = int(m.text) - maxlimit = PREMIUM_LIMIT if await is_premium_user(uid) else FREEMIUM_LIMIT + maxlimit = PREMIUM_LIMIT if (uid in OWNER_ID or await is_premium_user(uid)) else FREEMIUM_LIMIT if count > maxlimit: await m.reply_text(f'Maximum limit is {maxlimit}.') return - - Z[uid].update({'step': 'process', 'did': str(m.chat.id), 'num': count}) - i, s, n, lt = Z[uid]['cid'], Z[uid]['sid'], Z[uid]['num'], Z[uid]['lt'] - success = 0 - - pt = await m.reply_text('Processing batch...') - uc = await get_uclient(uid) - ubot = UB.get(uid) - - if not uc or not ubot: - await pt.edit('Missing client setup') - Z.pop(uid, None) - return - - if is_user_active(uid): - await pt.edit('Active task exists') - Z.pop(uid, None) - return - - await add_active_batch(uid, { - "total": n, - "current": 0, - "success": 0, - "cancel_requested": False, - "progress_message_id": pt.id - }) - - try: - for j in range(n): - - if should_cancel(uid): - await pt.edit(f'Cancelled at {j}/{n}. Success: {success}') - break - - await update_batch_progress(uid, j, success) - - mid = int(s) + j - - try: - msg = await get_msg(ubot, uc, i, mid, lt) - if msg: - res = await process_msg(ubot, uc, msg, str(m.chat.id), lt, uid, i) - if 'Done' in res or 'Copied' in res or 'Sent' in res: - success += 1 - else: - pass - except Exception as e: - try: await pt.edit(f'{j+1}/{n}: Error - {str(e)[:30]}') - except: pass - - await asyncio.sleep(10) - - if j+1 == n: - await m.reply_text(f'Batch Completed ✅ Success: {success}/{n}') - - finally: - await remove_active_batch(uid) - Z.pop(uid, None) + + Z[uid].update({'step': 'process', 'did': str(m.chat.id), 'num': count}) + i, s, n, lt = Z[uid]['cid'], Z[uid]['sid'], Z[uid]['num'], Z[uid]['lt'] + uc = await get_uclient(uid) + ubot = await get_ubot(uid) + + if not uc or not ubot: + await m.reply_text('Missing client setup') + Z.pop(uid, None) + return + try: + await run_batch_request(ubot, m, uid, ubot, uc, i, s, n, lt) + finally: + Z.pop(uid, None) + + +@X.on_message(filters.text & filters.private & ~login_in_progress & ~filters.command([ + 'start', 'batch', 'cancel', 'login', 'logout', 'stop', 'set', + 'pay', 'redeem', 'gencode', 'single', 'generate', 'keyinfo', 'encrypt', 'decrypt', 'keys', 'setbot', 'rembot', + 'mycollections' +])) +async def direct_input_handler(c, m): + uid = m.from_user.id + if uid not in OWNER_ID or uid in Z: + return + + parsed = parse_source_input(m.text) + if not parsed: + return + + await run_parsed_request(m, uid, parsed, force_single=False) diff --git a/plugins/start.py b/plugins/start.py index 81791507..8ce709fe 100644 --- a/plugins/start.py +++ b/plugins/start.py @@ -2,13 +2,20 @@ # Licensed under the GNU General Public License v3.0. # See LICENSE file in the repository root for full license text. -from shared_client import app -from pyrogram import filters -from pyrogram.errors import UserNotParticipant -from pyrogram.types import BotCommand, InlineKeyboardButton, InlineKeyboardMarkup -from config import LOG_GROUP, OWNER_ID, FORCE_SUB - -async def subscribe(app, message): +from shared_client import app +from pyrogram import filters +from pyrogram.errors import UserNotParticipant +from pyrogram.types import BotCommand, InlineKeyboardButton, InlineKeyboardMarkup +from config import LOG_GROUP, OWNER_ID, FORCE_SUB + +@app.on_message(filters.private, group=-10) +async def owner_only_guard(_, message): + if message.from_user and message.from_user.id not in OWNER_ID: + await message.reply_text("🔒 Private bot. Owner only.") + message.stop_propagation() + return + +async def subscribe(app, message): if FORCE_SUB: try: user = await app.get_chat_member(FORCE_SUB, message.from_user.id) @@ -21,8 +28,25 @@ async def subscribe(app, message): await message.reply_photo(photo="https://graph.org/file/d44f024a08ded19452152.jpg",caption=caption, reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Join Now...", url=f"{link}")]])) return 1 except Exception as ggn: - await message.reply_text(f"Something Went Wrong. Contact admins... with following message {ggn}") - return 1 + await message.reply_text(f"Something Went Wrong. Contact admins... with following message {ggn}") + return 1 + + +@app.on_message(filters.command("start") & filters.private) +async def start_handler(client, message): + if message.from_user.id not in OWNER_ID: + await message.reply_text("🔒 Private bot. Owner only.") + return + text = ( + "📦 **Private Save Bot**\n\n" + "直接发送以下任一内容即可：\n" + "• Telegram 链接\n" + "• `频道ID 消息ID`\n" + "• `频道ID 消息ID 数量`\n" + "• `file_store...` 合集密钥\n\n" + "也可用命令：`/batch`、`/single`、`/mycollections`" + ) + await message.reply_text(text, disable_web_page_preview=True) @app.on_message(filters.command("set")) async def set(_, message): diff --git a/plugins/vault.py b/plugins/vault.py new file mode 100644 index 00000000..48e8143c --- /dev/null +++ b/plugins/vault.py @@ -0,0 +1,330 @@ +import time +import re + +from pyrogram import filters +from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup, InputMediaPhoto, InputMediaVideo + +from shared_client import app +from config import OWNER_ID +from utils.func import ( + get_user_vault_collections, + get_vault_collection_by_key, + get_vault_collection_files, + get_vault_file_by_key, +) + +ACTIVE_VAULT_SENDS = set() +RECENT_VAULT_SENDS = {} +VAULT_SEND_COOLDOWN = 20 + + +def _is_owner(user_id: int) -> bool: + return user_id in OWNER_ID + + +def _extract_vault_key(text: str) -> str | None: + match = re.search(r"(file_store[A-Za-z0-9]+)", text or "") + return match.group(1) if match else None + + +def _is_recent_send(lock_key) -> bool: + now = time.time() + last = RECENT_VAULT_SENDS.get(lock_key) + if last and now - last < VAULT_SEND_COOLDOWN: + return True + RECENT_VAULT_SENDS.pop(lock_key, None) + return False + + +def _mark_recent_send(lock_key): + RECENT_VAULT_SENDS[lock_key] = time.time() + + +def _short_text(text: str, max_len: int = 40) -> str: + text = (text or "").replace("\n", " ").strip() + return text if len(text) <= max_len else text[: max_len - 1] + "..." + + +def _human_size(size: int) -> str: + size = size or 0 + units = ["B", "KB", "MB", "GB"] + value = float(size) + idx = 0 + while value >= 1024 and idx < len(units) - 1: + value /= 1024 + idx += 1 + return f"{value:.2f} {units[idx]}" + + +def _dedupe_files(files: list) -> list: + unique = [] + seen = set() + for file_info in files: + file_unique_id = file_info.get("file_unique_id") + if file_unique_id: + key = ("file_unique_id", file_unique_id) + else: + key = ( + "storage_message", + file_info.get("storage_chat_id"), + file_info.get("storage_message_id"), + file_info.get("file_name"), + file_info.get("file_size"), + ) + if key in seen: + continue + seen.add(key) + unique.append(file_info) + return unique + + +def _is_visual(file_info: dict) -> bool: + mime = (file_info.get("mime_type") or "").lower() + return mime.startswith("image") or mime.startswith("video") + + +async def _send_single_vault_file(chat_id: int, file_info: dict) -> bool: + try: + await app.copy_message(chat_id, file_info["storage_chat_id"], file_info["storage_message_id"]) + return True + except Exception: + pass + + try: + await app.send_cached_media(chat_id, file_info["file_id"], caption=file_info.get("caption") or "") + return True + except Exception: + return False + + +async def _send_visual_group(chat_id: int, files: list) -> bool: + media = [] + for index, file_info in enumerate(files): + mime = (file_info.get("mime_type") or "").lower() + caption = file_info.get("caption") or "" + caption = caption if index == 0 else "" + if mime.startswith("image"): + media.append(InputMediaPhoto(media=file_info["file_id"], caption=caption)) + elif mime.startswith("video"): + media.append(InputMediaVideo(media=file_info["file_id"], caption=caption)) + else: + return False + + try: + await app.send_media_group(chat_id, media) + return True + except Exception: + return False + + +async def _send_vault_files(chat_id: int, files: list, strict_visual_groups: bool = False) -> int: + files = _dedupe_files(files) + visual_files = [item for item in files if _is_visual(item)] + other_files = [item for item in files if not _is_visual(item)] + sent = 0 + + idx = 0 + while idx < len(visual_files): + group = visual_files[idx:idx + 10] + if len(group) >= 2: + if not await _send_visual_group(chat_id, group): + if strict_visual_groups: + raise RuntimeError(f"Failed to send media group chunk starting at item {idx + 1}") + current = group[0] + if await _send_single_vault_file(chat_id, current): + sent += 1 + idx += 1 + continue + sent += len(group) + idx += len(group) + continue + + current = group[0] + if strict_visual_groups: + raise RuntimeError("Only one visual file left; cannot send it as a media group") + if await _send_single_vault_file(chat_id, current): + sent += 1 + idx += 1 + + for item in other_files: + if strict_visual_groups: + raise RuntimeError("This collection page contains non-visual files; strict media-group mode refused single sends") + if await _send_single_vault_file(chat_id, item): + sent += 1 + + return sent + + +async def _show_collection_page(message, collection, files, page=1, edit=False): + files = _dedupe_files(files) + per_page = 10 + total_files = len(files) + total_pages = max(1, (total_files + per_page - 1) // per_page) + page = max(1, min(page, total_pages)) + + start_idx = (page - 1) * per_page + end_idx = start_idx + per_page + page_files = files[start_idx:end_idx] + + lines = [ + f"Folder: {collection['name']}", + f"Files: {total_files} (page {page}/{total_pages})", + "-------------------------", + ] + for item in page_files: + mime = (item.get("mime_type") or "").lower() + if mime.startswith("video"): + kind = "video" + elif mime.startswith("image"): + kind = "image" + elif mime.startswith("audio"): + kind = "audio" + else: + kind = "file" + lines.append(f"- `{_short_text(item.get('file_name') or 'unnamed file')}`") + lines.append(f" {kind} | {_human_size(item.get('file_size') or 0)}") + lines.append(f"Key: `{collection['access_key']}`") + text = "\n".join(lines) + + buttons = [ + [ + InlineKeyboardButton(f"Send This Page ({len(page_files)})", callback_data=f"vault_send_{collection['access_key']}_{page}"), + InlineKeyboardButton("Show File Codes", callback_data=f"vault_codes_{collection['access_key']}_{page}") + ] + ] + if total_files > per_page: + buttons.append([InlineKeyboardButton(f"Send All ({total_files})", callback_data=f"vault_all_{collection['access_key']}")]) + + nav = [] + if page > 1: + nav.append(InlineKeyboardButton("Prev", callback_data=f"vault_page_{collection['access_key']}_{page-1}")) + if page < total_pages: + nav.append(InlineKeyboardButton("Next", callback_data=f"vault_page_{collection['access_key']}_{page+1}")) + if nav: + buttons.append(nav) + + markup = InlineKeyboardMarkup(buttons) + if edit: + await message.edit_text(text, reply_markup=markup, disable_web_page_preview=True) + else: + await message.reply_text(text, reply_markup=markup, disable_web_page_preview=True) + + +@app.on_message(filters.command("mycollections") & filters.private) +async def mycollections_handler(_, message): + if not _is_owner(message.from_user.id): + await message.reply_text("Private bot. Owner only.") + return + + collections = await get_user_vault_collections(message.from_user.id) + if not collections: + await message.reply_text("No collections yet.") + return + + lines = ["My Collections\n"] + for col in collections[:30]: + files = _dedupe_files(await get_vault_collection_files(col["_id"])) + lines.append(f"- {col['name']}") + lines.append(f" Key: `{col['access_key']}`") + lines.append(f" Files: {len(files)}\n") + await message.reply_text("\n".join(lines)) + + +@app.on_callback_query(filters.regex(r"^vault_(page|send|all|codes)_")) +async def vault_callback_handler(_, callback): + if not _is_owner(callback.from_user.id): + await callback.answer("Private bot. Owner only.", show_alert=True) + return + + parts = callback.data.split("_") + action = parts[1] + access_key = "_".join(parts[2:-1]) if action in {"page", "send", "codes"} else "_".join(parts[2:]) + page = int(parts[-1]) if action in {"page", "send", "codes"} else 1 + + collection = await get_vault_collection_by_key(access_key) + if not collection: + await callback.answer("Collection not found.", show_alert=True) + return + files = _dedupe_files(await get_vault_collection_files(collection["_id"])) + + if action == "page": + await _show_collection_page(callback.message, collection, files, page=page, edit=True) + await callback.answer() + return + + if action == "send": + lock_key = ("send", callback.message.chat.id, access_key, page) + if lock_key in ACTIVE_VAULT_SENDS or _is_recent_send(lock_key): + await callback.answer("This page was just sent.") + return + ACTIVE_VAULT_SENDS.add(lock_key) + await callback.answer("Sending page...") + start_idx = (page - 1) * 10 + end_idx = start_idx + 10 + try: + sent = await _send_vault_files(callback.message.chat.id, files[start_idx:end_idx], strict_visual_groups=True) + await callback.message.reply_text(f"Sent page files: {sent}") + except Exception as e: + await callback.message.reply_text(f"Page send failed: {str(e)}") + finally: + ACTIVE_VAULT_SENDS.discard(lock_key) + _mark_recent_send(lock_key) + return + + if action == "codes": + start_idx = (page - 1) * 10 + end_idx = start_idx + 10 + page_files = files[start_idx:end_idx] + lines = [f"File Codes | {collection['name']} | page {page}\n"] + for idx, item in enumerate(page_files, start=1): + lines.append(f"{idx}. `{_short_text(item.get('file_name') or 'unnamed file', 28)}`") + lines.append(f" `{item['access_key']}`") + await callback.message.reply_text("\n".join(lines), disable_web_page_preview=True) + await callback.answer("File codes sent.") + return + + lock_key = ("all", callback.message.chat.id, access_key, 0) + if lock_key in ACTIVE_VAULT_SENDS or _is_recent_send(lock_key): + await callback.answer("All files were just sent.") + return + ACTIVE_VAULT_SENDS.add(lock_key) + await callback.answer("Sending all files...") + try: + sent = await _send_vault_files(callback.message.chat.id, files, strict_visual_groups=True) + await callback.message.reply_text(f"Sent all files: {sent}") + except Exception as e: + await callback.message.reply_text(f"Send all failed: {str(e)}") + finally: + ACTIVE_VAULT_SENDS.discard(lock_key) + _mark_recent_send(lock_key) + + +@app.on_message(filters.text & filters.private) +async def vault_key_handler(_, message): + if not _is_owner(message.from_user.id): + return + + key = _extract_vault_key(message.text.strip()) + if not key: + return + + collection = await get_vault_collection_by_key(key) + if collection: + files = _dedupe_files(await get_vault_collection_files(collection["_id"])) + if not files: + await message.reply_text(f"{collection['name']}\n\nNo files in this collection.") + return + await _show_collection_page(message, collection, files, page=1, edit=False) + return + + file_info = await get_vault_file_by_key(key) + if file_info: + sent = await _send_vault_files(message.chat.id, [file_info], strict_visual_groups=False) + await message.reply_text(f"Sent {sent}/1 file.") + return + + await message.reply_text(f"Key not found: `{key}`") + + +async def run_vault_plugin(): + return diff --git a/scripts/cleanup_vault_duplicates.py b/scripts/cleanup_vault_duplicates.py new file mode 100644 index 00000000..0d1cc092 --- /dev/null +++ b/scripts/cleanup_vault_duplicates.py @@ -0,0 +1,58 @@ +import os +from collections import defaultdict + +from dotenv import load_dotenv +from pymongo import MongoClient + + +def build_key(doc): + if doc.get("file_unique_id"): + return ("file_unique_id", doc.get("collection_id"), doc["file_unique_id"]) + return ( + "source_fallback", + doc.get("collection_id"), + doc.get("source_chat_id"), + doc.get("source_message_id"), + doc.get("file_name"), + doc.get("file_size"), + ) + + +def main(): + load_dotenv() + mongo_uri = os.getenv("MONGO_DB", "mongodb://127.0.0.1:27017") + db_name = os.getenv("DB_NAME", "telegram_downloader") + + client = MongoClient(mongo_uri) + db = client[db_name] + vault_files = db["vault_files"] + source_cache = db["vault_source_cache"] + + docs = list(vault_files.find().sort("created_at", 1)) + buckets = defaultdict(list) + for doc in docs: + buckets[build_key(doc)].append(doc) + + removed = 0 + updated_cache = 0 + for _, items in buckets.items(): + if len(items) < 2: + continue + + keep = items[0] + dupes = items[1:] + for dupe in dupes: + result = source_cache.update_many( + {"file_doc_id": dupe["_id"]}, + {"$set": {"file_doc_id": keep["_id"]}}, + ) + updated_cache += result.modified_count + vault_files.delete_one({"_id": dupe["_id"]}) + removed += 1 + + print(f"Removed duplicate vault files: {removed}") + print(f"Updated cache references: {updated_cache}") + + +if __name__ == "__main__": + main() diff --git a/shared_client.py b/shared_client.py index 4a2960df..9dfac3e3 100644 --- a/shared_client.py +++ b/shared_client.py @@ -2,25 +2,31 @@ # Licensed under the GNU General Public License v3.0. # See LICENSE file in the repository root for full license text. -from telethon import TelegramClient -from config import API_ID, API_HASH, BOT_TOKEN, STRING -from pyrogram import Client -import sys - -client = TelegramClient("telethonbot", API_ID, API_HASH) -app = Client("pyrogrambot", api_id=API_ID, api_hash=API_HASH, bot_token=BOT_TOKEN) -userbot = Client("4gbbot", api_id=API_ID, api_hash=API_HASH, session_string=STRING) - -async def start_client(): - if not client.is_connected(): - await client.start(bot_token=BOT_TOKEN) - print("SpyLib started...") - if STRING: - try: - await userbot.start() - print("Userbot started...") - except Exception as e: - print(f"Hey honey!! check your premium string session, it may be invalid of expire {e}") +from telethon import TelegramClient +from config import API_ID, API_HASH, BOT_TOKEN, STRING +from pyrogram import Client +import sys +import os + +client = TelegramClient("telethonbot", API_ID, API_HASH) +app = Client("pyrogrambot", api_id=API_ID, api_hash=API_HASH, bot_token=BOT_TOKEN) +if STRING: + userbot = Client("4gbbot", api_id=API_ID, api_hash=API_HASH, session_string=STRING) +elif os.path.exists("vault_user.session"): + userbot = Client("vault_user", api_id=API_ID, api_hash=API_HASH) +else: + userbot = None + +async def start_client(): + if not client.is_connected(): + await client.start(bot_token=BOT_TOKEN) + print("SpyLib started...") + if userbot: + try: + await userbot.start() + print("Userbot started...") + except Exception as e: + print(f"Hey honey!! check your premium string session, it may be invalid of expire {e}") sys.exit(1) await app.start() print("Pyro App Started...") diff --git a/utils/func.py b/utils/func.py index 4a4db951..9f50e739 100644 --- a/utils/func.py +++ b/utils/func.py @@ -9,6 +9,7 @@ import cv2 import logging import asyncio +import hashlib from datetime import datetime, timedelta from motor.motor_asyncio import AsyncIOMotorClient from config import MONGO_DB as MONGO_URI, DB_NAME @@ -26,6 +27,153 @@ premium_users_collection = db["premium_users"] statistics_collection = db["statistics"] codedb = db["redeem_code"] +vault_files_collection = db["vault_files"] +vault_collections_collection = db["vault_collections"] +vault_source_cache_collection = db["vault_source_cache"] + + +def _generate_vault_key(prefix="file_store", length=24): + import secrets + import string + chars = string.ascii_letters + string.digits + return prefix + "".join(secrets.choice(chars) for _ in range(length)) + + +async def ensure_vault_indexes(): + await vault_files_collection.create_index("access_key", unique=True) + await vault_files_collection.create_index("file_id") + await vault_files_collection.create_index("content_hash") + await vault_files_collection.create_index( + [("collection_id", 1), ("storage_chat_id", 1), ("storage_message_id", 1)], + unique=True + ) + await vault_collections_collection.create_index("access_key", unique=True) + await vault_collections_collection.create_index([("owner_id", 1), ("name", 1)], unique=True) + await vault_source_cache_collection.create_index([("source_chat_id", 1), ("source_message_id", 1)], unique=True) + + +async def create_vault_collection(owner_id, name, access_key=None): + access_key = access_key or _generate_vault_key() + doc = { + "owner_id": int(owner_id), + "name": name, + "access_key": access_key, + "created_at": datetime.now(), + } + try: + result = await vault_collections_collection.insert_one(doc) + doc["_id"] = result.inserted_id + doc["id"] = str(result.inserted_id) + return doc + except Exception as e: + logger.error(f"create_vault_collection error: {e}") + return None + + +async def get_vault_collection_by_key(access_key): + return await vault_collections_collection.find_one({"access_key": access_key}) + + +async def get_user_vault_collections(owner_id): + cursor = vault_collections_collection.find({"owner_id": int(owner_id)}).sort("created_at", -1) + return await cursor.to_list(length=200) + + +async def add_vault_file( + collection_id, + source_chat_id, + source_message_id, + storage_chat_id, + storage_message_id, + file_id, + file_unique_id, + file_name, + mime_type, + file_size, + caption="", + storage_mode="telegram_vault", + source_media_group_id=None, + content_hash=None, +): + doc = { + "collection_id": collection_id, + "source_chat_id": int(source_chat_id), + "source_message_id": int(source_message_id), + "storage_chat_id": int(storage_chat_id), + "storage_message_id": int(storage_message_id), + "file_id": file_id, + "file_unique_id": file_unique_id, + "file_name": file_name, + "mime_type": mime_type, + "file_size": file_size or 0, + "caption": caption or "", + "storage_mode": storage_mode, + "source_media_group_id": source_media_group_id, + "content_hash": content_hash, + "created_at": datetime.now(), + } + existing = await vault_files_collection.find_one( + { + "collection_id": collection_id, + "storage_chat_id": int(storage_chat_id), + "storage_message_id": int(storage_message_id), + } + ) + if existing: + return existing + + doc["access_key"] = _generate_vault_key() + result = await vault_files_collection.insert_one(doc) + doc["_id"] = result.inserted_id + return doc + + +async def get_vault_collection_files(collection_id): + cursor = vault_files_collection.find({"collection_id": collection_id}).sort("created_at", 1) + return await cursor.to_list(length=5000) + + +async def get_vault_file_by_key(access_key): + return await vault_files_collection.find_one({"access_key": access_key}) + + +async def get_vault_file_by_hash(content_hash): + if not content_hash: + return None + return await vault_files_collection.find_one( + {"content_hash": content_hash, "storage_mode": "telegram_vault"}, + sort=[("created_at", -1)] + ) + + +async def cache_source_file(source_chat_id, source_message_id, file_doc_id): + await vault_source_cache_collection.update_one( + {"source_chat_id": int(source_chat_id), "source_message_id": int(source_message_id)}, + {"$set": {"file_doc_id": file_doc_id, "updated_at": datetime.now()}}, + upsert=True + ) + + +async def get_cached_source_file(source_chat_id, source_message_id): + cache_doc = await vault_source_cache_collection.find_one( + {"source_chat_id": int(source_chat_id), "source_message_id": int(source_message_id)} + ) + if not cache_doc: + return None + return await vault_files_collection.find_one({"_id": cache_doc["file_doc_id"]}) + + +async def compute_file_hash(file_path): + loop = asyncio.get_running_loop() + + def _hash_file(): + sha256 = hashlib.sha256() + with open(file_path, "rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + sha256.update(chunk) + return sha256.hexdigest() + + return await loop.run_in_executor(None, _hash_file) # ------- < start > Session Encoder don't change -------