[REFACTOR] Enhancements and Code Cleanup

WebStreamer/__init__.py

@@ -3,8 +3,8 @@
 
 
 import time
-from .vars import Var
 from WebStreamer.bot.clients import StreamBot
+from .vars import Var
 
 __version__ = "2.2.4"
 StartTime = time.time()

WebStreamer/__main__.py

@@ -4,21 +4,22 @@
 import sys
 import asyncio
 import logging
-from .vars import Var
+from logging import handlers
 from aiohttp import web
 from pyrogram import idle
 from WebStreamer import utils
 from WebStreamer import StreamBot
 from WebStreamer.server import web_server
 from WebStreamer.bot.clients import initialize_clients
+from .vars import Var
 
 
 logging.basicConfig(
     level=logging.DEBUG if Var.DEBUG else logging.INFO,
     datefmt="%d/%m/%Y %H:%M:%S",
     format="[%(asctime)s][%(name)s][%(levelname)s] ==> %(message)s",
     handlers=[logging.StreamHandler(stream=sys.stdout),
-              logging.FileHandler("streambot.log", mode="a", encoding="utf-8")],)
+              handlers.RotatingFileHandler("streambot.log", mode="a", maxBytes=1048576*25, backupCount=2, encoding="utf-8")],)
 
 logging.getLogger("aiohttp").setLevel(logging.DEBUG if Var.DEBUG else logging.ERROR)
 logging.getLogger("pyrogram").setLevel(logging.INFO if Var.DEBUG else logging.ERROR)
@@ -43,10 +44,10 @@ async def start_services():
     await server.setup()
     await web.TCPSite(server, Var.BIND_ADDRESS, Var.PORT).start()
     logging.info("Service Started")
-    logging.info("bot =>> {}".format(bot_info.first_name))
+    logging.info("bot =>> %s", bot_info.first_name)
     if bot_info.dc_id:
-        logging.info("DC ID =>> {}".format(str(bot_info.dc_id)))
-    logging.info("URL =>> {}".format(Var.URL))
+        logging.info("DC ID =>> %d", bot_info.dc_id)
+    logging.info("URL =>> %s", Var.URL)
     await idle()
 
 async def cleanup():
@@ -63,4 +64,4 @@ async def cleanup():
     finally:
         loop.run_until_complete(cleanup())
         loop.stop()
-        logging.info("Stopped Services")
+        logging.info("Stopped Services")

WebStreamer/bot/__init__.py

@@ -4,16 +4,16 @@
 
 import os
 import os.path
-from ..vars import Var
 import logging
 from pyrogram import Client
+from ..vars import Var
 
 logger = logging.getLogger("bot")
 
 sessions_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sessions")
 if Var.USE_SESSION_FILE:
     logger.info("Using session files")
-    logger.info("Session folder path: {}".format(sessions_dir))
+    logger.info("Session folder path: %s", sessions_dir)
     if not os.path.isdir(sessions_dir):
         os.makedirs(sessions_dir)
 
@@ -29,5 +29,5 @@
     in_memory=not Var.USE_SESSION_FILE,
 )
 
-multi_clients = {}
-work_loads = {}
+multi_clients: dict[int, Client] = {}
+work_loads: dict[int, int] = {}

WebStreamer/bot/clients.py

@@ -4,8 +4,8 @@
 import asyncio
 import logging
 from os import environ
-from ..vars import Var
 from pyrogram import Client
+from ..vars import Var
 from . import multi_clients, work_loads, sessions_dir, StreamBot
 
 logger = logging.getLogger("multi_client")
@@ -24,10 +24,10 @@ async def initialize_clients():
     if not all_tokens:
         logger.info("No additional clients found, using default client")
         return
-    
+
     async def start_client(client_id, token):
         try:
-            logger.info(f"Starting - Client {client_id}")
+            logger.info("Starting - Client %s", client_id)
             if client_id == len(all_tokens):
                 await asyncio.sleep(2)
                 print("This will take some time, please wait...")
@@ -44,12 +44,12 @@ async def start_client(client_id, token):
             work_loads[client_id] = 0
             return client_id, client
         except Exception:
-            logger.error(f"Failed starting Client - {client_id} Error:", exc_info=True)
-    
+            logger.error("Failed starting Client - %s Error:", client_id, exc_info=True)
+
     clients = await asyncio.gather(*[start_client(i, token) for i, token in all_tokens.items()])
     multi_clients.update(dict(clients))
     if len(multi_clients) != 1:
         Var.MULTI_CLIENT = True
         logger.info("Multi-client mode enabled")
     else:
-        logger.info("No additional clients were initialized, using default client")
+        logger.info("No additional clients were initialized, using default client")

WebStreamer/bot/plugins/start.py

@@ -4,7 +4,7 @@
 from pyrogram import filters
 from pyrogram.types import Message
 
-from WebStreamer.vars import Var 
+from WebStreamer.vars import Var
 from WebStreamer.bot import StreamBot
 
 @StreamBot.on_message(filters.command(["start", "help"]) & filters.private)

WebStreamer/bot/plugins/stream.py

@@ -1,14 +1,13 @@
 # This file is a part of TG-FileStreamBot
 # Coding : Jyothis Jayanth [@EverythingSuckz]
 
-import logging
-from pyrogram import filters, errors
-from WebStreamer.vars import Var
 from urllib.parse import quote_plus
