leech option for direct download using aria2 (#408)

* leech option for direct download using aria2

* Improved and fixed some bugs of leech function

README.md

@@ -45,6 +45,7 @@ Just send a link directly to the bot.
 12. instagram posts(only available for my bot)
 13. 4 GiB file size support with Telegram Premium
 14. History and inline mode support
+15. Supports multiple download engines (yt-dlp, aria2, requests).
 
 > [!NOTE]
 > **For users of [my official bot](https://t.me/benny_ytdlbot)**\
@@ -297,7 +298,8 @@ about - What's this bot?
 help - Help
 spdl - Use to download specific link downloader links
 ytdl - Download video in group
-direct - Download file directly
+leech - Download file using aria2
+direct - Download file using requests
 settings - Set your preference
 buy - Buy token
 sub - Subscribe to YouTube Channel

ytdlbot/tasks.py

@@ -59,6 +59,7 @@ from utils import (
     get_revision,
     sizeof_fmt,
     shorten_url,
+    extract_filename,
 )
 
 customize_logger(["pyrogram.client", "pyrogram.session.session", "pyrogram.connection.connection"])
@@ -145,6 +146,14 @@ def direct_download_task(chat_id: int, message_id: int, url: str):
     logging.info("Direct download celery tasks ended.")
 
 
+@app.task()
+def leech_download_task(chat_id: int, message_id: int, url: str):
+    logging.info("Leech download celery tasks started for %s", url)
+    bot_msg = retrieve_message(chat_id, message_id)
+    leech_normal_download(bot, bot_msg, url)
+    logging.info("Leech download celery tasks ended.")
+
+
 def get_unique_clink(original_url: str, user_id: int):
     payment = Payment()
     settings = payment.get_user_settings(user_id)
@@ -213,6 +222,14 @@ def direct_download_entrance(client: Client, bot_msg: typing.Union[types.Message
         direct_normal_download(client, bot_msg, url)
 
 
+def leech_download_entrance(client: Client, bot_msg: typing.Union[types.Message, typing.Coroutine], url: str):
+    if ENABLE_CELERY:
+        leech_normal_download(client, bot_msg, url)
+        # leech_normal_download.delay(bot_msg.chat.id, bot_msg.id, url)
+    else:
+        leech_normal_download(client, bot_msg, url)
+
+
 def spdl_download_entrance(client: Client, bot_msg: types.Message, url: str, mode=None):
     payment = Payment()
     redis = Redis()
@@ -255,7 +272,7 @@ def audio_entrance(client: Client, bot_msg: types.Message):
 def direct_normal_download(client: Client, bot_msg: typing.Union[types.Message, typing.Coroutine], url: str):
     chat_id = bot_msg.chat.id
     headers = {
-        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.3987.149 Safari/537.36"
+        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36"
     }
     length = 0
 
@@ -263,15 +280,11 @@ def direct_normal_download(client: Client, bot_msg: typing.Union[types.Message,
     try:
         req = requests.get(url, headers=headers, stream=True)
         length = int(req.headers.get("content-length"))
-        filename = re.findall("filename=(.+)", req.headers.get("content-disposition"))[0]
-    except TypeError:
-        filename = getattr(req, "url", "").rsplit("/")[-1]
     except Exception as e:
         bot_msg.edit_text(f"Download failed!❌\n\n```{e}```", disable_web_page_preview=True)
         return
 
-    if not filename:
-        filename = quote_plus(url)
+    filename = extract_filename(req)
 
     with tempfile.TemporaryDirectory(prefix="ytdl-", dir=TMPFILE_PATH) as f:
         filepath = f"{f}/{filename}"
@@ -297,6 +310,66 @@ def direct_normal_download(client: Client, bot_msg: typing.Union[types.Message,
         bot_msg.edit_text("Download success!✅")
 
 
+def leech_normal_download(client: Client, bot_msg: typing.Union[types.Message, typing.Coroutine], url: str):
+    chat_id = bot_msg.chat.id
+    temp_dir = tempfile.TemporaryDirectory(prefix="leech_dl-", dir=TMPFILE_PATH)
+    tempdir = temp_dir.name
+    UA = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36"
+    response = None
+    video_paths = None
+    # Download process using aria2c
+    try:
+        bot_msg.edit_text(f"Download Starting...", disable_web_page_preview=True)
+        # Command to download the link using aria2c
+        command = [
+            "aria2c",
+            "-U",
+            UA,
+            "--max-tries=5",
+            "--console-log-level=warn",
+            "-d",
+            tempdir,
+            url,
+        ]
+        # Run the command using subprocess.Popen
+        process = subprocess.Popen(command, bufsize=0, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        line = ""
+        max_iterations = 100  # Set a reasonable maximum number of iterations
+        iteration = 0
+
+        while process.poll() is None and iteration < max_iterations:
+            line = process.stdout.readline().decode("utf-8")
+            if line.startswith("[#"):
+                line = line.strip()
+                bot_msg.edit_text(f"Downloading... \n\n`{line}`", disable_web_page_preview=True)
+                break
+            iteration += 1
+        
+        if iteration >= max_iterations:
+            bot_msg.edit_text("Something went wrong. Please try again.", disable_web_page_preview=True)
+    except Exception as e:
+        bot_msg.edit_text(f"Download failed!❌\n\n`{e}`", disable_web_page_preview=True)
+        return
+    # Get filename and extension correctly after download
+    filepath = list(pathlib.Path(tempdir).glob("*"))
+    file_path_obj = filepath[0]
+    path_obj = pathlib.Path(file_path_obj)
+    filename = path_obj.name
+    logging.info("Downloaded file %s", filename)
+    bot_msg.edit_text(f"Download Complete", disable_web_page_preview=True)
+    ext = filetype.guess_extension(file_path_obj)
+    # Rename file if it doesn't have extension
+    if ext is not None:
+        if not filename.endswith(ext):
+            new_filename = f"{tempdir}/{filename}.{ext}"
+            os.rename(file_path_obj, new_filename)
+    # Get file path of the downloaded file to upload
+    video_paths = list(pathlib.Path(tempdir).glob("*"))
+    client.send_chat_action(chat_id, enums.ChatAction.UPLOAD_DOCUMENT)
+    upload_processor(client, bot_msg, url, video_paths)
+    bot_msg.edit_text("Download success!✅")
+
+
 def normal_audio(client: Client, bot_msg: typing.Union[types.Message, typing.Coroutine]):
     chat_id = bot_msg.chat.id
     # fn = getattr(bot_msg.video, "file_name", None) or getattr(bot_msg.document, "file_name", None)

ytdlbot/utils.py

@@ -22,6 +22,7 @@ import re
 import coloredlogs
 import ffmpeg
 import psutil
+from urllib.parse import quote_plus
 from http.cookiejar import MozillaCookieJar
 
 from config import TMPFILE_PATH
@@ -249,5 +250,21 @@ def shorten_url(url, CAPTION_URL_LENGTH_LIMIT):
   return shortened_url
 
 
+def extract_filename(response):
+    try:
+        content_disposition = response.headers.get("content-disposition")
+        if content_disposition:
+            filename = re.findall("filename=(.+)", content_disposition)[0]
+            return filename
+    except (TypeError, IndexError):
+        pass  # Handle potential exceptions during extraction
+
+    # Fallback if Content-Disposition header is missing
+    filename = response.url.rsplit("/")[-1]
+    if not filename:
+        filename = quote_plus(response.url)
+    return filename
+
+
 if __name__ == "__main__":
     auto_restart()

ytdlbot/ytdl_bot.py

@@ -46,6 +46,7 @@ from config import (
     REQUIRED_MEMBERSHIP,
     TOKEN_PRICE,
     TRX_SIGNAL,
+    ENABLE_ARIA2,
 )
 from constant import BotText
 from database import InfluxDB, MySQL, Redis
@@ -54,6 +55,7 @@ from tasks import app as celery_app
 from tasks import (
     audio_entrance,
     direct_download_entrance,
+    leech_download_entrance,
     hot_patch,
     purge_tasks,
     ytdl_download_entrance,
@@ -267,6 +269,23 @@ def clear_history(client: Client, message: types.Message):
     message.reply_text("History cleared.", quote=True)
 
 
+@app.on_message(filters.command(["spdl"]))
+def spdl_handler(client: Client, message: types.Message):
+    redis = Redis()
+    chat_id = message.from_user.id
+    client.send_chat_action(chat_id, enums.ChatAction.TYPING)
+    url = re.sub(r"/spdl\s*", "", message.text)
+    logging.info("spdl start %s", url)
+    if not re.findall(r"^https?://", url.lower()):
+        redis.update_metrics("bad_request")
+        message.reply_text("Something wrong 🤔.\nCheck your URL and send me again.", quote=True)
+        return
+
+    bot_msg = message.reply_text("Request received.", quote=True)
+    redis.update_metrics("spdl_request")
+    spdl_download_entrance(client, bot_msg, url)
+
+
 @app.on_message(filters.command(["direct"]))
 def direct_handler(client: Client, message: types.Message):
     redis = Redis()
@@ -283,21 +302,25 @@ def direct_handler(client: Client, message: types.Message):
     redis.update_metrics("direct_request")
     direct_download_entrance(client, bot_msg, url)
 
-@app.on_message(filters.command(["spdl"]))
-def spdl_handler(client: Client, message: types.Message):
+
+@app.on_message(filters.command(["leech"]))
+def leech_handler(client: Client, message: types.Message):
+    if not ENABLE_ARIA2:
+        message.reply_text("Aria2 Not Enabled.", quote=True)
+        return
     redis = Redis()
     chat_id = message.from_user.id
     client.send_chat_action(chat_id, enums.ChatAction.TYPING)
-    url = re.sub(r"/spdl\s*", "", message.text)
-    logging.info("spdl start %s", url)
+    url = re.sub(r"/leech\s*", "", message.text)
+    logging.info("leech using aria2 start %s", url)
     if not re.findall(r"^https?://", url.lower()):
         redis.update_metrics("bad_request")
-        message.reply_text("Something wrong 🤔.\nCheck your URL and send me again.", quote=True)
+        message.reply_text("Send me a correct LINK.", quote=True)
         return
 
     bot_msg = message.reply_text("Request received.", quote=True)
-    redis.update_metrics("direct_request")
-    spdl_download_entrance(client, bot_msg, url)
+    redis.update_metrics("leech_request")
+    leech_download_entrance(client, bot_msg, url)
 
 
 @app.on_message(filters.command(["settings"]))