From f4abe3255c4ed6fef70cc3d6abe98d015c03ab41 Mon Sep 17 00:00:00 2001 From: BennyThink Date: Thu, 20 Jan 2022 21:01:58 +0800 Subject: [PATCH] add direct download support --- README.md | 1 + ytdlbot/tasks.py | 86 ++++++++++++++++++++++++++++++++++++++++----- ytdlbot/ytdl_bot.py | 21 +++++++++-- 3 files changed, 97 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 340861e..0ba293a 100644 --- a/README.md +++ b/README.md @@ -216,6 +216,7 @@ ytdl - Download video in group vip - Join VIP terms - View Terms of Service settings - Set your preference +direct - Download file directly ``` # Test data diff --git a/ytdlbot/tasks.py b/ytdlbot/tasks.py index 8c61f27..d9a3772 100644 --- a/ytdlbot/tasks.py +++ b/ytdlbot/tasks.py @@ -10,19 +10,23 @@ __author__ = "Benny " import logging import os import pathlib +import re import tempfile import threading import time +from urllib.parse import quote_plus +import requests from celery import Celery from pyrogram import idle from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup from client_init import create_app -from config import BROKER, ENABLE_CELERY, OWNER, WORKERS +from config import BROKER, ENABLE_CELERY, ENABLE_VIP, OWNER, WORKERS from constant import BotText from db import Redis from downloader import convert_flac, sizeof_fmt, upload_hook, ytdl_download +from limit import VIP from utils import (apply_log_formatter, customize_logger, get_metadata, get_user_settings) @@ -47,11 +51,11 @@ def get_messages(chat_id, message_id): @app.task() -def download_task(chat_id, message_id, url): - logging.info("celery tasks started for %s", url) +def ytdl_download_task(chat_id, message_id, url): + logging.info("YouTube celery tasks started for %s", url) bot_msg = get_messages(chat_id, message_id) - normal_download(bot_msg, celery_client, url) - logging.info("celery tasks ended.") + ytdl_normal_download(bot_msg, celery_client, url) + logging.info("YouTube celery tasks ended.") @app.task() @@ -62,11 +66,28 @@ def audio_task(chat_id, message_id): logging.info("Audio celery tasks ended.") -def download_entrance(bot_msg, client, url): +@app.task() +def direct_download_task(chat_id, message_id, url): + logging.info("Direct download celery tasks started for %s", url) + bot_msg = get_messages(chat_id, message_id) + direct_normal_download(bot_msg, celery_client, url) + logging.info("Direct download celery tasks ended.") + + +def ytdl_download_entrance(bot_msg, client, url): if ENABLE_CELERY: - download_task.delay(bot_msg.chat.id, bot_msg.message_id, url) + ytdl_download_task.delay(bot_msg.chat.id, bot_msg.message_id, url) else: - normal_download(bot_msg, client, url) + ytdl_normal_download(bot_msg, client, url) + + +def direct_download_entrance(bot_msg, client, url): + if ENABLE_CELERY: + # TODO disable it for now + direct_normal_download(bot_msg, client, url) + # direct_download_task.delay(bot_msg.chat.id, bot_msg.message_id, url) + else: + direct_normal_download(bot_msg, client, url) def audio_entrance(bot_msg): @@ -76,6 +97,53 @@ def audio_entrance(bot_msg): normal_audio(bot_msg) +def direct_normal_download(bot_msg, client, url): + 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/80.0.3987.149 Safari/537.36"} + vip = VIP() + + if ENABLE_VIP: + remain, _, _ = vip.check_remaining_quota(chat_id) + try: + head_req = requests.head(url, headers=headers) + length = int(head_req.headers.get("content-length")) + except: + length = 0 + if remain < length: + bot_msg.reply_text(f"Sorry, you have reached your quota.\n") + return + + req = "" + try: + req = requests.get(url, headers=headers) + filename = re.findall("filename=(.+)", req.headers.get("content-disposition"))[0] + except TypeError: + filename = getattr(req, "url", "").rsplit("/")[-1] + except Exception as e: + bot_msg.edit_text(f"Download failed!❌\n\n```{e}```", disable_web_page_preview=True) + return + + if not filename: + filename = quote_plus(url) + + with tempfile.TemporaryDirectory() as f: + filepath = f"{f}/{filename}" + print(filepath) + with open(filepath, "wb") as fp: + fp.write(req.content) + logging.info("Downloaded file %s", filename) + st_size = os.stat(filepath).st_size + if ENABLE_VIP: + vip.use_quota(chat_id, st_size) + client.send_chat_action(chat_id, "upload_document") + client.send_document(bot_msg.chat.id, filepath, + caption=f"filesize: {sizeof_fmt(st_size)}", + progress=upload_hook, progress_args=(bot_msg,), + ) + bot_msg.edit_text(f"Download success!✅") + + def normal_audio(bot_msg): chat_id = bot_msg.chat.id mp4_name = bot_msg.video.file_name # 'youtube-dl_test_video_a.mp4' @@ -103,7 +171,7 @@ def get_worker_status(username): return f"Downloaded by {me.mention()}" -def normal_download(bot_msg, client, url): +def ytdl_normal_download(bot_msg, client, url): chat_id = bot_msg.chat.id temp_dir = tempfile.TemporaryDirectory() diff --git a/ytdlbot/ytdl_bot.py b/ytdlbot/ytdl_bot.py index e9f5688..fe99d64 100644 --- a/ytdlbot/ytdl_bot.py +++ b/ytdlbot/ytdl_bot.py @@ -24,7 +24,8 @@ from config import (AUTHORIZED_USER, ENABLE_CELERY, ENABLE_VIP, OWNER, from constant import BotText from db import InfluxDB, MySQL, Redis from limit import verify_payment -from tasks import audio_entrance, download_entrance +from tasks import (audio_entrance, direct_download_entrance, + ytdl_download_entrance) from utils import (auto_restart, customize_logger, get_revision, get_user_settings, set_user_settings) @@ -119,6 +120,22 @@ def terms_handler(client: "Client", message: "types.Message"): client.send_message(chat_id, bot_text.terms) +@app.on_message(filters.command(["direct"])) +def direct_handler(client: "Client", message: "types.Message"): + chat_id = message.from_user.id + client.send_chat_action(chat_id, "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 + + bot_msg = message.reply_text("Request received.", quote=True) + Redis().update_metrics("direct_request") + direct_download_entrance(bot_msg, client, url) + + @app.on_message(filters.command(["settings"])) def settings_handler(client: "Client", message: "types.Message"): chat_id = message.chat.id @@ -176,7 +193,7 @@ def download_handler(client: "Client", message: "types.Message"): text = bot_text.get_receive_link_text() bot_msg: typing.Union["types.Message", "typing.Any"] = message.reply_text(text, quote=True) client.send_chat_action(chat_id, 'upload_video') - download_entrance(bot_msg, client, url) + ytdl_download_entrance(bot_msg, client, url) @app.on_callback_query(filters.regex(r"document|video"))