feat: control group or user permisions who can use a bot

Author: ovchynnikov

.env.example

@@ -1,2 +1,5 @@
 BOT_TOKEN="your_telegram_bot_token_here"
 DEBUG=False
+LIMIT_BOT_ACCESS=False # if True, bot will only work for users in ALLOWED_USERNAMES or ALLOWED_CHAT_IDS
+ALLOWED_USERNAMES= # list of allowed usernames as strings separated by commas
+ALLOWED_CHAT_IDS= # list of allowed chat ids as strings separated by commas

src/main.py

@@ -10,17 +10,9 @@
 from telegram.constants import MessageEntityType
 from logger import print_logs
 from video_utils import compress_video, download_video, cleanup_file
+from permissions import is_user_or_chat_not_allowed, supported_sites
 
 load_dotenv()
-supported_sites = [
-    "**https://",
-    "instagram.com/",
-    "tiktok.com/",
-    "reddit.com/",
-    "x.com/",
-    "youtube.com/shorts",
-]
-
 
 def load_responses():
     """Function loading bot responses."""
@@ -97,6 +89,11 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):  #
         await update.message.reply_text(random.choice(responses))
         return
 
+    # Check if user is not allowed
+    if is_user_or_chat_not_allowed(update.effective_user.username, update.effective_chat.id):
+        await update.message.reply_text("You are not allowed to use this bot")
+        return
+
     # Handle Instagram stories
     if "instagram.com/stories/" in message_text:
         await update.message.reply_text("Сторіз не можу скачати. Треба логін")
@@ -114,7 +111,7 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):  #
 
         # Download the video
         video_path = download_video(url)
-
+        # Check if video was downloaded
         if not video_path or not os.path.exists(video_path):
             return
 

src/permissions.py

@@ -0,0 +1,41 @@
+"""Module for handling user permissions and access control.
+This module provides functionality to check if users are allowed to access
+the bot based on their Telegram usernames and chat IDs configured in environment variables.
+"""
+
+import os
+from typing import Optional
+
+allowed_usernames = [x.strip() for x in os.getenv("ALLOWED_USERNAMES", "").split(",") if x]
+allowed_chat_ids = [int(x) for x in os.getenv("ALLOWED_CHAT_IDS", "").split(",") if x]
+limit_bot_access = os.getenv("LIMIT_BOT_ACCESS")
+
+# Check if user or chat is not allowed. Returns True if not allowed, False if allowed
+def is_user_or_chat_not_allowed(username: Optional[str], chat_id: int) -> bool:
+    """Check if username or chat_id is not in the allowed lists.
+
+    Args:
+        username: Telegram username to check
+
+    Returns:
+        True if neither user nor chat is allowed, False if either is allowed
+    """
+    # default case when no limits are set
+    if limit_bot_access == "False":
+        return False
+
+    # If chat_id is allowed, grant access regardless of username
+    if chat_id in allowed_chat_ids:
+        return False
+
+    # Otherwise check if username is allowed
+    return username not in allowed_usernames
+
+supported_sites = [
+    "**https://",
+    "instagram.com/",
+    "tiktok.com/",
+    "reddit.com/",
+    "x.com/",
+    "youtube.com/shorts",
+]