diff --git a/.gitignore b/.gitignore index 56388a92..778fa339 100644 --- a/.gitignore +++ b/.gitignore @@ -176,3 +176,4 @@ reinforcement/* /ytdlbot/premium.session /dump.rdb /ytdlbot/premium.session-journal +/ytdlbot/main.session-journal diff --git a/ytdlbot/constant.py b/ytdlbot/constant.py index 3938d13d..5a2f8adf 100644 --- a/ytdlbot/constant.py +++ b/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 diff --git a/ytdlbot/database.py b/ytdlbot/database.py index 0dd6f8d5..4f5df0fe 100644 --- a/ytdlbot/database.py +++ b/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) diff --git a/ytdlbot/premium.py b/ytdlbot/premium.py index 90346284..622b8936 100644 --- a/ytdlbot/premium.py +++ b/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__": diff --git a/ytdlbot/tasks.py b/ytdlbot/tasks.py index ea951c77..1ceb24a2 100644 --- a/ytdlbot/tasks.py +++ b/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) diff --git a/ytdlbot/ytdl_bot.py b/ytdlbot/ytdl_bot.py index 0ebaf5b3..4d51839a 100644 --- a/ytdlbot/ytdl_bot.py +++ b/ytdlbot/ytdl_bot.py @@ -8,6 +8,7 @@ __author__ = "Benny " 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(