set MERGE_MODE 4 to use /extract command

SHBGM · Oct 17, 2022 · f91b0a9 · f91b0a9
1 parent b620f6f
commit f91b0a9
Show file tree

Hide file tree

Showing 8 changed files with 303 additions and 13 deletions.
diff --git a/bot.py b/bot.py
@@ -415,6 +415,40 @@ async def photo_handler(c: Client, m: Message):
     del user
 
 
+@mergeApp.on_message(filters.command(["extract"]) & filters.private)
+async def media_extracter(c: Client, m: Message):
+    user = UserSettings(uid=m.from_user.id, name=m.from_user.first_name)
+    if not user.allowed:
+        return
+    if user.merge_mode == 4:
+        if m.reply_to_message is None:
+            await m.reply(text="Reply /extract to a video or document file")
+            return
+        rmess = m.reply_to_message
+        if rmess.video or rmess.document:
+            media = rmess.video or rmess.document
+            mid=rmess.id
+            file_name = media.file_name
+            if file_name is None:
+                await m.reply("File name not found; goto @yashoswalyo")
+                return
+            markup = bMaker.makebuttons(
+                set1=["Audio", "Subtitle", "All", "Cancel"],
+                set2=[f"extract_audio_{mid}", f"extract_subtitle_{mid}", f"extract_all_{mid}", 'cancel'],
+                isCallback=True,
+                rows=2,
+            )
+            await m.reply(
+                text="Choose from below what you want to extract?",
+                quote=True,
+                reply_markup=InlineKeyboardMarkup(markup),
+            )
+    else:
+        await m.reply(
+            text="Change settings and set mode to extract\nthen use /extract command"
+        )
+
+
 @mergeApp.on_message(filters.command(["help"]) & filters.private)
 async def help_msg(c: Client, m: Message):
     await m.reply_text(

diff --git a/config.py b/config.py
@@ -4,12 +4,13 @@
 class Config(object):
     API_HASH = os.environ.get("API_HASH")
     BOT_TOKEN = os.environ.get("BOT_TOKEN")
-    TELEGRAM_API = os.environ.get("TELEGRAM_API")
+    TELEGRAM_API = os.environ["TELEGRAM_API"]
     OWNER = os.environ.get("OWNER")
     OWNER_USERNAME = os.environ.get("OWNER_USERNAME")
     PASSWORD = os.environ.get("PASSWORD")
     DATABASE_URL = os.environ.get("DATABASE_URL")
     LOGCHANNEL = os.environ.get("LOGCHANNEL")  # Add channel id as -100 + Actual ID
+    GDRIVE_FOLDER_ID = os.environ.get("GDRIVE_FOLDER_ID","root")
     USER_SESSION_STRING = os.environ.get("USER_SESSION_STRING", None)
     IS_PREMIUM = False
-    MODES = ["video-video", "video-audio", "video-subtitle"]
+    MODES = ["video-video", "video-audio", "video-subtitle","extract-streams"]
diff --git a/helpers/display_progress.py b/helpers/display_progress.py
@@ -40,7 +40,7 @@ def is_cancelled(self):
             self._cancelled = True
         return self._cancelled
 
-    async def progress_for_pyrogram(self, current, total, ud_type, start):
+    async def progress_for_pyrogram(self, current, total, ud_type, start, count=""):
         chat_id = self._mess.chat.id
         mes_id = self._mess.id
         from_user = self._from_user
@@ -91,12 +91,13 @@ async def progress_for_pyrogram(self, current, total, ud_type, start):
             # cpu = "{psutil.cpu_percent()}%"
             tmp = (
                 progress
-                + "\n**⌧ Total 🗃:**` 〚{1}〛`\n**⌧ Done ✅ :**` 〚{0}〛`\n**⌧ Speed 📊 :** ` 〚{2}/s〛`\n**⌧ ETA 🔃 :**` 〚{3}〛`".format(
+                + "\n**⌧ Total 🗃:**` 〚{1}〛`\n**⌧ Done ✅ :**` 〚{0}〛`\n**⌧ Speed 📊 :** ` 〚{2}/s〛`\n**⌧ ETA 🔃 :**` 〚{3}〛`\n {4}".format(
                     humanbytes(current),
                     humanbytes(total),
                     humanbytes(speed),
                     # elapsed_time if elapsed_time != '' else "0 s",
                     estimated_total_time if estimated_total_time != "" else "0 s",
+                    count
                 )
             )
             try:

diff --git a/helpers/ffmpeg_helper.py b/helpers/ffmpeg_helper.py
@@ -8,6 +8,8 @@
 from config import Config
 from pyrogram.types import Message
 from __init__ import LOGGER
+from helpers.utils import get_path_size
+
 
 async def MergeVideo(input_file: str, user_id: int, message: Message, format_: str):
     """
@@ -96,7 +98,7 @@ async def MergeSub(filePath: str, subPath: str, user_id):
 
 def MergeSubNew(filePath: str, subPath: str, user_id, file_list):
     """
-    This is for Merging Video + Subtitle(s) Together.
+    This method is for Merging Video + Subtitle(s) Together.
 
     Parameters:
     - `filePath`: Path to Video file.
@@ -132,7 +134,7 @@ def MergeSubNew(filePath: str, subPath: str, user_id, file_list):
     return f"downloads/{str(user_id)}/[@yashoswalyo]_softmuxed_video.mkv"
 
 
-def MergeAudio(videoPath:str,files_list:list,user_id):
+def MergeAudio(videoPath: str, files_list: list, user_id):
     LOGGER.info("Generating Mux Command")
     inputfiles = ""
     maps = ""
@@ -144,18 +146,22 @@ def MergeAudio(videoPath:str,files_list:list,user_id):
     for i in range(len(videoStreamsData)):
         if videoStreamsData[i]["codec_type"] == "audio":
             rmDispositions += f"-disposition:a:{audioTracks} 0 "
-            audioTracks+=1
+            audioTracks += 1
     for i in files_list:
         inputfiles += f"-i '{i}' "
-    for j in range(1,len(files_list)):
+    for j in range(1, len(files_list)):
         maps += f"-map {j}:a "
-    LOGGER.info("Merging files now")
-    subprocess.call(
-        f"ffmpeg -hide_banner {inputfiles}-map 0:v:0 -map 0:a:? {maps}-map 0:s:? -c:v copy -c:a copy -c:s copy {rmDispositions}-disposition:a:{audioTracks} default './downloads/{str(user_id)}/[@yashoswalyo]_export.mkv'",
-        shell=True
+    LOGGER.info(
+        f"Command: ffmpeg -hide_banner {inputfiles}-map 0:v:0 -map 0:a:? {maps}-map 0:s:? -c:v copy -c:a copy -c:s copy {rmDispositions}-disposition:a:{audioTracks} default 'downloads/{str(user_id)}/[@yashoswalyo]_export.mkv'"
+    )
+    process = subprocess.call(
+        f"ffmpeg -hide_banner {inputfiles}-map 0:v:0 -map 0:a:? {maps}-map 0:s:? -c:v copy -c:a copy -c:s copy {rmDispositions}-disposition:a:{audioTracks} default 'downloads/{str(user_id)}/[@yashoswalyo]_export.mkv'",
+        shell=True,
     )
+    LOGGER.info(process)
     return f"downloads/{str(user_id)}/[@yashoswalyo]_export.mkv"
 
+
 async def cult_small_video(video_file, output_directory, start_time, end_time, format_):
     # https://stackoverflow.com/a/13891070/4723940
     out_put_file_name = (
@@ -248,3 +254,113 @@ async def take_screen_shot(video_file, output_directory, ttl):
         return out_put_file_name
     else:
         return None
+
+
+async def extractAudios(path_to_file, user_id):
+    '''
+    Thanks to Jitu Bhai and group.
+    '''
+    dir_name = os.path.dirname(os.path.dirname(path_to_file))
+    if not os.path.exists(path_to_file):
+        return None
+    if not os.path.exists(dir_name + "/extract"):
+        os.makedirs(dir_name + "/extract")
+    videoStreamsData = ffmpeg.probe(path_to_file)
+    # with open("data.json",'w') as f:
+    #     f.write(json.dumps(videoStreamsData))
+    extract_dir = dir_name + "/extract"
+    audios = []
+    for stream in videoStreamsData.get("streams"):
+        try:
+            if stream["codec_type"] == "audio":
+                audios.append(stream)
+        except Exception as e:
+            LOGGER.warning(e)
+    for audio in audios:
+        try:
+            index = audio["index"]
+            try:
+                output_file: str = (
+                    "("
+                    + audio["tags"]["language"]
+                    + ") "
+                    + audio["tags"]["title"]
+                    + "."
+                    + audio["codec_type"]
+                    + ".mka"
+                )
+                output_file = output_file.replace(" ", ".")
+            except:
+                output_file = str(audio["index"]) + "." + audio["codec_type"] + ".mka"
+            exec_command = f"ffmpeg -hide_banner -i '{path_to_file}' -map 0:{index} -c copy '{extract_dir}/{output_file}'"
+            LOGGER.info(exec_command)
+            subprocess.call(exec_command, shell=True)
+        except:
+            1
+    if get_path_size(extract_dir) > 0:
+        return extract_dir
+    else:
+        LOGGER.warning(f"{extract_dir} is empty")
+        return None
+
+
+async def extractSubtitles(path_to_file, user_id):
+    '''
+    Thanks to Jitu Bhai and group.
+    '''
+    dir_name = os.path.dirname(os.path.dirname(path_to_file))
+    if not os.path.exists(path_to_file):
+        return None
+    if not os.path.exists(dir_name + "/extract"):
+        os.makedirs(dir_name + "/extract")
+    videoStreamsData = ffmpeg.probe(path_to_file)
+    # with open("data.json",'w') as f:
+    #     f.write(json.dumps(videoStreamsData))
+    extract_dir = dir_name + "/extract"
+    subtitles = []
+    for stream in videoStreamsData.get("streams"):
+        try:
+            if stream["codec_type"] == "subtitle":
+                subtitles.append(stream)
+        except Exception as e:
+            LOGGER.warning(e)
+    for subtitle in subtitles:
+        try:
+            index = subtitle["index"]
+            try:
+                output_file: str = (
+                    "("
+                    + subtitle["tags"]["language"]
+                    + ") "
+                    + subtitle["tags"]["title"]
+                    + "."
+                    + subtitle["codec_type"]
+                    + ".mka"
+                )
+                output_file = output_file.replace(" ", ".")
+            except:
+                try:
+                    output_file = (
+                        str(subtitle["index"])
+                        + "."
+                        + subtitle["tags"]["language"]
+                        + "."
+                        + subtitle["codec_type"]
+                        + ".mka"
+                    )
+                except:
+                    output_file = (
+                        str(subtitle["index"]) + "." + subtitle["codec_type"] + ".mka"
+                    )
+
+            exec_command = f"ffmpeg -hide_banner -i '{path_to_file}' -map 0:{index} -c copy '{extract_dir}/{output_file}'"
+            LOGGER.info(exec_command)
+            subprocess.call(exec_command, shell=True)
+        except:
+            1
+    if get_path_size(extract_dir) > 0:
+        return extract_dir
+    else:
+        LOGGER.warning(f"{extract_dir} is empty")
+        return None
+    1
diff --git a/helpers/uploader.py b/helpers/uploader.py
@@ -107,3 +107,37 @@ async def uploadVideo(
                     chat_id=int(LOGCHANNEL),
                     caption=f"`{media.file_name}`\n\nMerged for: <a href='tg://user?id={cb.from_user.id}'>{cb.from_user.first_name}</a>",
                 )
+
+
+async def uploadFiles(
+    c: Client,
+    cb: CallbackQuery,
+    up_path,
+    n,
+    all
+):
+    try:
+        sent_ = None
+        prog = Progress(cb.from_user.id, c, cb.message)
+        c_time = time.time()
+        sent_: Message = await c.send_document(
+            chat_id=cb.message.chat.id,
+            document=up_path,
+            caption=f"`{up_path.rsplit('/',1)[-1]}`",
+            progress=prog.progress_for_pyrogram,
+            progress_args=(
+                f"Uploading: `{up_path.rsplit('/',1)[-1]}`",
+                c_time,
+                f"\n**Uploading: {n}/{all}**"
+            ),
+        )
+        if sent_ is not None:
+            if Config.LOGCHANNEL is not None:
+                media = sent_.video or sent_.document
+                await sent_.copy(
+                    chat_id=int(LOGCHANNEL),
+                    caption=f"`{media.file_name}`\n\nExtracted by: <a href='tg://user?id={cb.from_user.id}'>{cb.from_user.first_name}</a>",
+                )
+    except:
+        1    
+    1
diff --git a/plugins/cb_handler.py b/plugins/cb_handler.py
@@ -10,6 +10,7 @@
     gDict,
     queueDB,
     showQueue,
+    mergeApp
 )
 from helpers import database
 from helpers.utils import UserSettings
@@ -24,6 +25,7 @@
 from plugins.mergeVideo import mergeNow
 from plugins.mergeVideoAudio import mergeAudio
 from plugins.mergeVideoSub import mergeSub
+from plugins.streams_extractor import streamsExtractor
 from plugins.usettings import userSettings
 
 
@@ -406,3 +408,17 @@ async def callback_handler(c: Client, cb: CallbackQuery):
             cb.message, uid, cb.from_user.first_name, cb.from_user.last_name, user
         )
         return
+
+    elif cb.data.startswith('extract'):
+        edata = cb.data.split('_')[1]
+        media_mid = int(cb.data.split('_')[2])
+        try:
+            if edata == 'audio':
+                LOGGER.info('audio')
+                await streamsExtractor(c,cb,media_mid,exAudios=True)
+            elif edata == 'subtitle':
+                await streamsExtractor(c,cb,media_mid,exSubs=True)
+            elif edata == 'all':
+                await streamsExtractor(c,cb,media_mid,exAudios=True,exSubs=True)
+        except Exception as e:
+            LOGGER.error(e)