Dec 13

Author: 5hojib

bot/__init__.py

@@ -96,6 +96,7 @@ def format(self, record: LogRecord) -> str:
 rss_dict = {}
 auth_chats = {}
 excluded_extensions = ["aria2", "!qB"]
+included_extensions = []
 drives_names = []
 drives_ids = []
 index_urls = []

bot/core/config_manager.py

@@ -18,6 +18,7 @@ class Config:
     DATABASE_URL: str = ""
     DEFAULT_UPLOAD: str = "gd"
     EXCLUDED_EXTENSIONS: str = ""
+    INCLUDED_EXTENSIONS: str = ""
     FFMPEG_CMDS: ClassVar[dict[str, list[str]]] = {}
     FILELION_API: str = ""
     GDRIVE_ID: str = ""

bot/core/startup.py

@@ -14,6 +14,7 @@
     drives_ids,
     drives_names,
     excluded_extensions,
+    included_extensions,
     index_urls,
     nzb_options,
     qbit_options,
@@ -269,6 +270,11 @@ async def update_variables():
             x = x.lstrip(".")
             excluded_extensions.append(x.strip().lower())
 
+    if Config.INCLUDED_EXTENSIONS:
+        fx = Config.INCLUDED_EXTENSIONS.split()
+        for x in fx:
+            x = x.lstrip(".")
+            included_extensions.append(x.strip().lower())
     if Config.GDRIVE_ID:
         drives_names.append("Main")
         drives_ids.append(Config.GDRIVE_ID)

bot/helper/common.py

@@ -20,6 +20,7 @@
     cores,
     cpu_eater_lock,
     excluded_extensions,
+    included_extensions,
     intervals,
     multi_tags,
     task_dict,
@@ -137,6 +138,7 @@ def __init__(self):
         self.subproc = None
         self.thumb = None
         self.excluded_extensions = []
+        self.included_extensions = []
         self.files_to_proceed = []
         self.is_super_chat = self.message.chat.type.name in [
             "SUPERGROUP",
@@ -227,6 +229,9 @@ async def before_start(self):
             if "EXCLUDED_EXTENSIONS" not in self.user_dict
             else ["aria2", "!qB"]
         )
+        self.included_extensions = self.user_dict.get("INCLUDED_EXTENSIONS") or (
+            included_extensions if "INCLUDED_EXTENSIONS" not in self.user_dict else []
+        )
         if not self.rc_flags:
             if self.user_dict.get("RCLONE_FLAGS"):
                 self.rc_flags = self.user_dict["RCLONE_FLAGS"]

bot/helper/ext_utils/bot_utils.py

@@ -191,7 +191,10 @@ def arg_parser(items, arg_base):
             else:
                 sub_list = []
                 for j in range(i + 1, total):
-                    if items[j] in arg_base and part != "-c" != items[j]:
+                    if items[j] in arg_base:
+                        if part == "-c" and items[j] == "-c":
+                            sub_list.append(items[j])
+                            continue
                         if part in bool_arg_set and not sub_list:
                             arg_base[part] = True
                             break

bot/helper/ext_utils/files_utils.py

@@ -237,10 +237,20 @@ def get_mime_type(file_path: str) -> str:
 
 async def remove_excluded_files(fpath, ee):
     for root, _, files in await sync_to_async(walk, fpath):
+        if root.strip().endswith("/yt-dlp-thumb"):
+            continue
         for f in files:
             if f.strip().lower().endswith(tuple(ee)):
                 await remove(ospath.join(root, f))
 
+async def remove_non_included_files(fpath, ie):
+    for root, _, files in await sync_to_async(walk, fpath):
+        if root.strip().endswith("/yt-dlp-thumb"):
+            continue
+        for f in files:
+            if f.strip().lower().endswith(tuple(ie)):
+                continue
+            await remove(ospath.join(root, f))
 
 async def join_files(opath):
     files = await listdir(opath)
@@ -329,7 +339,7 @@ async def _sevenz_progress(self):
                 break
             line = line.decode().strip()
             if match := re_search(pattern, line):
