premium download consent

ytdlbot/config.py

@@ -56,7 +56,7 @@
 M3U8_SUPPORT = os.getenv("M3U8_SUPPORT", False)
 ENABLE_ARIA2 = os.getenv("ENABLE_ARIA2", False)
 
-RATE_LIMIT = os.getenv("RATE_LIMIT", 60)
+RATE_LIMIT = os.getenv("RATE_LIMIT", 120)
 RCLONE_PATH = os.getenv("RCLONE")
 # This will set the value for the tmpfile path(download path) if it is set.
 # If TMPFILE is not set, it will return None and use system’s default temporary file path.

ytdlbot/constant.py

@@ -83,6 +83,11 @@ class BotText:
 """
     custom_text = os.getenv("CUSTOM_TEXT", "")
 
+    premium_warning = """
+    Your file is too big, do you want me to try to send it as premium user? 
+    Owner will know who you are and what you are downloading. You may be banned if you abuse this feature.
+    """
+
     @staticmethod
     def get_receive_link_text() -> str:
         reserved = get_func_queue("reserved")

ytdlbot/tasks.py

@@ -8,7 +8,6 @@
 __author__ = "Benny <[email protected]>"
 
 import asyncio
-import json
 import logging
 import os
 import pathlib
@@ -40,7 +39,6 @@
     ENABLE_CELERY,
     ENABLE_VIP,
     OWNER,
-    PREMIUM_USER,
     RATE_LIMIT,
     RCLONE_PATH,
     TMPFILE_PATH,
@@ -53,6 +51,7 @@
 from limit import Payment
 from utils import (
     apply_log_formatter,
+    auto_restart,
     customize_logger,
     get_metadata,
     get_revision,
@@ -79,6 +78,18 @@ def retrieve_message(chat_id: int, message_id: int) -> types.Message | Any:
         return bot.get_messages(chat_id, message_id)
 
 
+def premium_button():
+    markup = types.InlineKeyboardMarkup(
+        [
+            [
+                types.InlineKeyboardButton("Yes", callback_data="premium-yes"),
+                types.InlineKeyboardButton("No", callback_data="premium-no"),
+            ]
+        ]
+    )
+    return markup
+
+
 @app.task(rate_limit=f"{RATE_LIMIT}/m")
 def ytdl_download_task(chat_id: int, message_id: int, url: str):
     logging.info("YouTube celery tasks started for %s", url)
@@ -88,9 +99,7 @@ def ytdl_download_task(chat_id: int, message_id: int, url: str):
     except FileTooBig as e:
         # if you can go there, that means you have premium users set up
         logging.warning("Seeking for help from premium user...")
-        bot_msg.edit_text(f"{e}\n\nPlease wait, I will try to send it as premium user")
-        data = {"url": url, "user_id": bot_msg.chat.id}
-        bot.send_message(PREMIUM_USER, json.dumps(data), disable_notification=True, disable_web_page_preview=True)
+        bot_msg.edit_text(f"{e}\n\n{bot_text.premium_warning}", reply_markup=premium_button())
     except Exception:
         bot_msg.edit_text(f"Download failed!❌\n\n`{traceback.format_exc()[-2000:]}`", disable_web_page_preview=True)
     logging.info("YouTube celery tasks ended.")
@@ -157,10 +166,8 @@ def ytdl_download_entrance(client: Client, bot_msg: types.Message, url: str, mod
             ytdl_normal_download(client, bot_msg, url)
     except FileTooBig as e:
         logging.warning("Seeking for help from premium user...")
-        bot_msg.edit_text(f"{e}\n\nPlease wait, I will try to send it as premium user")
         # this is only for normal node. Celery node will need to do it in celery tasks
-        data = {"url": url, "user_id": bot_msg.chat.id}
-        client.send_message(PREMIUM_USER, json.dumps(data), disable_notification=True, disable_web_page_preview=True)
+        bot_msg.edit_text(f"{e}\n\n{bot_text.premium_warning}", reply_markup=premium_button())
     except Exception as e:
         logging.error("Failed to download %s, error: %s", url, e)
         bot_msg.edit_text(f"Download failed!❌\n\n`{traceback.format_exc()[-2000:]}`", disable_web_page_preview=True)
@@ -502,7 +509,7 @@ def run_celery():
     threading.Thread(target=run_celery, daemon=True).start()
 
     scheduler = BackgroundScheduler(timezone="Europe/London")
-    # scheduler.add_job(auto_restart, "interval", seconds=900)
+    scheduler.add_job(auto_restart, "interval", seconds=900)
     scheduler.start()
 
     idle()

ytdlbot/ytdl_bot.py

@@ -8,6 +8,7 @@
 __author__ = "Benny <[email protected]>"
 
 import contextlib
+import json
 import logging
 import os
 import random
@@ -316,6 +317,17 @@ def tronpayment_btn_calback(client: Client, callback_query: types.CallbackQuery)
         client.send_photo(chat_id, bio, caption=f"Send any amount of TRX to `{addr}`")
 
 
+@app.on_callback_query(filters.regex(r"premium.*"))
+def premium_click(client: Client, callback_query: types.CallbackQuery):
+    data = callback_query.data
+    if data == "premium-yes":
+        replied = callback_query.message.reply_to_message
+        data = {"url": replied.text, "user_id": callback_query.message.chat.id}
+        client.send_message(PREMIUM_USER, json.dumps(data), disable_notification=True, disable_web_page_preview=True)
+    else:
+        callback_query.answer("Cancelled.")
+
+
 @app.on_callback_query(filters.regex(r"bot-payments-.*"))
 def bot_payment_btn_calback(client: Client, callback_query: types.CallbackQuery):
     callback_query.answer("Generating invoice...")