add token bucket

Author: BennyThink

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

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

ytdlbot/config.py

@@ -19,7 +19,7 @@
 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 @@
 
 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))

ytdlbot/constant.py

@@ -10,8 +10,8 @@
 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 @@ class BotText:
 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 @@ def ping_worker():
             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."

ytdlbot/utils.py

@@ -205,7 +205,7 @@ def next_salt_detector(self):
     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
 

ytdlbot/ytdl_bot.py

@@ -22,11 +22,12 @@
 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 @@
 
 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)