From 5bd60da7c3d7a5e008d03ceb93aeec533119ee74 Mon Sep 17 00:00:00 2001 From: ch3p4ll3 Date: Wed, 20 Dec 2023 19:47:14 +0100 Subject: [PATCH] add toggle speed limit --- src/bot/custom_filters.py | 1 + .../settings/client_settings_callbacks.py | 44 +++++++++++++++++-- src/client_manager/client_manager.py | 10 +++++ src/client_manager/qbittorrent_manager.py | 11 +++++ 4 files changed, 63 insertions(+), 3 deletions(-) diff --git a/src/bot/custom_filters.py b/src/bot/custom_filters.py index 3077ead..3479187 100644 --- a/src/bot/custom_filters.py +++ b/src/bot/custom_filters.py @@ -50,6 +50,7 @@ check_connection_filter = filters.regex(r"^check_connection$") edit_client_setting_filter = filters.regex(r'^edit_clt(#.+|$)?$') reload_settings_filter = filters.regex(r"^reload_settings$") +toggle_speed_limit_filter = filters.regex(r"^toggle_speed_limit$") # Other diff --git a/src/bot/plugins/callbacks/settings/client_settings_callbacks.py b/src/bot/plugins/callbacks/settings/client_settings_callbacks.py index 515bc30..0c119b9 100644 --- a/src/bot/plugins/callbacks/settings/client_settings_callbacks.py +++ b/src/bot/plugins/callbacks/settings/client_settings_callbacks.py @@ -12,6 +12,41 @@ async def edit_client_settings_callback(client: Client, callback_query: CallbackQuery) -> None: confs = '\n- '.join(iter([f"**{key.capitalize()}:** {item}" for key, item in Configs.config.client.model_dump().items()])) + repository = ClientRepo.get_client_manager(Configs.config.client.type) + speed_limit = repository.get_speed_limit_mode() + + confs += f"\n\n**Speed Limit**: {'Enabled' if speed_limit else 'Disabled'}" + + await callback_query.edit_message_text( + f"Edit Qbittorrent Client Settings \n\n**Current Settings:**\n- {confs}", + reply_markup=InlineKeyboardMarkup( + [ + [ + InlineKeyboardButton("📝 Edit Client Settings", "lst_client") + ], + [ + InlineKeyboardButton("🐢 Toggle Speed Limit", "toggle_speed_limit") + ], + [ + InlineKeyboardButton("✅ Check Client connection", "check_connection") + ], + [ + InlineKeyboardButton("🔙 Settings", "settings") + ] + ] + ) + ) + + +@Client.on_callback_query(custom_filters.toggle_speed_limit_filter & custom_filters.check_user_filter & custom_filters.user_is_administrator) +async def toggle_speed_limit_callback(client: Client, callback_query: CallbackQuery) -> None: + confs = '\n- '.join(iter([f"**{key.capitalize()}:** {item}" for key, item in Configs.config.client.model_dump().items()])) + + repository = ClientRepo.get_client_manager(Configs.config.client.type) + speed_limit = repository.toggle_speed_limit() + + confs += f"\n\n**Speed Limit**: {'Enabled' if speed_limit else 'Disabled'}" + await callback_query.edit_message_text( f"Edit Qbittorrent Client Settings \n\n**Current Settings:**\n- {confs}", reply_markup=InlineKeyboardMarkup( @@ -19,6 +54,9 @@ async def edit_client_settings_callback(client: Client, callback_query: Callback [ InlineKeyboardButton("📝 Edit Client Settings", "lst_client") ], + [ + InlineKeyboardButton("🐢 Toggle Speed Limit", "toggle_speed_limit") + ], [ InlineKeyboardButton("✅ Check Client connection", "check_connection") ], @@ -30,7 +68,7 @@ async def edit_client_settings_callback(client: Client, callback_query: Callback ) -@Client.on_callback_query(custom_filters.check_connection_filter) +@Client.on_callback_query(custom_filters.check_connection_filter & custom_filters.check_user_filter & custom_filters.user_is_administrator) async def check_connection_callback(client: Client, callback_query: CallbackQuery) -> None: try: repository = ClientRepo.get_client_manager(Configs.config.client.type) @@ -41,7 +79,7 @@ async def check_connection_callback(client: Client, callback_query: CallbackQuer await callback_query.answer("❌ Unable to establish connection with QBittorrent", show_alert=True) -@Client.on_callback_query(custom_filters.list_client_settings_filter) +@Client.on_callback_query(custom_filters.list_client_settings_filter & custom_filters.check_user_filter & custom_filters.user_is_administrator) async def list_client_settings_callback(client: Client, callback_query: CallbackQuery) -> None: # get all fields of the model dynamically fields = [ @@ -63,7 +101,7 @@ async def list_client_settings_callback(client: Client, callback_query: Callback ) -@Client.on_callback_query(custom_filters.edit_client_setting_filter) +@Client.on_callback_query(custom_filters.edit_client_setting_filter & custom_filters.check_user_filter & custom_filters.user_is_administrator) async def edit_client_setting_callback(client: Client, callback_query: CallbackQuery) -> None: data = callback_query.data.split("#")[1] field_to_edit = data.split("-")[0] diff --git a/src/client_manager/client_manager.py b/src/client_manager/client_manager.py index ae61095..be5633d 100644 --- a/src/client_manager/client_manager.py +++ b/src/client_manager/client_manager.py @@ -88,3 +88,13 @@ def check_connection(cls) -> str: def export_torrent(cls, torrent_hash: str) -> BytesIO: """Export a .torrent file for the torrent.""" raise NotImplementedError + + @classmethod + def get_speed_limit_mode(cls) -> bool: + """Get speed limit of the client, returns True if speed limit is active""" + raise NotImplementedError + + @classmethod + def toggle_speed_limit(cls) -> bool: + """Toggle speed limit of the client, returns True if speed limit is active""" + raise NotImplementedError diff --git a/src/client_manager/qbittorrent_manager.py b/src/client_manager/qbittorrent_manager.py index a3670e7..fe26d45 100644 --- a/src/client_manager/qbittorrent_manager.py +++ b/src/client_manager/qbittorrent_manager.py @@ -155,3 +155,14 @@ def export_torrent(cls, torrent_hash: str) -> BytesIO: file_to_return = BytesIO(torrent_bytes) file_to_return.name = f"{torrent_name}.torrent" return file_to_return + + @classmethod + def get_speed_limit_mode(cls) -> bool: + with qbittorrentapi.Client(**Configs.config.client.connection_string) as qbt_client: + return qbt_client.transfer.speedLimitsMode == "1" + + @classmethod + def toggle_speed_limit(cls) -> bool: + with qbittorrentapi.Client(**Configs.config.client.connection_string) as qbt_client: + qbt_client.transfer.setSpeedLimitsMode() + return qbt_client.transfer.speedLimitsMode == "1"