From eead16a63576d747f0acf5d0e40e3b36fc70148d Mon Sep 17 00:00:00 2001
From: jxxghp <jxxghp@gmail.com>
Date: Mon, 7 Aug 2023 12:01:07 +0800
Subject: [PATCH] =?UTF-8?q?fix=20=E6=92=AD=E6=94=BE=E9=99=90=E9=80=9F?=
 =?UTF-8?q?=E6=8F=92=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/modules/plex/plex.py             |   6 ++
 app/plugins/speedlimiter/__init__.py | 103 ++++++++++++++++++++++++++-
 2 files changed, 107 insertions(+), 2 deletions(-)

diff --git a/app/modules/plex/plex.py b/app/modules/plex/plex.py
index 545430a8..5c3e5c26 100644
--- a/app/modules/plex/plex.py
+++ b/app/modules/plex/plex.py
@@ -369,3 +369,9 @@ class Plex(metaclass=Singleton):
                                                               image_type="Backdrop")
 
         return eventItem
+
+    def get_plex(self):
+        """
+        获取plex对象，以便直接操作
+        """
+        return self._plex
diff --git a/app/plugins/speedlimiter/__init__.py b/app/plugins/speedlimiter/__init__.py
index 8da88fd9..d591e626 100644
--- a/app/plugins/speedlimiter/__init__.py
+++ b/app/plugins/speedlimiter/__init__.py
@@ -4,9 +4,13 @@ from apscheduler.schedulers.background import BackgroundScheduler
 
 from app.core.config import settings
 from app.log import logger
+from app.modules.plex import Plex
 from app.modules.qbittorrent import Qbittorrent
 from app.modules.transmission import Transmission
 from app.plugins import _PluginBase
+from app.schemas import NotificationType
+from app.utils.http import RequestUtils
+from app.utils.ip import IpUtils
 
 
 class SpeedLimiter(_PluginBase):
@@ -248,12 +252,107 @@ class SpeedLimiter(_PluginBase):
         """
         if not self._qb and not self._tr:
             return
+        # 当前播放的总比特率
+        total_bit_rate = 0
+        # 查询播放中会话
+        playing_sessions = []
+        if settings.MEDIASERVER == "emby":
+            req_url = "{HOST}emby/Sessions?api_key={APIKEY}"
+            try:
+                res = RequestUtils().get_res(req_url)
+                if res and res.status_code == 200:
+                    sessions = res.json()
+                    for session in sessions:
+                        if session.get("NowPlayingItem") and not session.get("PlayState", {}).get("IsPaused"):
+                            playing_sessions.append(session)
+            except Exception as e:
+                logger.error(f"获取Emby播放会话失败：{str(e)}")
+            # 计算有效比特率
+            for session in playing_sessions:
+                if not IpUtils.is_private_ip(session.get("RemoteEndPoint")) \
+                        and session.get("NowPlayingItem", {}).get("MediaType") == "Video":
+                    total_bit_rate += int(session.get("NowPlayingItem", {}).get("Bitrate") or 0)
+        elif settings.MEDIASERVER == "jellyfin":
+            req_url = "{HOST}Sessions?api_key={APIKEY}"
+            try:
+                res = RequestUtils().get_res(req_url)
+                if res and res.status_code == 200:
+                    sessions = res.json()
+                    for session in sessions:
+                        if session.get("NowPlayingItem") and not session.get("PlayState", {}).get("IsPaused"):
+                            playing_sessions.append(session)
+            except Exception as e:
+                logger.error(f"获取Jellyfin播放会话失败：{str(e)}")
+            # 计算有效比特率
+            for session in playing_sessions:
+                if not IpUtils.is_private_ip(session.get("RemoteEndPoint")) \
+                        and session.get("NowPlayingItem", {}).get("MediaType") == "Video":
+                    media_streams = session.get("NowPlayingItem", {}).get("MediaStreams") or []
+                    for media_stream in media_streams:
+                        total_bit_rate += int(media_stream.get("BitRate") or 0)
+        elif settings.MEDIASERVER == "plex":
+            _plex = Plex().get_plex()
+            if _plex:
+                sessions = _plex.sessions()
+                for session in sessions:
+                    bitrate = sum([m.bitrate or 0 for m in session.media])
+                    playing_sessions.append({
+                        "type": session.TAG,
+                        "bitrate": bitrate,
+                        "address": session.player.address
+                    })
+                # 计算有效比特率
+                for session in playing_sessions:
+                    if not IpUtils.is_private_ip(session.get("address")) \
+                            and session.get("type") == "Video":
+                        total_bit_rate += int(session.get("bitrate") or 0)
 
-    def __set_limiter(self):
+        if total_bit_rate:
+            # 当前正在播放，开始限速
+            self.__set_limiter(limit_type="播放", upload_limit=self._play_up_speed,
+                               download_limit=self._play_down_speed)
+        else:
+            # 当前没有播放，开始限速
+            self.__set_limiter(limit_type="未播放", upload_limit=self._noplay_up_speed,
+                               download_limit=self._noplay_down_speed)
+
+    def __set_limiter(self, limit_type: str, upload_limit: float, download_limit: float):
         """
         设置限速
         """
-        pass
+        if not self._qb and not self._tr:
+            return
+        if upload_limit:
+            text = f"上传：{upload_limit} KB/s"
+        else:
+            text = f"上传：未限速"
+        if download_limit:
+            text = f"{text}\n下载：{download_limit} KB/s"
+        else:
+            text = f"{text}\n下载：未限速"
+        try:
+            if self._qb:
+                self._qb.set_speed_limit(download_limit=download_limit, upload_limit=upload_limit)
+                # 发送通知
+                if self._notify:
+                    title = f"Qbittorrent 开始{limit_type}限速"
+                    self.post_message(
+                        mtype=NotificationType.MediaServer,
+                        title=title,
+                        text=text
+                    )
+            if self._tr:
+                self._tr.set_speed_limit(download_limit=download_limit, upload_limit=upload_limit)
+                # 发送通知
+                if self._notify:
+                    title = f"Transmission 开始{limit_type}限速"
+                    self.post_message(
+                        mtype=NotificationType.MediaServer,
+                        title=title,
+                        text=text
+                    )
+        except Exception as e:
+            logger.error(f"设置限速失败：{str(e)}")
 
     def stop_service(self):
         """