-from WebStreamer.bot import StreamBot, logger
-from WebStreamer.utils import get_hash, get_name
+from pyrogram import filters, errors
 from pyrogram.enums.parse_mode import ParseMode
 from pyrogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton
+from WebStreamer.vars import Var
+from WebStreamer.bot import StreamBot, logger
+from WebStreamer.utils import get_hash, get_name, get_mimetype
 
 
 @StreamBot.on_message(
@@ -30,25 +29,23 @@ async def media_receive_handler(_, m: Message):
         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)
+    mimetype = get_mimetype(log_msg)
     stream_link = f"{Var.URL}{log_msg.id}/{quote_plus(get_name(m))}?hash={file_hash}"
     short_link = f"{Var.URL}{file_hash}{log_msg.id}"
-    logger.info(f"Generated link: {stream_link} for {m.from_user.first_name}")
+    logger.info("Generated link: %s for %s", stream_link, m.from_user.first_name)
+    markup = [InlineKeyboardButton("Download", url=stream_link+"&d=true")]
+    if set(mimetype.split("/")) & {"video","audio","pdf"}:
+        markup.append(InlineKeyboardButton("Stream", url=stream_link))
     try:
         await m.reply_text(
-            text="<code>{}</code>\n(<a href='{}'>shortened</a>)".format(
-                stream_link, short_link
-            ),
+            text=f"<code>{stream_link}</code>\n(<a href='{short_link}'>shortened</a>)",
             quote=True,
             parse_mode=ParseMode.HTML,
-            reply_markup=InlineKeyboardMarkup(
-                [[InlineKeyboardButton("Open", url=stream_link)]]
-            ),
+            reply_markup=InlineKeyboardMarkup([markup]),
         )
     except errors.ButtonUrlInvalid:
         await m.reply_text(
-            text="<code>{}</code>\n\nshortened: {})".format(
-                stream_link, short_link
-            ),
+            text=f"<code>{stream_link}</code>\n(<a href='{short_link}'>shortened</a>)",
             quote=True,
             parse_mode=ParseMode.HTML,
         )

WebStreamer/server/exceptions.py

@@ -3,4 +3,4 @@ class InvalidHash(Exception):
     message = "Invalid hash"
 
 class FIleNotFound(Exception):
-    message = "File not found"
+    message = "File not found"

WebStreamer/server/stream_routes.py

@@ -5,7 +5,6 @@
 import time
 import math
 import logging
-import secrets
 import mimetypes
 from aiohttp import web
 from aiohttp.http_exceptions import BadStatusLine
@@ -62,30 +61,31 @@ async def stream_handler(request: web.Request):
 class_cache = {}
 
 async def media_streamer(request: web.Request, message_id: int, secure_hash: str):
+    head: bool = request.method == "HEAD"
     range_header = request.headers.get("Range", 0)
-    
+
     index = min(work_loads, key=work_loads.get)
     faster_client = multi_clients[index]
-    
+
     if Var.MULTI_CLIENT:
-        logger.info(f"Client {index} is now serving {request.remote}")
+        logger.info("Client %d is now serving %s", index, request.remote)
 
     if faster_client in class_cache:
         tg_connect = class_cache[faster_client]
-        logger.debug(f"Using cached ByteStreamer object for client {index}")
+        logger.debug("Using cached ByteStreamer object for client %d", index)
     else:
-        logger.debug(f"Creating new ByteStreamer object for client {index}")
+        logger.debug("Creating new ByteStreamer object for client %d", index)
         tg_connect = utils.ByteStreamer(faster_client)
         class_cache[faster_client] = tg_connect
     logger.debug("before calling get_file_properties")
     file_id = await tg_connect.get_file_properties(message_id)
     logger.debug("after calling get_file_properties")
-    
-    
+
+
     if utils.get_hash(file_id.unique_id, Var.HASH_LENGTH) != secure_hash:
-        logger.debug(f"Invalid hash for message with ID {message_id}")
+        logger.debug("Invalid hash for message with ID %d", message_id)
         raise InvalidHash
-    
+
     file_size = file_id.file_size
 
     if range_header:
@@ -112,18 +112,21 @@ async def media_streamer(request: web.Request, message_id: int, secure_hash: str
 
     req_length = until_bytes - from_bytes + 1
     part_count = math.ceil(until_bytes / chunk_size) - math.floor(offset / chunk_size)
-    body = tg_connect.yield_file(
-        file_id, index, offset, first_part_cut, last_part_cut, part_count, chunk_size
-    )
+    if head:
+        body=None
+    else:
+        body = tg_connect.yield_file(
+            file_id, index, offset, first_part_cut, last_part_cut, part_count, chunk_size
+        )
     mime_type = file_id.mime_type
     file_name = utils.get_name(file_id)
-    disposition = "attachment"
+    disposition = "inline"
 
     if not mime_type:
         mime_type = mimetypes.guess_type(file_name)[0] or "application/octet-stream"
 
-    if "video/" in mime_type or "audio/" in mime_type or "/html" in mime_type:
-        disposition = "inline"
+    if request.rel_url.query.get("d") == "true":
+        disposition = "attachment"
 
     return web.Response(
         status=206 if range_header else 200,

WebStreamer/utils/__init__.py

@@ -3,5 +3,5 @@
 
 from .keepalive import ping_server
 from .time_format import get_readable_time
-from .file_properties import get_hash, get_name
-from .custom_dl import ByteStreamer
+from .file_properties import get_hash, get_name, get_mimetype
+from .custom_dl import ByteStreamer