From 3f3831bdba51af45d8f585f20fb71622059ffab0 Mon Sep 17 00:00:00 2001
From: bakatrouble <bakatrouble@gmail.com>
Date: Fri, 14 Apr 2023 11:24:20 +0300
Subject: [PATCH] protect bot

---
 main.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/main.py b/main.py
index f6e2231..5f4a64a 100644
--- a/main.py
+++ b/main.py
@@ -7,6 +7,7 @@ from PIL import Image
 import httpx
 import redis.asyncio as aioredis
 from aiogram import Bot, Dispatcher
+from aiogram.dispatcher import filters
 from aiogram.types import Message, ParseMode
 from aiogram.utils import executor, exceptions
 import dotenv
@@ -22,6 +23,8 @@ logging.basicConfig(level=logging.INFO)
 bot = Bot(token=os.environ['BOT_TOKEN'])
 dp = Dispatcher(bot)
 
+user_ids = os.environ['USERS'].split(',')
+
 
 async def check_updates():
     logging.warning('Waiting for lock...')
@@ -99,7 +102,7 @@ async def check_updates():
                         logging.exception(e)
 
 
-@dp.message_handler(commands=['add'])
+@dp.message_handler(commands=['add'], filter=filters.IDFilter(chat_id=user_ids))
 async def add_tag(msg: Message):
     args = msg.get_args()
     if not args:
@@ -110,7 +113,7 @@ async def add_tag(msg: Message):
     await msg.reply(f'Tags {args} added')
 
 
-@dp.message_handler(regexp=r'^\/del_\S+$')
+@dp.message_handler(regexp=r'^\/del_\S+$', filter=filters.IDFilter(chat_id=user_ids))
 async def del_tag(msg: Message):
     args = msg.text[5:]
     if not args:
@@ -126,7 +129,7 @@ async def del_tag(msg: Message):
     await msg.reply(f'Tag {args} removed')
 
 
-@dp.message_handler(commands=['del'])
+@dp.message_handler(commands=['del'], filter=filters.IDFilter(chat_id=user_ids))
 async def del_command(msg: Message):
     args = msg.get_args()
     if not args:
@@ -137,7 +140,7 @@ async def del_command(msg: Message):
     await msg.reply(f'Tags {args} removed')
 
 
-@dp.message_handler(commands=['list'])
+@dp.message_handler(commands=['list'], filter=filters.IDFilter(chat_id=user_ids))
 async def list_tags(msg: Message):
     tags = [t.decode() for t in await redis.smembers('e621:subs')]
     tags.sort()
@@ -155,7 +158,7 @@ async def list_tags(msg: Message):
     await msg.reply(f'Monitored tags:\n\n{lines}')
 
 
-@dp.message_handler(commands=['update'])
+@dp.message_handler(commands=['update'], filter=filters.IDFilter(chat_id=user_ids))
 async def update(msg: Message):
     await check_updates()