limitations on premium download

.gitignore

@@ -176,3 +176,4 @@ reinforcement/*
 /ytdlbot/premium.session
 /dump.rdb
 /ytdlbot/premium.session-journal
+/ytdlbot/main.session-journal

ytdlbot/constant.py

@@ -85,7 +85,9 @@ class BotText:
 
     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.
+    This is an experimental feature so you can only use it once per day.
+    Also, the premium user will know who you are and what you are downloading. 
+    You may be banned if you abuse this feature.
     """
 
     @staticmethod

ytdlbot/database.py

@@ -182,6 +182,8 @@ def reset_today(self):
             if k.startswith("today"):
                 self.r.hdel("metrics", k)
 
+        self.r.delete("premium")
+
     def user_count(self, user_id):
         self.r.hincrby("metrics", user_id)
 

ytdlbot/premium.py

@@ -8,14 +8,13 @@
 import logging
 import pathlib
 import tempfile
-from unittest.mock import MagicMock
 
 import yt_dlp
 from pyrogram import Client, filters, types
 
 from config import APP_HASH, APP_ID, PYRO_WORKERS, TOKEN
-from limit import Payment
-from utils import apply_log_formatter
+from limit import Payment, Redis
+from utils import apply_log_formatter, sizeof_fmt
 
 apply_log_formatter()
 app = Client("premium", APP_ID, APP_HASH, workers=PYRO_WORKERS)
@@ -26,11 +25,11 @@
 def download_hook(d: dict):
     downloaded = d.get("downloaded_bytes", 0)
     total = d.get("total_bytes") or d.get("total_bytes_estimate", 0)
-    print(downloaded, total)
+    logging.info("Downloaded %s/%s, %.2f%% complete", downloaded, total, downloaded / total * 100)
 
 
 async def upload_hook(current, total):
-    print(current, total)
+    logging.info("Uploaded %s/%s, %.2f%% complete", current, total, current / total * 100)
 
 
 @app.on_message(filters.user(BOT_ID) & filters.incoming)
@@ -66,18 +65,8 @@ async def hello(client: Client, message: types.Message):
     payment = Payment()
     settings = payment.get_user_settings(user_id)
     video_path = next(pathlib.Path(tempdir.name).glob("*"))
-    logging.info("filesize: %s", video_path.stat().st_size)
-    if settings[2] == "video" or isinstance(settings[2], MagicMock):
-        logging.info("Sending as video")
-        await client.send_video(
-            BOT_ID,
-            video_path.as_posix(),
-            caption="Powered by ytdlbot",
-            supports_streaming=True,
-            file_name=f"{user_id}.mp4",
-            progress=upload_hook,
-        )
-    elif settings[2] == "audio":
+    logging.info("Final filesize is %s", sizeof_fmt(video_path.stat().st_size))
+    if settings[2] == "audio":
         logging.info("Sending as audio")
         await client.send_audio(
             BOT_ID,
@@ -96,9 +85,19 @@ async def hello(client: Client, message: types.Message):
             progress=upload_hook,
         )
     else:
-        logging.error("Send type is not video or audio")
+        logging.info("Sending as video")
+        await client.send_video(
+            BOT_ID,
+            video_path.as_posix(),
+            caption="Powered by ytdlbot",
+            supports_streaming=True,
+            file_name=f"{user_id}.mp4",
+            progress=upload_hook,
+        )
 
     tempdir.cleanup()
+    redis = Redis()
+    redis.r.hset("premium", user_id, 1)
 
 
 if __name__ == "__main__":

ytdlbot/tasks.py

@@ -78,7 +78,12 @@ def retrieve_message(chat_id: int, message_id: int) -> types.Message | Any:
         return bot.get_messages(chat_id, message_id)
 
 
-def premium_button():
+def premium_button(user_id):
+    redis = Redis()
+    used = redis.r.hget("premium", user_id)
+    ban = redis.r.hget("ban", user_id)
+    if used or ban:
+        return None
     markup = types.InlineKeyboardMarkup(
         [
             [
@@ -99,7 +104,11 @@ 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\n{bot_text.premium_warning}", reply_markup=premium_button())
+        markup = premium_button(chat_id)
+        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.")
     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.")
@@ -167,7 +176,11 @@ def ytdl_download_entrance(client: Client, bot_msg: types.Message, url: str, mod
     except FileTooBig as e:
         logging.warning("Seeking for help from premium user...")
         # this is only for normal node. Celery node will need to do it in celery tasks
-        bot_msg.edit_text(f"{e}\n\n{bot_text.premium_warning}", reply_markup=premium_button())
+        markup = premium_button(chat_id)
+        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.")
     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)

ytdlbot/ytdl_bot.py

@@ -8,6 +8,7 @@
 __author__ = "Benny <[email protected]>"
 
 import contextlib
+import json
 import logging
 import os
 import random
@@ -320,12 +321,15 @@ def tronpayment_btn_calback(client: Client, callback_query: types.CallbackQuery)
 def premium_click(client: Client, callback_query: types.CallbackQuery):
     data = callback_query.data
     if data == "premium-yes":
-        callback_query.answer("developing....please wait for next release")
-        # 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)
+        callback_query.answer("Seeking premium user...")
+        callback_query.message.edit_text("Please wait patiently...no progress bar will be shown.")
+        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.")
+        original_text = callback_query.message.text
+        callback_query.message.edit_text(original_text.split("\n")[0])
 
 
 @app.on_callback_query(filters.regex(r"bot-payments-.*"))
@@ -365,6 +369,15 @@ def premium_forward(client: Client, message: types.Message):
     client.forward_messages(target_user, message.chat.id, message.id)
 
 
+@app.on_message(filters.command(["ban"]) & filters.user(PREMIUM_USER))
+def ban_handler(client: Client, message: types.Message):
+    replied = message.reply_to_message.text
+    user_id = json.loads(replied).get("user_id")
+    redis = Redis()
+    redis.r.hset("ban", user_id, 1)
+    message.reply_text(f"Done, banned {user_id}.", quote=True)
+
+
 def generate_invoice(amount: int, title: str, description: str, payload: str):
     invoice = raw_types.input_media_invoice.InputMediaInvoice(
         invoice=raw_types.invoice.Invoice(