From 505bfb4c2ec3ebcc13334c06ca3a350b3f1ba775 Mon Sep 17 00:00:00 2001
From: jxxghp <jxxghp@gmail.com>
Date: Mon, 7 Aug 2023 11:01:55 +0800
Subject: [PATCH] =?UTF-8?q?feat=20=E4=B8=8B=E8=BD=BD=E5=99=A8=E9=99=90?=
 =?UTF-8?q?=E9=80=9F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/modules/qbittorrent/qbittorrent.py   | 19 ++++++++++++++++
 app/modules/transmission/transmission.py | 29 ++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/app/modules/qbittorrent/qbittorrent.py b/app/modules/qbittorrent/qbittorrent.py
index d341f38e..6cdd5cd7 100644
--- a/app/modules/qbittorrent/qbittorrent.py
+++ b/app/modules/qbittorrent/qbittorrent.py
@@ -306,3 +306,22 @@ class Qbittorrent(metaclass=Singleton):
         except Exception as err:
             logger.error(f"获取传输信息出错：{err}")
             return None
+
+    def set_speed_limit(self, download_limit: float = None, upload_limit: float = None) -> bool:
+        """
+        设置速度限制
+        :param download_limit: 下载速度限制，单位KB/s
+        :param upload_limit: 上传速度限制，单位kB/s
+        """
+        if not self.qbc:
+            return False
+        download_limit = download_limit * 1024
+        upload_limit = upload_limit * 1024
+        try:
+            if self.qbc.transfer.upload_limit != upload_limit:
+                self.qbc.transfer.upload_limit = upload_limit
+            if self.qbc.transfer.download_limit != download_limit:
+                self.qbc.transfer.download_limit = download_limit
+        except Exception as err:
+            logger.error(f"设置速度限制出错：{err}")
+            return False
diff --git a/app/modules/transmission/transmission.py b/app/modules/transmission/transmission.py
index e0dfe222..7f8ee5b2 100644
--- a/app/modules/transmission/transmission.py
+++ b/app/modules/transmission/transmission.py
@@ -235,3 +235,32 @@ class Transmission(metaclass=Singleton):
         except Exception as err:
             logger.error(f"获取传输信息出错：{err}")
             return None
+
+    def set_speed_limit(self, download_limit: float = None, upload_limit: float = None):
+        """
+        设置速度限制
+        :param download_limit: 下载速度限制，单位KB/s
+        :param upload_limit: 上传速度限制，单位kB/s
+        """
+        if not self.trc:
+            return
+        try:
+            session = self.trc.get_session()
+            download_limit_enabled = True if download_limit else False
+            upload_limit_enabled = True if upload_limit else False
+            if download_limit_enabled == session.speed_limit_down_enabled and \
+                    upload_limit_enabled == session.speed_limit_up_enabled and \
+                    download_limit == session.speed_limit_down and \
+                    upload_limit == session.speed_limit_up:
+                return
+            self.trc.set_session(
+                speed_limit_down=download_limit if download_limit != session.speed_limit_down
+                else session.speed_limit_down,
+                speed_limit_up=upload_limit if upload_limit != session.speed_limit_up
+                else session.speed_limit_up,
+                speed_limit_down_enabled=download_limit_enabled,
+                speed_limit_up_enabled=upload_limit_enabled
+            )
+        except Exception as err:
+            logger.error(f"设置速度限制出错：{err}")
+            return False