diff --git a/README.md b/README.md index 4cafcee..517c162 100644 --- a/README.md +++ b/README.md @@ -29,9 +29,10 @@ Websites [supported by youtube-dl](https://ytdl-org.github.io/youtube-dl/support I don't have unlimited servers and bandwidth, so I have to make some restrictions. -* 10 GiB one-way traffic per 24 hours for each user +* 5 GiB one-way traffic per 24 hours for each user * maximum 5 minutes streaming conversion support * maximum 3 subscriptions +* limited request in certain time range You can choose to become 'VIP' if you really need large traffic. And also, you could always deploy your own bot. @@ -128,6 +129,7 @@ you can configure all the following environment variables: above `AUTHORIZED_USER` * ENABLE_CELERY: Distribution mode, default: disable. You'll can setup workers in different locations. +* ENABLE_FFMPEG: enable ffmpeg so Telegram can stream * MYSQL_HOST: you'll have to setup MySQL if you enable VIP mode * MYSQL_USER * MYSQL_PASS diff --git a/requirements.txt b/requirements.txt index 871fb03..df17a3c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,7 +11,7 @@ flower==1.2.0 psutil==5.9.4 influxdb==5.3.1 beautifulsoup4==4.11.1 -fakeredis==1.10.0 +fakeredis==1.10.1 supervisor==4.2.4 tgbot-ping==1.0.4 redis==4.3.3 @@ -20,3 +20,4 @@ tqdm==4.64.1 requests-toolbelt==0.10.1 ffpb==0.4.1 youtube-search-python==1.6.6 +token-bucket==0.3.0 diff --git a/ytdlbot/config.py b/ytdlbot/config.py index de406b1..7d31439 100644 --- a/ytdlbot/config.py +++ b/ytdlbot/config.py @@ -19,7 +19,7 @@ TOKEN = os.getenv("TOKEN", "3703WLI") REDIS = os.getenv("REDIS") # quota settings -QUOTA = int(os.getenv("QUOTA", 10 * 1024 * 1024 * 1024)) # 10G +QUOTA = int(os.getenv("QUOTA", 5 * 1024 * 1024 * 1024)) # 5G if os.uname().sysname == "Darwin": QUOTA = 10 * 1024 * 1024 # 10M @@ -58,3 +58,5 @@ ARCHIVE_ID = os.getenv("ARCHIVE_ID") IPv6 = os.getenv("IPv6", False) ENABLE_FFMPEG = os.getenv("ENABLE_FFMPEG", False) +RATE = float(os.getenv("RATE", 60 * 5)) +BURST = int(os.getenv("BURST", 3)) diff --git a/ytdlbot/constant.py b/ytdlbot/constant.py index bc8ceec..517084c 100644 --- a/ytdlbot/constant.py +++ b/ytdlbot/constant.py @@ -10,8 +10,8 @@ __author__ = "Benny " import os import time -from config import (AFD_LINK, COFFEE_LINK, ENABLE_CELERY, ENABLE_VIP, EX, - MULTIPLY, REQUIRED_MEMBERSHIP, USD2CNY) +from config import (AFD_LINK, BURST, COFFEE_LINK, ENABLE_CELERY, ENABLE_VIP, + EX, MULTIPLY, RATE, REQUIRED_MEMBERSHIP, USD2CNY) from db import InfluxDB from downloader import sizeof_fmt from limit import QUOTA, VIP @@ -33,6 +33,8 @@ every one can use this bot within **{sizeof_fmt(QUOTA)} of quota for every {int( 4. You can optionally choose to become 'VIP' user if you need more traffic. Type /vip for more information. 5. Source code for this bot will always stay open, here-> https://github.com/tgbot-collection/ytdlbot + +6. Request limit is applied for everyone, excluding VIP users. """ if ENABLE_VIP else "Help text" about = "YouTube-DL by @BennyThink. Open source on GitHub: https://github.com/tgbot-collection/ytdlbot" @@ -149,3 +151,5 @@ Sending format: **{1}** text += f"{status}{hostname} **{active}** {load} {rev}\n" return text + + too_fast = f"You have reached rate limit. Current rate limit is 1 request per {RATE} seconds, {BURST - 1} bursts." diff --git a/ytdlbot/utils.py b/ytdlbot/utils.py index 642d67f..37de1ab 100644 --- a/ytdlbot/utils.py +++ b/ytdlbot/utils.py @@ -205,7 +205,7 @@ class Detector: def idle_detector(self): mtime = os.stat("/var/log/ytdl.log").st_mtime cur_ts = time.time() - if cur_ts - mtime > 300: + if cur_ts - mtime > 1800: logging.warning("Potential crash detected by %s, it's time to commit suicide...", self.func_name()) return True diff --git a/ytdlbot/ytdl_bot.py b/ytdlbot/ytdl_bot.py index 7ce887a..1bac9b3 100644 --- a/ytdlbot/ytdl_bot.py +++ b/ytdlbot/ytdl_bot.py @@ -22,11 +22,12 @@ from pyrogram import Client, filters, types from pyrogram.errors.exceptions.bad_request_400 import UserNotParticipant from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup from tgbot_ping import get_runtime +from token_bucket import Limiter, MemoryStorage from youtubesearchpython import VideosSearch from client_init import create_app -from config import (AUTHORIZED_USER, ENABLE_CELERY, ENABLE_VIP, OWNER, - REQUIRED_MEMBERSHIP) +from config import (AUTHORIZED_USER, BURST, ENABLE_CELERY, ENABLE_FFMPEG, + ENABLE_VIP, OWNER, RATE, REQUIRED_MEMBERSHIP) from constant import BotText from db import InfluxDB, MySQL, Redis from limit import VIP, verify_payment @@ -44,6 +45,11 @@ bot_text = BotText() logging.info("Authorized users are %s", AUTHORIZED_USER) +# rate, capacity +mem = MemoryStorage() +# 5 minutes, 2 bursts +lim = Limiter(1 / RATE, BURST, mem) + def private_use(func): def wrapper(client: "Client", message: "types.Message"): @@ -145,7 +151,7 @@ def patch_handler(client: "Client", message: "types.Message"): @app.on_message(filters.command(["uncache"])) -def patch_handler(client: "Client", message: "types.Message"): +def uncache_handler(client: "Client", message: "types.Message"): username = message.from_user.username link = message.text.split()[1] if username == OWNER: @@ -169,7 +175,7 @@ def ping_handler(client: "Client", message: "types.Message"): @app.on_message(filters.command(["about"])) -def help_handler(client: "Client", message: "types.Message"): +def about_handler(client: "Client", message: "types.Message"): chat_id = message.chat.id client.send_chat_action(chat_id, "typing") client.send_message(chat_id, bot_text.about) @@ -296,6 +302,11 @@ def download_handler(client: "Client", message: "types.Message"): message.reply_text("Channel download is disabled now. Please send me individual video link.", quote=True) red.update_metrics("reject_channel") return + # non vip user, consume too many token + if (not VIP().check_vip(chat_id)) and (not lim.consume(str(chat_id).encode(), 1)): + red.update_metrics("rate_limit") + message.reply_text(bot_text.too_fast, quote=True) + return red.update_metrics("video_request") text = bot_text.get_receive_link_text() @@ -338,9 +349,13 @@ def download_resolution_callback(client: "Client", callback_query: types.Callbac @app.on_callback_query(filters.regex(r"convert")) def audio_callback(client: "Client", callback_query: types.CallbackQuery): + if not ENABLE_FFMPEG: + callback_query.answer("Audio conversion is disabled now.") + callback_query.message.reply_text("Audio conversion is disabled now.") + return + callback_query.answer(f"Converting to audio...please wait patiently") Redis().update_metrics("audio_request") - vmsg = callback_query.message audio_entrance(vmsg, client)