From d951513a3e0ae6a1b4f5d804ae06a98736e808b4 Mon Sep 17 00:00:00 2001 From: BennyThink Date: Wed, 25 Aug 2021 22:13:25 +0800 Subject: [PATCH] quota bug fix --- README.md | 2 +- downloader.py | 71 ++++++++++++++++++++++++++++++++++++++------------- ytdl.py | 19 ++++---------- 3 files changed, 59 insertions(+), 33 deletions(-) diff --git a/README.md b/README.md index adbe510..50940eb 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ see [here](https://github.com/tgbot-collection/BotsRunner) start - Let's start about - Want to contribute? ping - Bot running status -help - Anything troubles you/ +help - Anything troubles you? ytdl - Download video in group vip - Join VIP terms - View Terms of Service diff --git a/downloader.py b/downloader.py index da98335..2d3f4e7 100644 --- a/downloader.py +++ b/downloader.py @@ -11,6 +11,7 @@ import logging import os import pathlib import subprocess +import time import traceback import fakeredis @@ -18,7 +19,7 @@ import filetype import youtube_dl from youtube_dl import DownloadError -from limit import VIP +from limit import VIP, Redis r = fakeredis.FakeStrictRedis() EXPIRE = 5 @@ -33,7 +34,7 @@ def sizeof_fmt(num: int, suffix='B'): def edit_text(bot_msg, text): - key = bot_msg.message_id + key = f"{bot_msg.chat.id}-{bot_msg.message_id}" # if the key exists, we shouldn't send edit message if not r.exists(key): r.set(key, "ok", ex=EXPIRE) @@ -44,12 +45,16 @@ 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) + # total = 0 filesize = sizeof_fmt(total) if total > 2 * 1024 * 1024 * 1024: raise Exception("\n\nYour video is too large. %s will exceed Telegram's max limit 2GiB" % filesize) percent = d.get("_percent_str", "N/A") speed = d.get("_speed_str", "N/A") + result, err_msg = check_quota(total, bot_msg.chat.id) + if result is False: + raise ValueError(err_msg) text = f'[{filesize}]: Downloading {percent} - {downloaded}/{total} @ {speed}' edit_text(bot_msg, text) @@ -60,6 +65,20 @@ def upload_hook(current, total, bot_msg): edit_text(bot_msg, text) +def check_quota(file_size, chat_id) -> ("bool", "str"): + remain, _, ttl = VIP().check_remaining_quota(chat_id) + if file_size > remain: + refresh_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(ttl + time.time())) + err = f"Quota exceed, you have {sizeof_fmt(remain)} remaining, " \ + f"but you want to download a video with {sizeof_fmt(file_size)} in size. \n" \ + f"Try again in {ttl} seconds({refresh_time})" + logging.warning(err) + Redis().update_metrics("quota_exceed") + return False, err + else: + return True, "" + + def convert_to_mp4(resp: dict): default_type = ["video/x-flv"] if resp["status"]: @@ -79,12 +98,13 @@ def convert_to_mp4(resp: dict): def ytdl_download(url, tempdir, bm) -> dict: - response = dict(status=None, error=None, filepath=[]) + chat_id = bm.chat.id + response = {"status": True, "error": "", "filepath": []} output = os.path.join(tempdir, '%(title)s.%(ext)s') ydl_opts = { 'progress_hooks': [lambda d: download_hook(d, bm)], 'outtmpl': output, - 'restrictfilenames': True, + 'restrictfilenames': False, 'quiet': True } formats = [ @@ -92,7 +112,7 @@ def ytdl_download(url, tempdir, bm) -> dict: "bestvideo[vcodec^=avc]+bestaudio[acodec^=mp4a]/best[vcodec^=avc]/best", "" ] - success, err = None, None + # TODO it appears twitter download on macOS will fail. Don't know why...Linux's fine. for f in formats: if f: ydl_opts["format"] = f @@ -100,22 +120,37 @@ def ytdl_download(url, tempdir, bm) -> dict: logging.info("Downloading for %s with format %s", url, f) with youtube_dl.YoutubeDL(ydl_opts) as ydl: ydl.download([url]) - success = True + response["status"] = True + response["error"] = "" break - except DownloadError: - err = traceback.format_exc() - logging.error("Download failed for %s ", url) - if success: - response["status"] = True - for i in os.listdir(tempdir): - p = os.path.join(tempdir, i) + except DownloadError as e: + err = str(e) + logging.error("Download failed for %s ", url) + response["status"] = False + response["error"] = err + # can't return here + except ValueError as e: + response["status"] = False + response["error"] = str(e) + + logging.info(response) + if response["status"] is False: + return response + + for i in os.listdir(tempdir): + remain, _, ttl = VIP().check_remaining_quota(chat_id) + p: "str" = os.path.join(tempdir, i) + file_size = os.stat(p).st_size + result, err_msg = check_quota(file_size, chat_id) + if result is False: + response["status"] = False + response["error"] = err_msg + else: + VIP().use_quota(bm.chat.id, file_size) + response["status"] = True response["filepath"].append(p) - VIP().use_quota(bm.chat.id, os.stat(p).st_size) - # break - else: - response["status"] = False - response["error"] = err + # convert format if necessary convert_to_mp4(response) return response diff --git a/ytdl.py b/ytdl.py index a10588b..c4a4235 100644 --- a/ytdl.py +++ b/ytdl.py @@ -12,7 +12,6 @@ import os import pathlib import re import tempfile -import time import typing import ffmpeg @@ -27,7 +26,7 @@ from limit import VIP, Redis, verify_payment def get_metadata(video_path): - height, width, duration = 1280, 720, 0 + width, height, duration = 1280, 720, 0 try: video_streams = ffmpeg.probe(video_path, select_streams="v") for item in video_streams.get("streams", []): @@ -120,16 +119,7 @@ def download_handler(client: "Client", message: "types.Message"): # check remaining quota chat_id = message.chat.id Redis().user_count(chat_id) - used, _, ttl = bot_text.return_remaining_quota(chat_id) - # TODO bug here: if user have 10MB of quota, and he is downloading a playlist toal 10G - # then it won't stop him from downloading - # the same applies to 10MB of quota, but try to download 20MB video - if used <= 0: - refresh_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(ttl + time.time())) - logging.error("quota exceed for %s, try again in %s seconds(%s)", chat_id, ttl, refresh_time) - message.reply_text(f"Quota exceed, try again in {ttl} seconds({refresh_time})", quote=True) - Redis().update_metrics("quota_exceed") - return + if message.chat.type != "private" and not message.text.lower().startswith("/ytdl"): logging.warning("%s, it's annoying me...🙄️ ", message.text) return @@ -147,6 +137,7 @@ def download_handler(client: "Client", message: "types.Message"): if not VIP().check_vip(chat_id): message.reply_text("Playlist download is only available to VIP users. Join /vip now.", quote=True) return + Redis().update_metrics("video_request") bot_msg: typing.Union["types.Message", "typing.Any"] = message.reply_text("Processing", quote=True) client.send_chat_action(chat_id, 'upload_video') @@ -169,9 +160,9 @@ def download_handler(client: "Client", message: "types.Message"): if result["status"]: client.send_chat_action(chat_id, 'upload_document') video_paths = result["filepath"] + bot_msg.edit_text('Download complete. Sending now...') for video_path in video_paths: filename = pathlib.Path(video_path).name - bot_msg.edit_text('Download complete. Sending now...') remain = bot_text.remaining_quota_caption(chat_id) size = sizeof_fmt(os.stat(video_path).st_size) meta = get_metadata(video_path) @@ -187,7 +178,7 @@ def download_handler(client: "Client", message: "types.Message"): else: client.send_chat_action(chat_id, 'typing') tb = result["error"][0:4000] - bot_msg.edit_text(f"{url} download failed❌：\n```{tb}```") + bot_msg.edit_text(f"Download failed!❌\n\n```{tb}```", disable_web_page_preview=True) temp_dir.cleanup()