Maintenance update

Author: Divkix

WebStreamer/__main__.py

@@ -19,6 +19,9 @@
 
 
 async def start_services():
+    """
+    Start the bot and the web server
+    """
     LOGGER.info("------------------- Initializing Telegram Bot -------------------")
     await StreamBot.start()
     LOGGER.info("----------------------------- DONE -----------------------------")

WebStreamer/bot/plugins/admin.py

@@ -21,6 +21,9 @@
     filters.command("status") & filters.private & filters.user(Vars.OWNER_ID),
 )
 async def status(_, m: Message):
+    """
+    Get status of the bot, number of users, number of files, etc.
+    """
     dl = Downloads()
     filename = "downloadList.txt"
     total_users = await Users().total_users_count()
@@ -59,6 +62,9 @@ async def status(_, m: Message):
     & filters.reply,
 )
 async def broadcast_(_, m: Message):
+    """
+    Broadcast a message to all users
+    """
     all_users = await Users().get_all_users()
     broadcast_msg = m.reply_to_message
     while 1:
@@ -79,7 +85,7 @@ async def broadcast_(_, m: Message):
     )
     async with open_aiofiles("broadcast.txt", "w") as broadcast_log_file:
         for user in all_users:
-            sts, msg = await send_msg(user_id=int(user), message=broadcast_msg)
+            sts, msg = await send_msg(user_id=int(user), m=broadcast_msg)
             if msg is not None:
                 await broadcast_log_file.write(msg)
             if sts == 200:

WebStreamer/bot/plugins/ban.py

@@ -7,6 +7,9 @@
 
 @StreamBot.on_callback_query(filters.regex("^ban_"))
 async def ban_user(c: StreamBot, q: CallbackQuery):
+    """
+    Ban a user from using the bot
+    """
     user_id = int(q.data.split("_", 1)[1])
     await c.ban_chat_member(Vars.AUTH_CHANNEL, user_id)
     await q.answer("User Banned from Updates Channel!", show_alert=True)

WebStreamer/bot/plugins/start.py

@@ -46,7 +46,7 @@
 
 class Btns:
     channel_and_group = [
-        ("Support Group", "https://t.me/DivideProjectsDiscussion", "url"),
+        ("Support Group", "https://t.me/DivideSupport", "url"),
         ("Channel", "https://t.me/DivideProjects", "url"),
     ]
     about_me = ("About Me", "aboutbot")
@@ -57,6 +57,9 @@ class Btns:
 @StreamBot.on_message(filters.command("start") & filters.private)
 @joinCheck()
 async def start(_, m: Message):
