feat: allow only specific users (#93)

Author: EverythingSuckz

README.md

@@ -160,6 +160,10 @@ Before running the bot, you will need to set up the following mandatory variable
 ### Optional Vars
 In addition to the mandatory variables, you can also set the following optional variables:
 
+- `ALLOWED_USERS`: The user Telegram IDs of users to which the bot only reply to.
+> **Note**
+> Leave this field empty and anyone will be able to use your bot instance.
+> You may also add multiple users by adding the IDs separated by comma (,)
 
 - `HASH_LENGTH` : This is the custom hash length for generated URLs. The hash length must be greater than 5 and less than 64.
 

WebStreamer/bot/plugins/start.py

@@ -3,11 +3,14 @@
 
 from pyrogram import filters
 from pyrogram.types import Message
-from WebStreamer.bot import StreamBot
 
+from WebStreamer.vars import Var 
+from WebStreamer.bot import StreamBot
 
-@StreamBot.on_message(filters.command(["start", "help"]))
+@StreamBot.on_message(filters.command(["start", "help"]) & filters.private)
 async def start(_, m: Message):
+    if Var.ALLOWED_USERS and not ((str(m.from_user.id) in Var.ALLOWED_USERS) or (m.from_user.username in Var.ALLOWED_USERS)):
+        await m.reply("You are not in the allowed list of users who can use me. Check <a href='https://github.com/EverythingSuckz/TG-FileStreamBot#optional-vars'>this link</a> for more info.")
     await m.reply(
         f'Hi {m.from_user.mention(style="md")}, Send me a file to get an instant stream link.'
     )

WebStreamer/bot/plugins/stream.py

@@ -26,6 +26,8 @@
     group=4,
 )
 async def media_receive_handler(_, m: Message):
+    if Var.ALLOWED_USERS and not ((str(m.from_user.id) in Var.ALLOWED_USERS) or (m.from_user.username in Var.ALLOWED_USERS)):
+        return await m.reply("You are not <b>allowed to use</b> this <a href='https://github.com/EverythingSuckz/TG-FileStreamBot'>bot</a>.", quote=True)
     log_msg = await m.forward(chat_id=Var.BIN_CHANNEL)
     file_hash = get_hash(log_msg, Var.HASH_LENGTH)
     stream_link = f"{Var.URL}{log_msg.id}/{quote_plus(get_name(m))}?hash={file_hash}"

WebStreamer/vars.py

@@ -32,4 +32,5 @@ class Var(object):
         )
     KEEP_ALIVE = str(environ.get("KEEP_ALIVE", "0").lower()) in  ("1", "true", "t", "yes", "y")
     DEBUG = str(environ.get("DEBUG", "0").lower()) in ("1", "true", "t", "yes", "y")
-    USE_SESSION_FILE = str(environ.get("USE_SESSION_FILE", "0").lower()) in ("1", "true", "t", "yes", "y")
+    USE_SESSION_FILE = str(environ.get("USE_SESSION_FILE", "0").lower()) in ("1", "true", "t", "yes", "y")
+    ALLOWED_USERS = [x.strip("@ ") for x in str(environ.get("ALLOWED_USERS", "") or "").split(",") if x.strip("@ ")]