-                self._listener.subsize = int(match[1] or match[2])
+                self._listener.subsize = int(match[1] or match[2] or match[3])
         s = b""
         while not (
             self._listener.is_cancelled

bot/helper/ext_utils/help_messages.py

@@ -392,6 +392,7 @@
     "INDEX_URL": "Send Index URL. Timeout: 60 sec",
     "UPLOAD_PATHS": "Send Dict of keys that have path values. Example: {'path 1': 'remote:rclonefolder', 'path 2': 'gdrive1 id', 'path 3': 'tg chat id', 'path 4': 'mrcc:remote:', 'path 5': b:@username} . Timeout: 60 sec",
     "EXCLUDED_EXTENSIONS": "Send excluded extensions separated by space without dot at beginning. Timeout: 60 sec",
+    "INCLUDED_EXTENSIONS": "Send included extensions separated by space without dot at beginning. Timeout: 60 sec",
     "NAME_SUBSTITUTE": r"""Word Subtitions. You can add pattern instead of normal text. Timeout: 60 sec
 NOTE: You must add \ before any character, those are the characters: \^$.|?*+()[]{}-
 Example: script/code/s | mirror/leech | tea/ /s | clone | cpu/ | \[mltb\]/mltb | \\text\\/text/s

bot/helper/ext_utils/media_utils.py

@@ -1,5 +1,5 @@
 import contextlib
-from asyncio import create_subprocess_exec, gather, wait_for
+from asyncio import create_subprocess_exec, gather, sleep, wait_for
 from asyncio.subprocess import PIPE
 from os import path as ospath
 from re import escape
@@ -11,7 +11,7 @@
 from aioshutil import rmtree
 from PIL import Image
 
-from bot import DOWNLOAD_DIR, LOGGER, cpu_no
+from bot import DOWNLOAD_DIR, LOGGER, cpu_no, threads, cores
 
 from .bot_utils import cmd_exec, sync_to_async
 from .files_utils import get_mime_type, is_archive, is_archive_split
@@ -131,8 +131,8 @@ async def take_ss(video_file, ss_nb) -> bool:
             output = f"{dirpath}/SS.{name}_{i:02}.png"
             cmd = [
                 "taskset",
-                "-c",
-                f"{cores}",
+                    "-c",
+                    f"{cores}",
                 "xtra",
                 "-hide_banner",
                 "-loglevel",
@@ -146,7 +146,7 @@ async def take_ss(video_file, ss_nb) -> bool:
                 "-frames:v",
                 "1",
                 "-threads",
-                f"{max(1, cpu_no // 2)}",
+                f"{threads}",
                 output,
             ]
             cap_time += interval
@@ -176,8 +176,8 @@ async def get_audio_thumbnail(audio_file):
     output = ospath.join(output_dir, f"{time()}.jpg")
     cmd = [
         "taskset",
-        "-c",
-        f"{cores}",
+                    "-c",
+                    f"{cores}",
         "xtra",
         "-hide_banner",
         "-loglevel",
@@ -188,7 +188,7 @@ async def get_audio_thumbnail(audio_file):
         "-vcodec",
         "copy",
         "-threads",
-        f"{max(1, cpu_no // 2)}",
+        f"{threads}",
         output,
     ]
     try:
@@ -218,8 +218,8 @@ async def get_video_thumbnail(video_file, duration):
     duration = duration // 2
     cmd = [
         "taskset",
-        "-c",
-        f"{cores}",
+                    "-c",
+                    f"{cores}",
         "xtra",
         "-hide_banner",
         "-loglevel",
@@ -264,8 +264,8 @@ async def get_multiple_frames_thumbnail(video_file, layout, keep_screenshots):
     output = ospath.join(output_dir, f"{time()}.jpg")
     cmd = [
         "taskset",
-        "-c",
-        f"{cores}",
+                    "-c",
+                    f"{cores}",
         "xtra",
         "-hide_banner",
         "-loglevel",
@@ -283,7 +283,7 @@ async def get_multiple_frames_thumbnail(video_file, layout, keep_screenshots):
         "-f",
         "mjpeg",
         "-threads",
-        f"{max(1, cpu_no // 2)}",
+        f"{threads}",
         output,
     ]
     try:
@@ -388,6 +388,7 @@ async def _ffmpeg_progress(self):
                         except Exception:
                             self._progress_raw = 0
                             self._eta_raw = 0
+        
 
     async def ffmpeg_cmds(self, ffmpeg, f_path):
         self.clear()
@@ -483,8 +484,8 @@ async def convert_video(self, video_file, ext, retry=False):
         if retry:
             cmd = [
                 "taskset",
-                "-c",
-                f"{cores}",
+                    "-c",
+                    f"{cores}",
                 "xtra",
                 "-hide_banner",
                 "-loglevel",
@@ -500,7 +501,7 @@ async def convert_video(self, video_file, ext, retry=False):
                 "-c:a",
                 "aac",
                 "-threads",
-                f"{max(1, cpu_no // 2)}",
+                f"{threads}",
                 output,
             ]
             if ext == "mp4":
@@ -512,8 +513,8 @@ async def convert_video(self, video_file, ext, retry=False):
         else:
             cmd = [
                 "taskset",
-                "-c",
-                f"{cores}",
+                    "-c",
+                    f"{cores}",
                 "xtra",
                 "-hide_banner",
                 "-loglevel",
@@ -527,7 +528,7 @@ async def convert_video(self, video_file, ext, retry=False):
                 "-c",
                 "copy",
                 "-threads",
-                f"{max(1, cpu_no // 2)}",
+                f"{threads}",
                 output,
             ]
         if self._listener.is_cancelled:
@@ -567,8 +568,8 @@ async def convert_audio(self, audio_file, ext):
         output = f"{base_name}.{ext}"
         cmd = [
             "taskset",
-            "-c",
-            f"{cores}",
+                    "-c",
+                    f"{cores}",
             "xtra",
             "-hide_banner",
             "-loglevel",
@@ -578,7 +579,7 @@ async def convert_audio(self, audio_file, ext):
             "-i",
             audio_file,
             "-threads",
-            f"{max(1, cpu_no // 2)}",
+            f"{threads}",
             output,
         ]
         if self._listener.is_cancelled:
@@ -641,8 +642,8 @@ async def sample_video(self, video_file, sample_duration, part_duration):
 
         cmd = [
             "taskset",
-            "-c",
-            f"{cores}",
+                    "-c",
+                    f"{cores}",
             "xtra",
             "-hide_banner",
             "-loglevel",
@@ -662,7 +663,7 @@ async def sample_video(self, video_file, sample_duration, part_duration):
             "-c:a",
             "aac",
             "-threads",
-            f"{max(1, cpu_no // 2)}",
+            f"{threads}",
             output_file,
         ]
 
@@ -706,8 +707,8 @@ async def split(self, f_path, file_, parts, split_size):
             out_path = f_path.replace(file_, f"{base_name}.part{i:03}{extension}")
             cmd = [
                 "taskset",
-                "-c",
-                f"{cores}",
+                    "-c",
+                    f"{cores}",
                 "xtra",
                 "-hide_banner",
                 "-loglevel",
@@ -731,7 +732,7 @@ async def split(self, f_path, file_, parts, split_size):
                 "-c",
                 "copy",
                 "-threads",
-                f"{max(1, cpu_no // 2)}",
+                f"{threads}",
                 out_path,
             ]
             if not multi_streams:

bot/helper/listeners/task_listener.py

@@ -32,6 +32,7 @@
     get_path_size,
     join_files,
     remove_excluded_files,
+    remove_non_included_files
 )
 from bot.helper.ext_utils.links_utils import is_gdrive_id
 from bot.helper.ext_utils.status_utils import get_readable_file_size
@@ -196,10 +197,14 @@ async def on_download_complete(self):
         else:
             up_dir = self.dir
             up_path = dl_path
-        await remove_excluded_files(
-            self.up_dir or self.dir,
-            self.excluded_extensions,
-        )
+        if not self.included_extensions:
+            await remove_excluded_files(
+                self.up_dir or self.dir, self.excluded_extensions
+            )
+        else:
+            await remove_non_included_files(
+                self.up_dir or self.dir, self.included_extensions
+            )
         if not Config.QUEUE_ALL:
             async with queue_dict_lock:
                 if self.mid in non_queued_dl:
@@ -217,7 +222,10 @@ async def on_download_complete(self):
             self.name = up_path.replace(f"{up_dir}/", "").split("/", 1)[0]
             self.size = await get_path_size(up_dir)
             self.clear()
-            await remove_excluded_files(up_dir, self.excluded_extensions)
+            if not self.included_extensions:
+                await remove_excluded_files(up_dir, self.excluded_extensions)
+            else:
+                await remove_non_included_files(up_dir, self.included_extensions)
 
         if self.watermark:
             up_path = await self.proceed_watermark(