| import os |
| import time |
| from pyrogram import Client, filters |
| from pyrogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton, CallbackQuery |
| from helper.database import codeflixbots as db |
| from utils.ffmpeg import get_media_info, extract_stream |
| import math |
|
|
| EXTRACT_STATE = {} |
|
|
| def humanbytes(size): |
| if not size: return "" |
| power = 2**10 |
| n = 0 |
| dic_powerN = {0: ' ', 1: 'K', 2: 'M', 3: 'G', 4: 'T'} |
| while size > power: |
| size /= power |
| n += 1 |
| return f"{size:.2f} {dic_powerN[n]}B" |
|
|
| async def extract_progress(current, total, text, message, start_time): |
| now = time.time() |
| if not hasattr(extract_progress, "last_update_time"): |
| extract_progress.last_update_time = 0 |
| |
| diff = now - start_time |
| if round(now - extract_progress.last_update_time) < 3 and current != total: |
| return |
| |
| extract_progress.last_update_time = now |
| |
| speed = current / diff if diff > 0 else 0 |
| percent = (current / total) * 100 if total > 0 else 0 |
| |
| bar_length = 15 |
| filled_length = int(bar_length * current / total) if total > 0 else 0 |
| bar = "█" * filled_length + "░" * (bar_length - filled_length) |
| |
| eta = round((total - current) / speed) if speed > 0 and total > 0 else 0 |
| eta_text = f"{eta}s" if eta > 0 else "0s" |
| |
| msg_text = f"**{text}**\n\n[{bar}] {percent:.1f}%\n" |
| msg_text += f"{humanbytes(current)} of {humanbytes(total)}\n" |
| msg_text += f"Speed: {humanbytes(speed)}/s\nETA: {eta_text}" |
| |
| try: |
| await message.edit_text(msg_text) |
| except Exception: |
| pass |
|
|
| async def initiate_extraction(client: Client, message: Message, source_message=None, url=None): |
| user_id = message.from_user.id |
| source_message = source_message or message |
| m_reply = await message.reply_text("📥 **Downloading file to fetch streams...**\nPlease wait.") |
| os.makedirs("temp", exist_ok=True) |
| start_time = time.time() |
|
|
| if url: |
| from plugins.file_rename import resolve_remote_filename |
| from utils.downloader import download_url_with_progress |
| try: |
| orig_name = await resolve_remote_filename(url, fallback="downloaded_file.mkv") |
| orig_name = os.path.basename(orig_name.replace("\\", "/")) |
| if not os.path.splitext(orig_name)[1]: |
| orig_name += ".mkv" |
| file_path = f"temp/{user_id}_{time.time()}_{orig_name}" |
|
|
| async def progress(current, total): |
| await extract_progress(current, total, "📥 **Downloading File for Extraction...**", m_reply, start_time) |
|
|
| await download_url_with_progress(url, file_path, progress_callback=progress) |
| except Exception as exc: |
| try: |
| if "file_path" in locals() and os.path.exists(file_path): |
| os.remove(file_path) |
| except Exception: |
| pass |
| await m_reply.edit_text(f"❌ Download failed: `{exc}`") |
| return |
| else: |
| media = source_message.document or source_message.video or source_message.audio |
| orig_name = getattr(media, "file_name", None) or "Media_File.mkv" |
| orig_name = os.path.basename(orig_name.replace("\\", "/")) |
| file_path = f"temp/{user_id}_{time.time()}_{orig_name}" |
| try: |
| await client.download_media( |
| source_message, file_name=file_path, progress=extract_progress, |
| progress_args=("📥 **Downloading File for Extraction...**", m_reply, start_time) |
| ) |
| except Exception as exc: |
| await m_reply.edit_text(f"❌ Download failed: `{exc}`") |
| return |
| |
| await m_reply.edit_text("🔍 **Scanning video for audio and subtitle streams...**") |
| try: |
| info = await get_media_info(file_path) |
| except Exception as exc: |
| try: |
| os.remove(file_path) |
| except Exception: |
| pass |
| await m_reply.edit_text(f"❌ Could not scan this file for audio/subtitle tracks: `{exc}`") |
| return |
| streams = info.get("streams", []) |
| |
| valid_streams = [] |
| |
| for s in streams: |
| codec_type = s.get("codec_type") |
| if codec_type in ["audio", "subtitle"]: |
| valid_streams.append(s) |
| |
| if not valid_streams: |
| await m_reply.edit_text("❌ No extractable **Audio** or **Subtitle** streams found in this video.") |
| try: os.remove(file_path) |
| except Exception: pass |
| return |
| |
| EXTRACT_STATE[user_id] = { |
| "file": file_path, |
| "original_filename": orig_name, |
| "streams": valid_streams, |
| "selected": set(), |
| "msg_id": m_reply.id, |
| "source_msg": source_message |
| } |
| |
| markup = build_extract_markup(user_id) |
| await m_reply.edit_text("✂️ **Extract Mode Active**\n\nSelect the tracks you want to extract from the list below, then click **Finish & Extract**.", reply_markup=markup) |
|
|
|
|
| @Client.on_message(filters.private & filters.command("extract")) |
| async def extract_command(client: Client, message: Message): |
| user_id = message.from_user.id |
| if not await db.is_user_authorized(user_id): |
| return await message.reply_text("❌ You are not authorized to use this bot.") |
| reply = message.reply_to_message |
| url = message.command[1].strip() if len(message.command) > 1 else None |
| if not url and reply and reply.text and reply.text.strip().lower().startswith(("http://", "https://")): |
| url = reply.text.strip() |
| if url and not url.lower().startswith(("http://", "https://")): |
| return await message.reply_text("Usage: `/extract https://direct-media-link` or reply to a Telegram media file with `/extract`.") |
| if not url and not (reply and (reply.document or reply.video or reply.audio)): |
| return await message.reply_text("Usage: `/extract https://direct-media-link` or reply to a video/audio/document with `/extract`.") |
| await initiate_extraction(client, message, source_message=reply, url=url) |
|
|
| def build_extract_markup(user_id): |
| state = EXTRACT_STATE.get(user_id) |
| if not state: return None |
| |
| buttons = [] |
| selected = state["selected"] |
| |
| for i, s in enumerate(state["streams"]): |
| codec_type = s.get("codec_type", "unknown").capitalize() |
| lang = s.get("tags", {}).get("language", s.get("tags", {}).get("LANGUAGE", "und")).upper() |
| title = s.get("tags", {}).get("title", s.get("tags", {}).get("TITLE", f"Track {s.get('index')}")) |
| |
| display_text = f"{codec_type} | {lang} | {title}" |
| if i in selected: |
| display_text = f"✅ {display_text}" |
| else: |
| display_text = f"❌ {display_text}" |
| |
| buttons.append([InlineKeyboardButton(display_text, callback_data=f"ext_toggle_{i}")]) |
| |
| buttons.append([InlineKeyboardButton("🚀 Finish & Extract", callback_data="ext_process")]) |
| buttons.append([InlineKeyboardButton("🗑️ Cancel", callback_data="ext_cancel")]) |
| |
| return InlineKeyboardMarkup(buttons) |
|
|
| @Client.on_callback_query(filters.regex(r"^ext_toggle_(\d+)")) |
| async def extract_toggle_cb(client: Client, query: CallbackQuery): |
| user_id = query.from_user.id |
| idx = int(query.matches[0].group(1)) |
| |
| state = EXTRACT_STATE.get(user_id) |
| if not state: |
| await query.answer("State lost.", show_alert=True) |
| return |
| |
| if idx in state["selected"]: |
| state["selected"].remove(idx) |
| await query.answer("Removed from selection.") |
| else: |
| state["selected"].add(idx) |
| await query.answer("Added to selection.") |
| |
| markup = build_extract_markup(user_id) |
| await query.message.edit_reply_markup(markup) |
|
|
| @Client.on_callback_query(filters.regex(r"^ext_cancel$")) |
| async def extract_cancel_cb(client: Client, query: CallbackQuery): |
| user_id = query.from_user.id |
| state = EXTRACT_STATE.get(user_id) |
| if state: |
| try: os.remove(state["file"]) |
| except Exception: pass |
| del EXTRACT_STATE[user_id] |
| |
| await query.message.edit_text("❌ Extraction cancelled.") |
|
|
| @Client.on_callback_query(filters.regex(r"^ext_process$")) |
| async def extract_process_cb(client: Client, query: CallbackQuery): |
| user_id = query.from_user.id |
| state = EXTRACT_STATE.get(user_id) |
| if not state: |
| await query.answer("State lost.", show_alert=True) |
| return |
| |
| selected = state["selected"] |
| if not selected: |
| await query.answer("Please select at least 1 track!", show_alert=True) |
| return |
| |
| await query.message.edit_text("⏳ **Extracting selected tracks...**\nProcessing via FFmpeg, please wait...") |
| |
| file_path = state["file"] |
| streams = state["streams"] |
| source_msg = state["source_msg"] |
| orig_name = state.get("original_filename", "Media_File") |
| base_name, _ = os.path.splitext(orig_name) |
| |
| extracted_files = [] |
| |
| for idx in selected: |
| s = streams[idx] |
| codec_type = s.get("codec_type") |
| codec_name = s.get("codec_name", "") |
| lang = s.get("tags", {}).get("language", s.get("tags", {}).get("LANGUAGE", "und")).upper() |
| stream_index = s.get("index") |
| |
| |
| ext = ".mka" if codec_type == "audio" else ".srt" |
| if codec_type == "audio": |
| if "aac" in codec_name: ext = ".m4a" |
| elif "mp3" in codec_name: ext = ".mp3" |
| elif "ac3" in codec_name: ext = ".ac3" |
| elif "flac" in codec_name: ext = ".flac" |
| elif codec_type == "subtitle": |
| if "ass" in codec_name: ext = ".ass" |
| elif "vtt" in codec_name: ext = ".vtt" |
| elif "subrip" in codec_name: ext = ".srt" |
| |
| track_ty_label = "Audio" if codec_type == "audio" else "Sub" |
| target_file_name = f"{base_name} - {track_ty_label} ({lang}){ext}" |
| |
| |
| tmp_out_path = f"temp/{user_id}_{time.time()}_track_{stream_index}{ext}" |
| |
| success = await extract_stream(file_path, stream_index, tmp_out_path) |
| if success and os.path.exists(tmp_out_path): |
| final_out_path = f"temp/{user_id}_{target_file_name}" |
| os.rename(tmp_out_path, final_out_path) |
| extracted_files.append((final_out_path, target_file_name)) |
| |
| if not extracted_files: |
| await query.message.edit_text("❌ Failed to extract any tracks. Format might not be supported directly.") |
| else: |
| await query.message.edit_text("✅ Extraction successful! Uploading individual tracks now...") |
| for file_to_send, custom_file_name in extracted_files: |
| start_time = time.time() |
| try: |
| await client.send_document( |
| chat_id=query.message.chat.id, |
| document=file_to_send, |
| file_name=custom_file_name, |
| caption=f"📁 **Extracted Track:** `{custom_file_name}`", |
| reply_to_message_id=source_msg.id, |
| progress=extract_progress, |
| progress_args=(f"📤 **Uploading {custom_file_name}...**", query.message, start_time) |
| ) |
| except Exception as e: |
| print(f"Failed to send extracted track: {e}") |
| try: os.remove(file_to_send) |
| except Exception: pass |
| |
| await query.message.edit_text("🎉 **All selected tracks have been extracted and uploaded successfully!**") |
| |
| try: os.remove(file_path) |
| except Exception: pass |
| if user_id in EXTRACT_STATE: |
| del EXTRACT_STATE[user_id] |
|
|
