From aa0638414277ff94223c17ab1bab1299ada3d15f Mon Sep 17 00:00:00 2001
From: Igor Popov <lentrog@gmail.com>
Date: Fri, 28 Oct 2022 00:07:36 +0300
Subject: [PATCH] Add USER_ID environment variable

---
 README.md                | 13 +++++++++++++
 bot/handlers/messages.py | 13 +++++++++++--
 settings.py              |  3 ++-
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 65b68c5..a95a81c 100644
--- a/README.md
+++ b/README.md
@@ -18,6 +18,19 @@ Built on top of [aiogram](https://github.com/aiogram/aiogram)
 
 [![Deploy](https://www.herokucdn.com/deploy/button.svg)](https://heroku.com/deploy)
 
+## Env
+
+(*REQUIRED*)
+
+- API_TOKEN - Bot token from BotFather 
+
+(*OPTIONAL*)
+- USER_ID - To give access only to specific user id (default: empty = all users)
+- SENTRY_DSN - To send unhandled exceptions to Sentry (default: empty = no reports)
+- ENVIRONMENT - Sentry environment variable (default: Local)
+- USER_AGENT - To override user-agent used to download videos (default: random every time)
+
+
 ## Local
 
 ```bash
diff --git a/bot/handlers/messages.py b/bot/handlers/messages.py
index e21595e..5defb6d 100644
--- a/bot/handlers/messages.py
+++ b/bot/handlers/messages.py
@@ -1,6 +1,8 @@
+from aiogram.dispatcher.filters import IDFilter
 from aiogram.types import Message
-from bot import dp, bot
+from bot import bot, dp
 from bot.api import TikTokAPI
+from settings import USER_ID
 
 
 TikTok = TikTokAPI(
@@ -12,7 +14,14 @@
 )
 
 
-@dp.message_handler()
+def filter_message_handler(user_id: str):
+    if user_id:
+        return dp.message_handler(IDFilter(user_id=user_id))
+    else:
+        return dp.message_handler()
+
+
+@filter_message_handler(USER_ID)
 async def get_message(message: Message):
     async for video in TikTok.handle_message(message):
         if not video: continue
diff --git a/settings.py b/settings.py
index d80a88c..9ee0962 100644
--- a/settings.py
+++ b/settings.py
@@ -1,6 +1,7 @@
 import os
 
-API_TOKEN = os.environ.get('API_TOKEN')
+API_TOKEN = os.getenv('API_TOKEN')
+USER_ID = os.getenv('USER_ID')
 SENTRY_DSN = os.getenv('SENTRY_DSN')
 ENVIRONMENT = os.getenv('ENVIRONMENT', 'Local')
 USER_AGENT = os.getenv('USER_AGENT')