premium 4GB max fix

ytdlbot/config.py

@@ -20,8 +20,10 @@
 
 REDIS = os.getenv("REDIS", "redis")
 
-TG_MAX_SIZE = 2000 * 1024 * 1024
-# TG_MAX_SIZE = 10 * 1024 * 1024
+TG_PREMIUM_MAX_SIZE = 4000 * 1024 * 1024
+TG_NORMAL_MAX_SIZE = 2000 * 1024 * 1024
+# TG_NORMAL_MAX_SIZE = 10 * 1024 * 1024
+
 
 EXPIRE = 24 * 3600
 

ytdlbot/constant.py

@@ -45,6 +45,8 @@ class BotText:
 
 4. Download for paid user will be automatically changed to Local mode to avoid queuing.
 
+5. Paid user can download files larger than 2GB.
+
 **Price:**
 valid permanently
 1. 1 USD == {TOKEN_PRICE} tokens

ytdlbot/downloader.py

@@ -32,7 +32,8 @@
     ENABLE_ARIA2,
     ENABLE_FFMPEG,
     PREMIUM_USER,
-    TG_MAX_SIZE,
+    TG_NORMAL_MAX_SIZE,
+    TG_PREMIUM_MAX_SIZE,
     FileTooBig,
     IPv6,
 )
@@ -123,7 +124,9 @@ def download_hook(d: dict, bot_msg):
     if d["status"] == "downloading":
         downloaded = d.get("downloaded_bytes", 0)
         total = d.get("total_bytes") or d.get("total_bytes_estimate", 0)
-        if total > TG_MAX_SIZE:
+        if total > TG_PREMIUM_MAX_SIZE:
+            raise Exception(f"There's no way to handle a file of {sizeof_fmt(total)}.")
+        if total > TG_NORMAL_MAX_SIZE:
             msg = f"Your download file size {sizeof_fmt(total)} is too large for Telegram."
             if PREMIUM_USER:
                 raise FileTooBig(msg)
@@ -292,10 +295,10 @@ def split_large_video(video_paths: list):
     split = False
     for original_video in video_paths:
         size = os.stat(original_video).st_size
-        if size > TG_MAX_SIZE:
+        if size > TG_NORMAL_MAX_SIZE:
             split = True
             logging.warning("file is too large %s, splitting...", size)
-            subprocess.check_output(f"sh split-video.sh {original_video} {TG_MAX_SIZE * 0.95} ".split())
+            subprocess.check_output(f"sh split-video.sh {original_video} {TG_NORMAL_MAX_SIZE * 0.95} ".split())
             os.remove(original_video)
 
     if split and original_video:

ytdlbot/premium.py

@@ -68,12 +68,13 @@ async def hello(client: Client, message: types.Message):
     settings = payment.get_user_settings(user_id)
     video_path = next(pathlib.Path(tempdir.name).glob("*"))
     logging.info("Final filesize is %s", sizeof_fmt(video_path.stat().st_size))
+    caption = "Powered by @benny_ytdlbot "
     if settings[2] == "audio":
         logging.info("Sending as audio")
         await client.send_audio(
             BOT_ID,
             video_path.as_posix(),
-            caption="Powered by ytdlbot ",
+            caption=caption,
             file_name=f"{user_id}.mp3",
             progress=upload_hook,
         )
@@ -82,7 +83,7 @@ async def hello(client: Client, message: types.Message):
         await client.send_document(
             BOT_ID,
             video_path.as_posix(),
-            caption="Powered by ytdlbot",
+            caption=caption,
             file_name=f"{user_id}.mp4",
             progress=upload_hook,
         )
@@ -91,13 +92,14 @@ async def hello(client: Client, message: types.Message):
         await client.send_video(
             BOT_ID,
             video_path.as_posix(),
-            caption="Powered by ytdlbot",
+            caption=caption,
             supports_streaming=True,
             file_name=f"{user_id}.mp4",
             progress=upload_hook,
         )
 
     tempdir.cleanup()
+    logging.info("Finished sending %s", url)
 
 
 if __name__ == "__main__":

ytdlbot/tasks.py

@@ -80,10 +80,18 @@ def retrieve_message(chat_id: int, message_id: int) -> types.Message | Any:
 
 def premium_button(user_id):
     redis = Redis()
+    payment = Payment()
     used = redis.r.hget("premium", user_id)
     ban = redis.r.hget("ban", user_id)
-    if used or ban:
+    paid_token = payment.get_pay_token(user_id)
+
+    if ban:
         return None
+    # vip mode: vip user can use once per day, normal user can't use
+    # non vip mode: everyone can use once per day
+    if used or (ENABLE_VIP and paid_token == 0):
+        return None
+
     markup = types.InlineKeyboardMarkup(
         [
             [
@@ -108,7 +116,7 @@ def ytdl_download_task(chat_id: int, message_id: int, url: str):
         if markup:
             bot_msg.edit_text(f"{e}\n\n{bot_text.premium_warning}", reply_markup=markup)
         else:
-            bot_msg.edit_text(f"{e}\n\n Big file download is not available now, please try again later.")
+            bot_msg.edit_text(f"{e}\nBig file download is not available now. Please /buy or try again later ")
     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.")
@@ -180,7 +188,7 @@ def ytdl_download_entrance(client: Client, bot_msg: types.Message, url: str, mod
         if markup:
             bot_msg.edit_text(f"{e}\n\n{bot_text.premium_warning}", reply_markup=markup)
         else:
-            bot_msg.edit_text(f"{e}\n\n Big file download is not available now, please try again later.")
+            bot_msg.edit_text(f"{e}\nBig file download is not available now. Please /buy or try again later ")
     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)