From 912f1baeb30fec9386bd7a046a1a73ff7b45affb Mon Sep 17 00:00:00 2001
From: ch3p4ll3 <vidonimattia@gmail.com>
Date: Wed, 20 Dec 2023 19:25:18 +0100
Subject: [PATCH] check torrent add status

---
 src/bot/plugins/commands.py               |  2 ++
 src/bot/plugins/on_message.py             | 12 ++++++++++--
 src/client_manager/client_manager.py      |  8 ++++----
 src/client_manager/qbittorrent_manager.py | 11 +++++++----
 4 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/src/bot/plugins/commands.py b/src/bot/plugins/commands.py
index acae72c..13ace3c 100644
--- a/src/bot/plugins/commands.py
+++ b/src/bot/plugins/commands.py
@@ -3,6 +3,7 @@
 import psutil
 
 from .. import custom_filters
+from ...db_management import write_support
 from ...utils import convert_size
 from ...configs import Configs
 from .common import send_menu
@@ -18,6 +19,7 @@ async def access_denied_message(client: Client, message: Message) -> None:
 @Client.on_message(filters.command("start") & custom_filters.check_user_filter)
 async def start_command(client: Client, message: Message) -> None:
     """Start the bot."""
+    write_support("None", message.chat.id)
     await send_menu(client, message.id, message.chat.id)
 
 
diff --git a/src/bot/plugins/on_message.py b/src/bot/plugins/on_message.py
index 303a42e..37e045f 100644
--- a/src/bot/plugins/on_message.py
+++ b/src/bot/plugins/on_message.py
@@ -24,11 +24,15 @@ async def on_text(client: Client, message: Message) -> None:
             category = db_management.read_support(message.from_user.id).split("#")[1]
 
             repository = ClientRepo.get_client_manager(Configs.config.client.type)
-            repository.add_magnet(
+            response = repository.add_magnet(
                 magnet_link=magnet_link,
                 category=category
             )
 
+            if not response:
+                await message.reply_text("Unable to add magnet link")
+                return
+
             await send_menu(client, message.id, message.from_user.id)
             db_management.write_support("None", message.from_user.id)
 
@@ -43,7 +47,11 @@ async def on_text(client: Client, message: Message) -> None:
                 await message.download(name)
 
                 repository = ClientRepo.get_client_manager(Configs.config.client.type)
-                repository.add_torrent(file_name=name, category=category)
+                response = repository.add_torrent(file_name=name, category=category)
+
+                if not response:
+                    await message.reply_text("Unable to add magnet link")
+                    return
 
             await send_menu(client, message.id, message.from_user.id)
             db_management.write_support("None", message.from_user.id)
diff --git a/src/client_manager/client_manager.py b/src/client_manager/client_manager.py
index bc94577..ae61095 100644
--- a/src/client_manager/client_manager.py
+++ b/src/client_manager/client_manager.py
@@ -5,13 +5,13 @@
 
 class ClientManager(ABC):
     @classmethod
-    def add_magnet(cls, magnet_link: Union[str, List[str]], category: str = None) -> None:
-        """Add one or multiple magnet links with or without a category"""
+    def add_magnet(cls, magnet_link: Union[str, List[str]], category: str = None) -> bool:
+        """Add one or multiple magnet links with or without a category, return true if successful"""
         raise NotImplementedError
 
     @classmethod
-    def add_torrent(cls, file_name: str, category: str = None) -> None:
-        """Add one torrent file with or without a category"""
+    def add_torrent(cls, file_name: str, category: str = None) -> bool:
+        """Add one torrent file with or without a category, return true if successful"""
         raise NotImplementedError
 
     @classmethod
diff --git a/src/client_manager/qbittorrent_manager.py b/src/client_manager/qbittorrent_manager.py
index 295c43a..a3670e7 100644
--- a/src/client_manager/qbittorrent_manager.py
+++ b/src/client_manager/qbittorrent_manager.py
@@ -11,23 +11,26 @@
 
 class QbittorrentManager(ClientManager):
     @classmethod
-    def add_magnet(cls, magnet_link: Union[str, List[str]], category: str = None) -> None:
+    def add_magnet(cls, magnet_link: Union[str, List[str]], category: str = None) -> bool:
         if category == "None":
             category = None
 
         with qbittorrentapi.Client(**Configs.config.client.connection_string) as qbt_client:
             logger.debug(f"Adding magnet with category {category}")
-            qbt_client.torrents_add(urls=magnet_link, category=category)
+            result = qbt_client.torrents_add(urls=magnet_link, category=category)
+
+        return result == "Ok."
 
     @classmethod
-    def add_torrent(cls, file_name: str, category: str = None) -> None:
+    def add_torrent(cls, file_name: str, category: str = None) -> bool:
         if category == "None":
             category = None
 
         try:
             with qbittorrentapi.Client(**Configs.config.client.connection_string) as qbt_client:
                 logger.debug(f"Adding torrent with category {category}")
-                qbt_client.torrents_add(torrent_files=file_name, category=category)
+                result = qbt_client.torrents_add(torrent_files=file_name, category=category)
+            return result == "Ok."
 
         except qbittorrentapi.exceptions.UnsupportedMediaType415Error:
             pass