diff --git a/ytdlbot/tasks.py b/ytdlbot/tasks.py index 7db8c73c..00a1c379 100644 --- a/ytdlbot/tasks.py +++ b/ytdlbot/tasks.py @@ -214,12 +214,12 @@ def ytdl_download_entrance(client: Client, bot_msg: types.Message, url: str, mod bot_msg.edit_text(f"Download failed!❌\n\n`{traceback.format_exc()[-2000:]}`", disable_web_page_preview=True) -def direct_download_entrance(client: Client, bot_msg: typing.Union[types.Message, typing.Coroutine], url: str, custom_filename): +def direct_download_entrance(client: Client, bot_msg: typing.Union[types.Message, typing.Coroutine], url: str, new_name): if ENABLE_CELERY: - direct_normal_download(client, bot_msg, url, custom_filename) + direct_normal_download(client, bot_msg, url, new_name) # direct_download_task.delay(bot_msg.chat.id, bot_msg.id, url) else: - direct_normal_download(client, bot_msg, url, custom_filename) + direct_normal_download(client, bot_msg, url, new_name) def leech_download_entrance(client: Client, bot_msg: typing.Union[types.Message, typing.Coroutine], url: str): @@ -269,7 +269,7 @@ def audio_entrance(client: Client, bot_msg: types.Message): normal_audio(client, bot_msg) -def direct_normal_download(client: Client, bot_msg: typing.Union[types.Message, typing.Coroutine], url: str, custom_filename): +def direct_normal_download(client: Client, bot_msg: typing.Union[types.Message, typing.Coroutine], url: str, new_name): 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/126.0.0.0 Safari/537.36" @@ -284,8 +284,8 @@ def direct_normal_download(client: Client, bot_msg: typing.Union[types.Message, bot_msg.edit_text(f"Download failed!❌\n\n```{e}```", disable_web_page_preview=True) return - if custom_filename: - filename = custom_filename + if new_name: + filename = new_name else: filename = extract_filename(req) diff --git a/ytdlbot/utils.py b/ytdlbot/utils.py index 1d74c006..1075578c 100644 --- a/ytdlbot/utils.py +++ b/ytdlbot/utils.py @@ -266,5 +266,22 @@ def extract_filename(response): return filename +def extract_url_and_name(message_text): + # Regular expression to match the URL + url_pattern = r'(https?://[^\s]+)' + # Regular expression to match the new name after '-n' + name_pattern = r'-n\s+([^\s]+)' + + # Find the URL in the message_text + url_match = re.search(url_pattern, message_text) + url = url_match.group(0) if url_match else None + + # Find the new name in the message_text + name_match = re.search(name_pattern, message_text) + new_name = name_match.group(1) if name_match else None + + return url, new_name + + if __name__ == "__main__": auto_restart() diff --git a/ytdlbot/ytdl_bot.py b/ytdlbot/ytdl_bot.py index ce16f280..ebcdc0ad 100644 --- a/ytdlbot/ytdl_bot.py +++ b/ytdlbot/ytdl_bot.py @@ -61,7 +61,13 @@ ytdl_download_entrance, spdl_download_entrance, ) -from utils import auto_restart, clean_tempfile, customize_logger, get_revision +from utils import ( + auto_restart, + clean_tempfile, + customize_logger, + get_revision, + extract_url_and_name +) logging.info("Authorized users are %s", AUTHORIZED_USER) customize_logger(["pyrogram.client", "pyrogram.session.session", "pyrogram.connection.connection"]) @@ -219,7 +225,7 @@ def send_message_and_measure_ping(): client.edit_message_text(chat_id=reply.chat.id, message_id=reply.id, text="Ping Calculation Complete.") time.sleep(1) client.delete_messages(chat_id=reply.chat.id, message_ids=reply.id) - + thread = threading.Thread(target=send_message_and_measure_ping) thread.start() @@ -269,63 +275,6 @@ def clear_history(client: Client, message: types.Message): message.reply_text("History cleared.", quote=True) -@app.on_message(filters.command(["spdl"])) -def spdl_handler(client: Client, message: types.Message): - redis = Redis() - chat_id = message.from_user.id - client.send_chat_action(chat_id, enums.ChatAction.TYPING) - url = re.sub(r"/spdl\s*", "", message.text) - logging.info("spdl start %s", url) - if not re.findall(r"^https?://", url.lower()): - redis.update_metrics("bad_request") - message.reply_text("Something wrong 🤔.\nCheck your URL and send me again.", quote=True) - return - - bot_msg = message.reply_text("Request received.", quote=True) - redis.update_metrics("spdl_request") - spdl_download_entrance(client, bot_msg, url) - - -@app.on_message(filters.command(["direct"])) -def direct_handler(client: Client, message: types.Message): - redis = Redis() - chat_id = message.from_user.id - client.send_chat_action(chat_id, enums.ChatAction.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 - - url_parts = url.split(" -n ", maxsplit=1) - url = url_parts[0].strip() # Assuming space after -n - custom_filename = url_parts[1].strip() if len(url_parts) > 1 else None - bot_msg = message.reply_text("Request received.", quote=True) - redis.update_metrics("direct_request") - direct_download_entrance(client, bot_msg, url, custom_filename) - - -@app.on_message(filters.command(["leech"])) -def leech_handler(client: Client, message: types.Message): - if not ENABLE_ARIA2: - message.reply_text("Aria2 Not Enabled.", quote=True) - return - redis = Redis() - chat_id = message.from_user.id - client.send_chat_action(chat_id, enums.ChatAction.TYPING) - url = re.sub(r"/leech\s*", "", message.text) - logging.info("leech using aria2 start %s", url) - if not re.findall(r"^https?://", url.lower()): - redis.update_metrics("bad_request") - message.reply_text("Send me a correct LINK.", quote=True) - return - - bot_msg = message.reply_text("Request received.", quote=True) - redis.update_metrics("leech_request") - leech_download_entrance(client, bot_msg, url) - - @app.on_message(filters.command(["settings"])) def settings_handler(client: Client, message: types.Message): chat_id = message.chat.id @@ -509,6 +458,81 @@ def search_ytb(kw: str): return text +@app.on_message(filters.command(["spdl"])) +def spdl_handler(client: Client, message: types.Message): + redis = Redis() + chat_id = message.from_user.id + client.send_chat_action(chat_id, enums.ChatAction.TYPING) + message_text = message.text + url, new_name = extract_url_and_name(message_text) + logging.info("spdl start %s", url) + if url is None or not re.findall(r"^https?://", url.lower()): + redis.update_metrics("bad_request") + message.reply_text("Something wrong 🤔.\nCheck your URL and send me again.", quote=True) + return + + bot_msg = message.reply_text("Request received.", quote=True) + redis.update_metrics("spdl_request") + spdl_download_entrance(client, bot_msg, url) + + +@app.on_message(filters.command(["direct"])) +def direct_handler(client: Client, message: types.Message): + redis = Redis() + chat_id = message.from_user.id + client.send_chat_action(chat_id, enums.ChatAction.TYPING) + message_text = message.text + url, new_name = extract_url_and_name(message_text) + logging.info("direct start %s", url) + if url is None or 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(client, bot_msg, url, new_name) + + +@app.on_message(filters.command(["leech"])) +def leech_handler(client: Client, message: types.Message): + if not ENABLE_ARIA2: + message.reply_text("Aria2 Not Enabled.", quote=True) + return + redis = Redis() + chat_id = message.from_user.id + client.send_chat_action(chat_id, enums.ChatAction.TYPING) + message_text = message.text + url, new_name = extract_url_and_name(message_text) + logging.info("leech using aria2 start %s", url) + if url is None or not re.findall(r"^https?://", url.lower()): + redis.update_metrics("bad_request") + message.reply_text("Send me a correct LINK.", quote=True) + return + + bot_msg = message.reply_text("Request received.", quote=True) + redis.update_metrics("leech_request") + leech_download_entrance(client, bot_msg, url) + + +@app.on_message(filters.command(["ytdl"])) +def ytdl_handler(client: Client, message: types.Message): + redis = Redis() + chat_id = message.from_user.id + client.send_chat_action(chat_id, enums.ChatAction.TYPING) + message_text = message.text + url, new_name = extract_url_and_name(message_text) + logging.info("ytdl start %s", url) + if url is None or not re.findall(r"^https?://", url.lower()): + redis.update_metrics("bad_request") + message.reply_text("Something wrong 🤔.\nCheck your URL and send me again.", quote=True) + return + + bot_msg = message.reply_text("Request received.", quote=True) + redis.update_metrics("ytdl_request") + ytdl_download_entrance(client, bot_msg, url) + + @app.on_message(filters.incoming & (filters.text | filters.document)) @private_use def download_handler(client: Client, message: types.Message): @@ -524,7 +548,7 @@ def download_handler(client: Client, message: types.Message): contents = open(tf.name, "r").read() # don't know why urls = contents.split() else: - urls = [re.sub(r"/ytdl\s*", "", message.text)] + urls = [message.text] logging.info("start %s", urls) for url in urls: