diff --git a/telegram_downloader_bot/main.py b/telegram_downloader_bot/main.py index 7eec879..3e5191d 100644 --- a/telegram_downloader_bot/main.py +++ b/telegram_downloader_bot/main.py @@ -40,6 +40,15 @@ app = Client("downloader_bot", bot_token=BOT_TOKEN) +def protected(func): + @wraps(func) + async def wrapper(client, message): + if message.from_user.id not in ALLOWED_IDS: + return await message.reply_text("You are not on the list!") + return await func(client, message) + return wrapper + + async def get_user_folder(message: Message) -> os.path: # Determine folder name based on whether the message was forwarded # and who it was forwarded from @@ -150,6 +159,7 @@ def download_tt_video(url: str) -> bool: @app.on_message(filters.command("start")) +@protected async def start_handler(_, message: Message): await message.reply_text( "This bot downloads TikTok videos to my personal server" @@ -157,11 +167,13 @@ async def start_handler(_, message: Message): @app.on_message(filters.command("help")) +@protected async def help_handler(_, message: Message): await message.reply_text("I won't help you!") @app.on_message(filters.text) +@protected async def message_handler(_, message: Message): urls = re.findall(r"\bhttps?://[^\s]+", message.text) @@ -183,6 +195,7 @@ async def message_handler(_, message: Message): @app.on_message(filters.media) +@protected async def media_handler(client, message: Message): await message.reply_text("Downloading media...")