Implement private mode

.env.sample

@@ -2,10 +2,11 @@ APP_ID=app-id-from-my.telegram.org
 API_HASH=api-hash-from-my.telegram.org
 BOT_TOKEN=bot-token-from-@BotFather
 
-USE_ENV=False
 MEGA_EMAIL=[email protected]
 MEGA_PASSWORD=mypassword
+AUTH_USERS=1340254734
 CIPHER_KEY=vJmDXO4xria6SYVcOfVYd3k3YM8WiiGBfRjbQ8MBsvI=
+USE_ENV=False
 
 DOWNLOAD_LOCATION=${PWD}/NexaBots
 CHUNK_SIZE=524288

README.md

@@ -60,7 +60,7 @@ python3 -m megadl
 Please refer to [documentation](https://megabot.hirusha.codes/config-vars)
 
 # Roadmap
-- [ ] Implement private mode
+- [x] Implement private mode
 - [x] Implement DDL to Mega.nz uploader
 - [ ] Better CLI output parser
 - [x] Heroku support

megadl/helpers/mclient.py

@@ -37,14 +37,20 @@ class MeganzClient(Client):
     """
 
     version = "v2-nightly"
-    database = Users() if os.getenv("MONGO_URI") else None
     dl_loc = None
     tmp_loc = None
+    database = Users() if os.getenv("MONGO_URI") else None
+    auth_users = (
+        set(map(int, os.getenv("AUTH_USERS").split()))
+        if os.getenv("AUTH_USERS")
+        else None
+    )
 
     def __init__(self):
         # set DOWNLOAD_LOCATION variable
         # if USE_ENV is True it'll use currend dir + NexaBots
         # otherwise use location defined in .env file by user
+        print("> Setting up file locations")
         self.cwd = os.getcwd()
         if os.getenv("USE_ENV") in ["True", "true"]:
             self.dl_loc = f"{self.cwd}/NexaBots"
@@ -54,6 +60,12 @@ def __init__(self):
         self.tmp_loc = f"{self.dl_loc}/temps"
         self.mx_size = int(os.getenv("TG_MAX_SIZE", 2040108421))
 
+        if not os.path.isdir(self.dl_loc):
+            os.makedirs(self.dl_loc)
+
+        if not os.path.isdir(self.tmp_loc):
+            os.makedirs(self.tmp_loc)
+
         # Initializing pyrogram
         print("> Initializing client")
         super().__init__(
@@ -91,19 +103,11 @@ def __init__(self):
         else:
             print("     Warning: Mongodb url not found")
 
-        # creating directories
-        print("> Creating directories")
-        if not os.path.isdir(self.dl_loc):
-            os.makedirs(self.dl_loc)
-
-        if not os.path.isdir(self.tmp_loc):
-            os.makedirs(self.tmp_loc)
-
         # other stuff
         print("> Setting up additional functions")
-        self.add_handler(MessageHandler(self.use_listner))
         self.listening = {}
         self.mega_running = {}
+        self.add_handler(MessageHandler(self.use_listner))
 
         print("--------------------")
 
@@ -114,10 +118,22 @@ def handle_checks(self, func: Callable) -> Callable:
         """
 
         async def fn_run(client: Client, msg: Message):
+            is_auth = False
+            uid = msg.from_user.id
             try:
-                # db functions
-                if self.database:
-                    await self.database.add(msg.from_user.id)
+                if self.auth_users and self.database:
+                    await self.database.add(uid)
+                    is_auth = uid in self.auth_users
+
+                elif self.database:
+                    await self.database.add(uid)
+
+                elif self.auth_users:
+                    is_auth = uid in self.auth_users
+
+                if not is_auth:
+                    await msg.reply("You're not authorized to use this bot 😬")
+                    return msg.stop_propagation()
 
                 return await func(client, msg)
             # Floodwait handling