+    """
+    Start the bot
+    """
     return await m.reply_text(
         text=PMTEXT.format(m.from_user.mention),
         parse_mode=ParseMode.HTML,
@@ -68,6 +71,9 @@ async def start(_, m: Message):
 @StreamBot.on_message(filters.command("help") & filters.private)
 @joinCheck()
 async def help_handler(_, m: Message):
+    """
+    Help message handler
+    """
     return await m.reply_text(
         HELPTEXT,
         parse_mode=ParseMode.HTML,
@@ -77,28 +83,36 @@ async def help_handler(_, m: Message):
 
 @StreamBot.on_callback_query()
 async def button(_, m: CallbackQuery):
+    """
+    handle button presses
+    """
     cb_data = m.data
     msg = m.message
 
-    if cb_data == "aboutbot":
-        await msg.edit(
-            text=ABOUT,
-            parse_mode=ParseMode.HTML,
-            disable_web_page_preview=True,
-            reply_markup=ikb([[Btns.back]]),
-        )
-    elif cb_data == "helptext":
-        await msg.edit(
-            text=HELPTEXT,
-            parse_mode=ParseMode.HTML,
-            disable_web_page_preview=True,
-            reply_markup=ikb([[Btns.back]]),
-        )
-    elif cb_data == "gotohome":
-        await msg.edit(
-            text=PMTEXT.format(msg.from_user.mention),
-            parse_mode=ParseMode.HTML,
-            disable_web_page_preview=True,
-            reply_markup=ikb([Btns.channel_and_group, [Btns.about_me, Btns.help_me]]),
-        )
+    match cb_data:
+        case "aboutbot":
+            await msg.edit(
+                text=ABOUT,
+                parse_mode=ParseMode.HTML,
+                disable_web_page_preview=True,
+                reply_markup=ikb([[Btns.back]]),
+            )
+        case "helptext":
+            await msg.edit(
+                text=HELPTEXT,
+                parse_mode=ParseMode.HTML,
+                disable_web_page_preview=True,
+                reply_markup=ikb([[Btns.back]]),
+            )
+        case "gotohome":
+            await msg.edit(
+                text=PMTEXT.format(msg.from_user.mention),
+                parse_mode=ParseMode.HTML,
+                disable_web_page_preview=True,
+                reply_markup=ikb(
+                    [Btns.channel_and_group, [Btns.about_me, Btns.help_me]],
+                ),
+            )
+        case _:
+            await msg.edit("Invalid Button Pressed!")
     await m.answer()

WebStreamer/bot/plugins/stream.py

@@ -27,6 +27,7 @@
 <i>@DivideProjects </i>
 """
 
+# Cache for storing how many times a user has used the bot, takes number of mimuted from Vars
 ttl_dict = TTLCache(maxsize=512, ttl=(Vars.FLOODCONTROL_TIME_MINUTES * 60))
 
 
@@ -35,7 +36,7 @@
     & (filters.document | filters.video | filters.audio | filters.photo),
     group=4,
 )
-@joinCheck()
+@joinCheck()  # Check if user has joined the channel
 async def private_receive_handler(c: Client, m: Message):
     user = m.from_user
     user_id = user.id
@@ -122,6 +123,9 @@ async def private_receive_handler(c: Client, m: Message):
 
 @StreamBot.on_callback_query(filters.regex("^delete_url."))
 async def delete_download(_, q: CallbackQuery):
+    """
+    Delete the download link from the database using a callback query
+    """
     user_id = q.from_user.id
     msg = q.message
     url = str(q.data.split(".")[-1])

WebStreamer/db/__init__.py

@@ -5,6 +5,9 @@
 
 
 def __connect_first():
+    """
+    Connect to the database before importing the models
+    """
     _ = MongoDB("test")
     LOGGER.info("Initialized Database!")
 

WebStreamer/db/downloads.py

@@ -1,17 +1,28 @@
 from datetime import datetime, timedelta
 from secrets import token_urlsafe
+from typing import Tuple, Union
 
 from WebStreamer.db.mongo import MongoDB
 from WebStreamer.logger import LOGGER
 
 
 class Downloads(MongoDB):
+    """
+    Define downloads collection here
+    """
+
     db_name = "filestreamerbot_downloads"
 
     def __init__(self):
+        """
+        Initialize the collection
+        """
         super().__init__(self.db_name)
 
     async def add_download(self, message_id: int, random_url: str, user_id: int) -> str:
+        """
+        Add a download to the database
+        """
         LOGGER.info(f"Added {random_url}: {message_id}")
         real_link = token_urlsafe(16)
         await self.insert_one(
@@ -25,21 +36,30 @@ async def add_download(self, message_id: int, random_url: str, user_id: int) ->
         )
         return real_link
 
-    async def get_actual_link(self, link: str):
+    async def get_actual_link(self, link: str) -> Union[str, None]:
+        """
+        Get the actual link from the database
+        """
         document = await self.find_one({"random_link": link})
         if not document:
             return None
         return document["link"]
 
-    async def get_msg_id(self, link: str):
+    async def get_msg_id(self, link: str) -> Tuple[int, bool, datetime]:
+        """
+        Get the message id from the database
+        """
         document = await self.find_one({"link": link})
         if not document:
             return 0, False, datetime.now()
         valid_upto = document["valid_upto"]
         valid = valid_upto > datetime.now()
         return document["message_id"], valid, valid_upto
 
-    async def total_downloads(self):
+    async def total_downloads(self) -> int:
+        """
+        Get the total number of downloads
+        """
         return await self.count()
 
     async def valid_downloads_list(self):

WebStreamer/db/mongo.py

@@ -1,3 +1,5 @@
+from typing import Any, Tuple, Union
+
 from motor.motor_asyncio import AsyncIOMotorClient
 
 from WebStreamer.vars import Vars
@@ -7,55 +9,57 @@
 
 
 class MongoDB:
-    """Class for interacting with Bot database."""
+    """
+    Class for interacting with Bot database.
+    """
 
     def __init__(self, collection) -> None:
         self.collection = main_db[collection]
 
     # Insert one entry into collection
-    async def insert_one(self, document):
+    async def insert_one(self, document) -> str:
         result = await self.collection.insert_one(document)
         return repr(result.inserted_id)
 
     # Find one entry from collection
-    async def find_one(self, query):
+    async def find_one(self, query) -> Union[bool, None, Any]:
         result = await self.collection.find_one(query)
         if result:
             return result
         return False
 
     # Find entries from collection
-    async def find_all(self, query=None):
+    async def find_all(self, query=None) -> Union[bool, None, Any]:
         if query is None:
             query = {}
         return [document async for document in self.collection.find(query)]
 
     # Count entries from collection
-    async def count(self, query=None):
+    async def count(self, query=None) -> int:
         if query is None:
             query = {}
         return await self.collection.count_documents(query)
 
     # Delete entry/entries from collection
-    async def delete_one(self, query):
+    async def delete_one(self, query) -> int:
         await self.collection.delete_many(query)
         after_delete = await self.collection.count_documents({})
         return after_delete
 
     # Replace one entry in collection
-    async def replace(self, query, new_data):
+    async def replace(self, query, new_data) -> Tuple[int, int]:
         old = await self.collection.find_one(query)
         _id = old["_id"]
         await self.collection.replace_one({"_id": _id}, new_data)
         new = await self.collection.find_one({"_id": _id})
         return old, new
 
     # Update one entry from collection
-    async def update(self, query, update):
+    async def update(self, query, update) -> Tuple[int, int]:
         result = await self.collection.update_one(query, {"$set": update})
         new_document = await self.collection.find_one(query)
         return result.modified_count, new_document
 
     @staticmethod
-    async def db_command(command: str):
+    async def db_command(command: str) -> Any:
         return await main_db.command(command)

WebStreamer/db/users.py

@@ -1,27 +1,48 @@
 from datetime import date
+from typing import List, Union
 
 from WebStreamer.db.mongo import MongoDB
 from WebStreamer.logger import LOGGER
 
 
-def new_user(uid):
-    return {"id": uid, "join_date": date.today().isoformat(), "downloads": []}
+def new_user(uid: int):
+    """
+    Creates a new user in the database
+    """
+    return {
+        "id": uid,
+        "join_date": date.today().isoformat(),
+        "downloads": [],
+    }
 
 
 class Users(MongoDB):
+    """
+    Users collections to be made in the database
+    """
+
     db_name = "filestreamerbot_users"
 
     def __init__(self):
         super().__init__(self.db_name)
 
-    async def total_users_count(self):
+    async def total_users_count(self) -> int:
+        """
+        Returns the total number of users in the database
+        """
         return await self.count({})
 
-    async def get_all_users(self):
+    async def get_all_users(self) -> List[int]:
+        """
+        Returns a list of all users in the database
+        """
         users = await self.find_all({})
         return [user["id"] for user in users]
 
-    async def user_exists(self, user_id: int):
+    async def user_exists(self, user_id: int) -> bool:
+        """
+        Checks if a user exists in the database
+        """
         user = await self.find_one({"id": user_id})
         if not user:
             user_data = {
@@ -33,5 +54,8 @@ async def user_exists(self, user_id: int):
             return False
         return True
 
-    async def delete_user(self, user_id: int):
+    async def delete_user(self, user_id: int) -> Union[bool, int]:
+        """
+        Deletes a user from the database
+        """
         return await self.delete_one({"id": user_id})

WebStreamer/server/__init__.py

@@ -6,11 +6,16 @@
 
 
 async def web_server():
+    """
+    Create the web server and return it
+    """
     web_app = web.Application(client_max_size=30000000)
+    # setup jinja2 with the web templates from templates folder
     setup_jinja2(
         web_app,
         enable_async=True,
         loader=FileSystemLoader("/app/WebStreamer/html/templates"),
     )
+    # add the routes to the web app
     web_app.add_routes(routes)
     return web_app