add direct download support

README.md

@@ -216,6 +216,7 @@ ytdl - Download video in group
 vip - Join VIP
 terms - View Terms of Service
 settings - Set your preference
+direct - Download file directly
 ```
 
 # Test data

ytdlbot/tasks.py

@@ -10,19 +10,23 @@ __author__ = "Benny <benny.think@gmail.com>"
 import logging
 import os
 import pathlib
+import re
 import tempfile
 import threading
 import time
+from urllib.parse import quote_plus
 
+import requests
 from celery import Celery
 from pyrogram import idle
 from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup
 
 from client_init import create_app
-from config import BROKER, ENABLE_CELERY, OWNER, WORKERS
+from config import BROKER, ENABLE_CELERY, ENABLE_VIP, OWNER, WORKERS
 from constant import BotText
 from db import Redis
 from downloader import convert_flac, sizeof_fmt, upload_hook, ytdl_download
+from limit import VIP
 from utils import (apply_log_formatter, customize_logger, get_metadata,
                    get_user_settings)
 
@@ -47,11 +51,11 @@ def get_messages(chat_id, message_id):
 
 
 @app.task()
-def download_task(chat_id, message_id, url):
-    logging.info("celery tasks started for %s", url)
+def ytdl_download_task(chat_id, message_id, url):
+    logging.info("YouTube celery tasks started for %s", url)
     bot_msg = get_messages(chat_id, message_id)
-    normal_download(bot_msg, celery_client, url)
-    logging.info("celery tasks ended.")
+    ytdl_normal_download(bot_msg, celery_client, url)
+    logging.info("YouTube celery tasks ended.")
 
 
 @app.task()
@@ -62,11 +66,28 @@ def audio_task(chat_id, message_id):
     logging.info("Audio celery tasks ended.")
 
 
-def download_entrance(bot_msg, client, url):
+@app.task()
+def direct_download_task(chat_id, message_id, url):
+    logging.info("Direct download celery tasks started for %s", url)
+    bot_msg = get_messages(chat_id, message_id)
+    direct_normal_download(bot_msg, celery_client, url)
+    logging.info("Direct download celery tasks ended.")
+
+
+def ytdl_download_entrance(bot_msg, client, url):
     if ENABLE_CELERY:
-        download_task.delay(bot_msg.chat.id, bot_msg.message_id, url)
+        ytdl_download_task.delay(bot_msg.chat.id, bot_msg.message_id, url)
     else:
-        normal_download(bot_msg, client, url)
+        ytdl_normal_download(bot_msg, client, url)
+
+
+def direct_download_entrance(bot_msg, client, url):
+    if ENABLE_CELERY:
+        # TODO disable it for now
+        direct_normal_download(bot_msg, client, url)
+        # direct_download_task.delay(bot_msg.chat.id, bot_msg.message_id, url)
+    else:
+        direct_normal_download(bot_msg, client, url)
 
 
 def audio_entrance(bot_msg):
@@ -76,6 +97,53 @@ def audio_entrance(bot_msg):
         normal_audio(bot_msg)
 
 
+def direct_normal_download(bot_msg, client, url):
+    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/80.0.3987.149 Safari/537.36"}
+    vip = VIP()
+
+    if ENABLE_VIP:
+        remain, _, _ = vip.check_remaining_quota(chat_id)
+        try:
+            head_req = requests.head(url, headers=headers)
+            length = int(head_req.headers.get("content-length"))
+        except:
+            length = 0
+        if remain < length:
+            bot_msg.reply_text(f"Sorry, you have reached your quota.\n")
+            return
+
+    req = ""
+    try:
+        req = requests.get(url, headers=headers)
+        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)
+
+    with tempfile.TemporaryDirectory() as f:
+        filepath = f"{f}/{filename}"
+        print(filepath)
+        with open(filepath, "wb") as fp:
+            fp.write(req.content)
+        logging.info("Downloaded file %s", filename)
+        st_size = os.stat(filepath).st_size
+        if ENABLE_VIP:
+            vip.use_quota(chat_id, st_size)
+        client.send_chat_action(chat_id, "upload_document")
+        client.send_document(bot_msg.chat.id, filepath,
+                             caption=f"filesize: {sizeof_fmt(st_size)}",
+                             progress=upload_hook, progress_args=(bot_msg,),
+                             )
+        bot_msg.edit_text(f"Download success!✅")
+
+
 def normal_audio(bot_msg):
     chat_id = bot_msg.chat.id
     mp4_name = bot_msg.video.file_name  # 'youtube-dl_test_video_a.mp4'
@@ -103,7 +171,7 @@ def get_worker_status(username):
     return f"Downloaded by {me.mention()}"
 
 
-def normal_download(bot_msg, client, url):
+def ytdl_normal_download(bot_msg, client, url):
     chat_id = bot_msg.chat.id
     temp_dir = tempfile.TemporaryDirectory()
 

ytdlbot/ytdl_bot.py

@@ -24,7 +24,8 @@ from config import (AUTHORIZED_USER, ENABLE_CELERY, ENABLE_VIP, OWNER,
 from constant import BotText
 from db import InfluxDB, MySQL, Redis
 from limit import verify_payment
-from tasks import audio_entrance, download_entrance
+from tasks import (audio_entrance, direct_download_entrance,
+                   ytdl_download_entrance)
 from utils import (auto_restart, customize_logger, get_revision,
                    get_user_settings, set_user_settings)
 
@@ -119,6 +120,22 @@ def terms_handler(client: "Client", message: "types.Message"):
     client.send_message(chat_id, bot_text.terms)
 
 
+@app.on_message(filters.command(["direct"]))
+def direct_handler(client: "Client", message: "types.Message"):
+    chat_id = message.from_user.id
+    client.send_chat_action(chat_id, "typing")
+    url = re.sub(r'/direct\s*', '', message.text)
+    logging.info("direct start %s", url)
+    if not re.findall(r"^https?://", url.lower()):
+        Redis().update_metrics("bad_request")
+        message.reply_text("Send me a DIRECT LINK.", quote=True)
+        return
+
+    bot_msg = message.reply_text("Request received.", quote=True)
+    Redis().update_metrics("direct_request")
+    direct_download_entrance(bot_msg, client, url)
+
+
 @app.on_message(filters.command(["settings"]))
 def settings_handler(client: "Client", message: "types.Message"):
     chat_id = message.chat.id
@@ -176,7 +193,7 @@ def download_handler(client: "Client", message: "types.Message"):
     text = bot_text.get_receive_link_text()
     bot_msg: typing.Union["types.Message", "typing.Any"] = message.reply_text(text, quote=True)
     client.send_chat_action(chat_id, 'upload_video')
-    download_entrance(bot_msg, client, url)
+    ytdl_download_entrance(bot_msg, client, url)
 
 
 @app.on_callback_query(filters.regex(r"document|video"))