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") # Determine logical extension 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}" # Temp path for ffmpeg extraction MUST HAVE proper extension for muxer to work 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]