diff --git a/public/screenshot.png b/public/screenshot.png index d2dd565..8ba5632 100644 Binary files a/public/screenshot.png and b/public/screenshot.png differ diff --git a/src/bot.py b/src/bot.py index bb580ac..e9c0797 100644 --- a/src/bot.py +++ b/src/bot.py @@ -1,25 +1,15 @@ # Environment variables from dotenv import load_dotenv -load_dotenv() +load_dotenv(override=True) import logging import os -import platform -import shutil -import time -from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update -from telegram.ext import ( - Application, - CallbackQueryHandler, - CommandHandler, - ContextTypes, - MessageHandler, - filters, -) +from telegram import Update +from telegram.ext import Application, CommandHandler, ContextTypes -from .auth import auth_required +from .cogs import downloader_commands, error_handler, general_commands # Environment variables BOT_TOKEN = os.getenv("BOT_TOKEN") @@ -30,9 +20,6 @@ if not all([BOT_TOKEN, LOCAL_BOT_API_URL, BOT_API_DIR, DOWNLOAD_TO_DIR]): raise ValueError("Please set all environment variables in .env file") -# Replacing colons with a different character for Windows -TOKEN_SUB_DIR = BOT_TOKEN.replace(":", "") if os.name == "nt" else BOT_TOKEN - # Enable logging logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO @@ -40,138 +27,10 @@ logging.getLogger("httpx").setLevel(logging.WARNING) logger = logging.getLogger(__name__) -# List of available commands -commands = { - "/start": "Start the bot", - "/help": "Get help", - "/info": "Get user and chat info", -} - - -async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - """Send a list of available commands to the user.""" - commands_list = "The following commands are available:\n" + "\n".join( - [f"{key} - {value}" for key, value in commands.items()] - ) - await update.message.reply_text( - f"{commands_list}\n\nSend me a file and I'll download it to `{DOWNLOAD_TO_DIR}`.", - parse_mode="markdown", - ) - - -async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - """Send a start message to the user.""" - user = update.effective_user - await update.message.reply_html( - f"Hi {user.mention_html()}! I'm a bot that can download files for you. " - "Send me a file and I'll download it for you.\n\n" - "Use /help to see available commands." - ) - - -async def info(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - """Send user and chat IDs to the user.""" - user = update.effective_user - await update.message.reply_text( - f"*User ID*: {user.id}\n*Chat ID*: {update.effective_chat.id}", - parse_mode="markdown", - ) - - -@auth_required -async def download(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - """Download the file sent by the user.""" - logger.info("Download command received") - - # Check if file exists in DOWNLOAD_TO_DIR already - if os.path.exists(DOWNLOAD_TO_DIR + update.message.document.file_name): - await update.message.reply_text("File already exists in downloads folder.") - return - - # File details - file_id = update.message.document.file_id - file_name = update.message.document.file_name - file_size = update.message.document.file_size / 1024 / 1024 # in MB - - # Confirmation message - ok = await update.message.reply_text( - f"Download file?\nFile name: {file_name}\nFile size: {file_size:.2f} MB", - reply_markup=InlineKeyboardMarkup( - [ - [ - InlineKeyboardButton("Yes", callback_data="yes"), - InlineKeyboardButton("No", callback_data="no"), - ] - ] - ), - ) - - -@auth_required -async def button(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - """Handle the confirmation button click for downloading the file.""" - logger.info("Button command received") - query = update.callback_query - - await query.answer() - - # Replied to message - message = update.effective_message.reply_to_message - file_name = message.document.file_name - file_size = message.document.file_size / 1024 / 1024 # in MB - - if query.data == "yes": - logger.info("Downloading file...") - - # Check if file exists in DOWNLOAD_TO_DIR already - if os.path.exists(DOWNLOAD_TO_DIR + file_name): - await message.reply_text("File already exists in downloads folder.") - return - - await message.reply_text( - f"Downloading file...\nFile name: {file_name}\nFile size: {file_size:.2f} MB" - ) - start_time = time.time() - - try: - newFile = await context.bot.get_file( - message.document.file_id, read_timeout=1000 - ) - except Exception as e: - await message.reply_text(f"File size: {file_size}\nError: {e}") - return - - # Work out time taken to download file - end_time = time.time() - download_time = end_time - start_time - download_time_mins = download_time / 60 - file_path = newFile.file_path.split("/")[-1] - downloaded_file_size = newFile.file_size / 1024 / 1024 - - # Rename the file to the original file name - current_file_path = f"{BOT_API_DIR}{TOKEN_SUB_DIR}/documents/{file_path}" - move_to_path = f"{DOWNLOAD_TO_DIR}{file_name}" - - # Make DOWNLOAD_TO_DIR if it doesn't exist - os.makedirs(DOWNLOAD_TO_DIR, exist_ok=True) - shutil.move(current_file_path, move_to_path) - - # If linux, give file correct permissions - if platform.system() == "Linux": - os.chmod(move_to_path, 0o664) - - response_message = ( - f"File downloaded successfully.\n" - f"Time taken: {download_time:.5f} secs ({download_time_mins:.2f} mins)\n" - f"File path: {file_path}\n" - f"File name: {file_name}\n" - f"File size: {downloaded_file_size:.2f} MB" - ) - await message.reply_text(response_message) - else: - logger.info("Download cancelled") - await message.reply_text("Download cancelled.") +async def bad_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """Raise an error to trigger the error handler.""" + await context.bot.wrong_method_name() # type: ignore[attr-defined] def main() -> None: @@ -179,6 +38,7 @@ def main() -> None: application = ( Application.builder() .token(BOT_TOKEN) + .concurrent_updates(True) .local_mode(True) .base_url(f"{LOCAL_BOT_API_URL}/bot") .base_file_url(f"{LOCAL_BOT_API_URL}/file/bot") @@ -186,13 +46,12 @@ def main() -> None: ) # on different commands - answer in Telegram - application.add_handler(CommandHandler("start", start)) - application.add_handler(CommandHandler("help", help_command)) - application.add_handler(CommandHandler("info", info)) + application.add_handler(CommandHandler("bad_command", bad_command)) + application.add_handlers(general_commands) + application.add_handlers(downloader_commands) - # on file upload - download the file - application.add_handler(MessageHandler(filters.Document.VIDEO, download)) - application.add_handler(CallbackQueryHandler(button)) + # error handler + application.add_error_handler(error_handler) # Run the bot until the user presses Ctrl-C application.run_polling(allowed_updates=Update.ALL_TYPES) diff --git a/src/cogs/__init__.py b/src/cogs/__init__.py new file mode 100644 index 0000000..e66e611 --- /dev/null +++ b/src/cogs/__init__.py @@ -0,0 +1,17 @@ +from .downloader import button, download, status +from .error_handler import error_handler +from .general import help_command, info, start, storage + +# Specify the commands for the bot +general_commands: list = [ + help_command, + info, + start, + storage +] + +downloader_commands: list = [ + button, + download, + status, +] diff --git a/src/cogs/downloader.py b/src/cogs/downloader.py new file mode 100644 index 0000000..184c906 --- /dev/null +++ b/src/cogs/downloader.py @@ -0,0 +1,208 @@ +import logging +import os +import platform +import shutil +import time + +from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update +from telegram.ext import ContextTypes, filters + +from ..middlewares.auth import auth_required +from ..middlewares.handlers import ( + callback_query_handler, + command_handler, + message_handler, +) +from ..models import DownloadingFile + +logger = logging.getLogger(__name__) + +# Environment variables +BOT_TOKEN = os.getenv("BOT_TOKEN") +BOT_API_DIR = os.getenv("BOT_API_DIR") +DOWNLOAD_TO_DIR = os.getenv("DOWNLOAD_TO_DIR") + +# Replacing colons with a different character for Windows +TOKEN_SUB_DIR = BOT_TOKEN.replace(":", "") if os.name == "nt" else BOT_TOKEN + +# Current downloading files +downloading_files: dict[str, DownloadingFile] = {} + + +def check_file_exists(file_id: str, file_name: str) -> tuple[bool, str]: + """ + Check if a file exists in the download directory or is currently being downloaded. + + Args: + file_id (str): The ID of the file to check. + file_name (str): The name of the file to check. + + Returns: + tuple[bool, str]: A tuple containing a boolean value and a message. + """ + if os.path.exists(DOWNLOAD_TO_DIR + file_name): + return True, "File already exists in downloads folder." + + if file_id in downloading_files: + return True, "File is already being downloaded." + + # Check file_name in downloading_files + if any(file.file_name == file_name for file in downloading_files.values()): + return True, "File is already being downloaded." + + return False + + +@command_handler("status") +async def status(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """Send downloading files status to the user.""" + if not downloading_files: + await update.message.reply_text("No files are being downloaded at the moment.") + return + + status_message = "*Downloading files status:*\n\n" + for file in downloading_files.values(): + status_message += ( + f"> 📄 *File name:* `{file.file_name}`\n" + f"> 💾 *File size:* `{file.file_size_mb}`\n" + f"> ⏰ *Start time:* `{file.start_datetime}`\n" + f"> ⏱ *Duration:* `{file.download_time}`\n\n" + ) + + await update.message.reply_text(status_message, parse_mode="MarkdownV2") + + +@message_handler(filters.Document.VIDEO) +@auth_required +async def download(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """Download the file sent by the user.""" + logger.info("Download command received") + + # Check if file already exists or is being downloaded + if check := check_file_exists( + update.message.document.file_id, update.message.document.file_name + ): + await update.message.reply_text(check[1]) + return + + # File details + file_id = update.message.document.file_id + file_name = update.message.document.file_name + file_size = DownloadingFile.convert_size(update.message.document.file_size) + + response_message = ( + f"Are you sure you want to download the file?\n\n" + f"> 📄 *File name:* `{file_name}`\n" + f"> 💾 *File size:* `{file_size}`\n" + ) + + # Confirmation message + await context.bot.send_message( + chat_id=update.message.chat_id, + text=response_message, + reply_to_message_id=update.message.message_id, + parse_mode="MarkdownV2", + reply_markup=InlineKeyboardMarkup( + [ + [ + InlineKeyboardButton("Yes", callback_data="yes"), + InlineKeyboardButton("No", callback_data="no"), + ] + ] + ), + ) + + +@callback_query_handler() +@auth_required +async def button(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """Handle the confirmation button click for downloading the file.""" + logger.info("Button command received") + query = update.callback_query + + await query.answer() + + # Replied to message + message = update.effective_message.reply_to_message + file_name = message.document.file_name + + if query.data == "yes": + logger.info("Downloading file...") + + # Check if file already exists or is being downloaded + if check := check_file_exists(message.document.file_id, file_name): + await message.reply_text(check[1]) + return + + start_time = time.time() + + # Add file to downloading_files + downloading_file = DownloadingFile( + file_name=file_name, + file_size=message.document.file_size, + start_time=start_time, + ) + downloading_files[message.document.file_id] = downloading_file + + # Send downloading message + await message.reply_text("⬇️ Downloading file...") + + try: + newFile = await context.bot.get_file( + message.document.file_id, read_timeout=1000 + ) + except Exception as e: + await message.reply_text( + ( + f"⛔ Error downloading file.\n\n" + f"> 📄 *File name:* `{downloading_file.file_name}`\n" + f"> 💾 *File size:* `{downloading_file.file_size_mb}`\n" + f"Error: ```{e}```" + ), + parse_mode="MarkdownV2", + ) + return + + # Remove file from downloading_files + downloading_files.pop(message.document.file_id) + + # Work out time taken to download file + download_complete_time = time.time() + dowload_duration = DownloadingFile.convert_duration( + download_complete_time - start_time + ) + file_path = newFile.file_path.split("/")[-1] + + # Rename the file to the original file name + current_file_path = f"{BOT_API_DIR}{TOKEN_SUB_DIR}/documents/{file_path}" + move_to_path = f"{DOWNLOAD_TO_DIR}{file_name}" + + # Make DOWNLOAD_TO_DIR if it doesn't exist + os.makedirs(DOWNLOAD_TO_DIR, exist_ok=True) + shutil.move(current_file_path, move_to_path) + + # If linux, give file correct permissions + if platform.system() == "Linux": + os.chmod(move_to_path, 0o664) + + # Calculate durations + complete_time = time.time() + moving_duration = DownloadingFile.convert_duration( + complete_time - download_complete_time + ) + total_duration = DownloadingFile.convert_duration(complete_time - start_time) + + response_message = ( + f"✅ File downloaded successfully\\.\n\n" + f"> 📄 *File name:* `{downloading_file.file_name}`\n" + f"> 📂 *File path:* `{file_path}`\n" + f"> 💾 *File size:* `{downloading_file.file_size_mb}`\n" + f"> ⏱ *Download Duration:* `{dowload_duration}`\n" + f"> ⏱ *Moving Duration:* `{moving_duration}`\n" + f"> ⏱ *Total Duration:* `{total_duration}`\n" + ) + + await message.reply_text(response_message, parse_mode="MarkdownV2") + else: + logger.info("Download cancelled") + await message.reply_text("Download cancelled.") diff --git a/src/cogs/error_handler.py b/src/cogs/error_handler.py new file mode 100644 index 0000000..ab236bd --- /dev/null +++ b/src/cogs/error_handler.py @@ -0,0 +1,41 @@ +import html +import logging +import os +import traceback + +from telegram import Update +from telegram.ext import ContextTypes + +logger = logging.getLogger(__name__) + +USER_ID = os.getenv("USER_ID") + + +async def error_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """Log the error and send a telegram message to notify the developer.""" + logger.error("Exception while handling an update:", exc_info=context.error) + + # Join the traceback error + tb_list = traceback.format_exception( + None, context.error, context.error.__traceback__ + ) + tb_string = "".join(tb_list) + + # Build the error message + update_str = update.to_dict() if isinstance(update, Update) else str(update) + message = ( + "An exception was raised while handling an update\n" + # f"
update = {html.escape(json.dumps(update_str, indent=2, ensure_ascii=False))}"
+        # "
\n\n" + f"
context.chat_data = {html.escape(str(context.chat_data))}
\n\n" + f"
context.user_data = {html.escape(str(context.user_data))}
\n\n" + f"
{html.escape(tb_string)}
" + ) + + # Send the error message to the developer + await context.bot.send_message(chat_id=USER_ID, text=message, parse_mode="HTML") + + # Send error message in chat + await update.message.reply_text( + "An error occurred while processing the request. Please check the logs." + ) diff --git a/src/cogs/general.py b/src/cogs/general.py new file mode 100644 index 0000000..14c6937 --- /dev/null +++ b/src/cogs/general.py @@ -0,0 +1,70 @@ +import logging +import os +import shutil + +from telegram import Update +from telegram.ext import ContextTypes + +from ..middlewares.handlers import command_handler + +logger = logging.getLogger(__name__) + +DOWNLOAD_TO_DIR = os.getenv("DOWNLOAD_TO_DIR") + +# List of available commands +commands = { + "/start": "Start the bot", + "/help": "Get help", + "/info": "Get user and chat info", + "/storage": "Get available storage information", + "/status": "Get downloading files status", +} + + +@command_handler("help") +async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """Send a list of available commands to the user.""" + commands_list = "The following commands are available:\n" + "\n".join( + [f"{key} - {value}" for key, value in commands.items()] + ) + await update.message.reply_text( + f"{commands_list}\n\nSend me a file and I'll download it to `{DOWNLOAD_TO_DIR}`.", + parse_mode="markdown", + ) + + +@command_handler("start") +async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """Send a start message to the user.""" + user = update.effective_user + await update.message.reply_html( + f"Hi {user.mention_html()}! I'm a bot that can download files for you. " + "Send me a file and I'll download it for you.\n\n" + "Use /help to see available commands." + ) + + +@command_handler("info") +async def info(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """Send user and chat IDs to the user.""" + user = update.effective_user + await update.message.reply_text( + f"*User ID*: {user.id}\n*Chat ID*: {update.effective_chat.id}", + parse_mode="markdown", + ) + + +@command_handler("storage") +async def storage(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """Send available storage information of the specified folder.""" + if os.path.exists(DOWNLOAD_TO_DIR): + total, used, free = shutil.disk_usage(DOWNLOAD_TO_DIR) + await update.message.reply_text( + f"📂 *Folder*: `{DOWNLOAD_TO_DIR}`\n" + f"🟣 *Total Space*: `{total // (2**30)} GB`\n" + f"🟠 *Used Space*: `{used // (2**30)} GB`\n" + f"🟢 *Free Space*: `{free // (2**30)} GB`", + parse_mode="markdown", + ) + else: + await update.message.reply_text("The specified folder does not exist.") diff --git a/src/middlewares/__init__.py b/src/middlewares/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/auth.py b/src/middlewares/auth.py similarity index 100% rename from src/auth.py rename to src/middlewares/auth.py diff --git a/src/middlewares/handlers.py b/src/middlewares/handlers.py new file mode 100644 index 0000000..2439144 --- /dev/null +++ b/src/middlewares/handlers.py @@ -0,0 +1,48 @@ +from typing import Any, Callable, Coroutine, Dict + +from telegram import Update +from telegram.ext import ( + CallbackContext, + CallbackQueryHandler, + CommandHandler, + ExtBot, + MessageHandler, + filters, +) +from telegram.ext.filters import BaseFilter + +ApplicationContext = CallbackContext[ + ExtBot[None], Dict[Any, Any], Dict[Any, Any], Dict[Any, Any] +] + + +# https://github.com/Lur1an/python-telegram-bot-template/blob/master/src/bot/common/context.py +def command_handler(command: str | list[str], *, filters: BaseFilter = filters.ALL): + def inner_decorator( + f: Callable[[Update, ApplicationContext], Coroutine[Any, Any, Any]] + ) -> CommandHandler: + return CommandHandler( + filters=filters, + command=command, + callback=f, + ) + + return inner_decorator + + +def message_handler(filters: BaseFilter): + def inner_decorator( + f: Callable[[Update, ApplicationContext], Coroutine[Any, Any, Any]] + ) -> MessageHandler: + return MessageHandler(filters=filters, callback=f) + + return inner_decorator + + +def callback_query_handler(): + def inner_decorator( + f: Callable[[Update, ApplicationContext], Coroutine[Any, Any, Any]] + ) -> CallbackQueryHandler: + return CallbackQueryHandler(callback=f) + + return inner_decorator diff --git a/src/models/__init__.py b/src/models/__init__.py new file mode 100644 index 0000000..0b89467 --- /dev/null +++ b/src/models/__init__.py @@ -0,0 +1 @@ +from .downloading_file import DownloadingFile diff --git a/src/models/downloading_file.py b/src/models/downloading_file.py new file mode 100644 index 0000000..391032b --- /dev/null +++ b/src/models/downloading_file.py @@ -0,0 +1,35 @@ +import time +from dataclasses import InitVar, dataclass +from datetime import datetime + + +@dataclass +class DownloadingFile: + file_name: str + file_size: int + start_time: float + _start_datetime: InitVar[datetime] = None + + def __post_init__(self, _start_datetime): + self._start_datetime = _start_datetime or datetime.now() + + @property + def start_datetime(self) -> str: + return self._start_datetime.strftime("%H:%M:%S %d/%m/%Y") + + @property + def file_size_mb(self) -> str: + return f"{self.file_size / 1024 / 1024:.2f} MB" + + @property + def download_time(self) -> str: + time_taken = time.time() - self.start_time + return self.convert_duration(time_taken) + + @staticmethod + def convert_duration(time_taken: float) -> str: + return f"{time_taken:.2f} secs ({time_taken / 60:.2f} mins)" + + @staticmethod + def convert_size(size: int) -> str: + return f"{size / 1024 / 1024:.2f} MB" diff --git a/src/utils/__init__.py b/src/utils/__init__.py new file mode 100644 index 0000000..e69de29