Add Plex token caching for users

plex_auto_languages/plex_alert_listener.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+from typing import Callable
+from websocket import WebSocketApp
+from plexapi.alert import AlertListener
+from plexapi.server import PlexServer as BasePlexServer
+
+from plex_auto_languages.utils.logger import get_logger
+
+
+logger = get_logger()
+
+
+class PlexAlertListener(AlertListener):
+
+    def __init__(self, server: BasePlexServer, callback: Callable = None, callbackError: Callable = None):
+        super().__init__(server, callback, callbackError)
+
+    def run(self):
+        url = self._server.url(self.key, includeToken=True).replace("http", "ws")
+        self._ws = WebSocketApp(url, on_message=self._onMessage, on_error=self._onError)
+        self._ws.run_forever(skip_utf8_validation=True)

plex_auto_languages/plex_server.py

@@ -1,4 +1,5 @@
 import time
+import requests
 import itertools
 from typing import Union, Callable
 from datetime import datetime, timedelta
@@ -13,6 +14,7 @@
 from plex_auto_languages.utils.logger import get_logger
 from plex_auto_languages.utils.configuration import Configuration
 from plex_auto_languages.plex_alert_handler import PlexAlertHandler
+from plex_auto_languages.plex_alert_listener import PlexAlertListener
 from plex_auto_languages.track_changes import TrackChanges, NewOrUpdatedTrackChanges
 from plex_auto_languages.utils.notifier import Notifier
 from plex_auto_languages.plex_server_cache import PlexServerCache
@@ -25,9 +27,10 @@
 
 class UnprivilegedPlexServer():
 
-    def __init__(self, url: str, token: str):
+    def __init__(self, url: str, token: str, session: requests.Session = requests.Session()):
+        self._session = session
         self._plex_url = url
-        self._plex = self._get_server(url, token)
+        self._plex = self._get_server(url, token, self._session)
 
     @property
     def connected(self):
@@ -44,9 +47,9 @@ def unique_id(self):
         return self._plex.machineIdentifier
 
     @staticmethod
-    def _get_server(url, token):
+    def _get_server(url: str, token: str, session: requests.Session):
         try:
-            return BasePlexServer(url, token)
+            return BasePlexServer(url, token, session=session)
         except (RequestsConnectionError, Unauthorized):
             return None
 
@@ -120,10 +123,10 @@ def is_alive(self):
         return self.connected and self._alert_listener is not None and self._alert_listener.is_alive()
 
     @staticmethod
-    def _get_server(url: str, token: str, max_tries: int = 5000):
+    def _get_server(url: str, token: str, session: requests.Session, max_tries: int = 5000):
         for _ in range(max_tries):
             try:
-                return BasePlexServer(url, token)
+                return BasePlexServer(url, token, session=session)
             except RequestsConnectionError:
                 logger.warning("ConnectionError: Unable to connect to Plex server, retrying...")
             except Unauthorized:
@@ -146,7 +149,7 @@ def start_alert_listener(self, error_callback: Callable):
         trigger_on_scan = self.config.get("trigger_on_scan")
         trigger_on_activity = self.config.get("trigger_on_activity")
         self._alert_handler = PlexAlertHandler(self, trigger_on_play, trigger_on_scan, trigger_on_activity)
-        self._alert_listener = AlertListener(self._plex, self._alert_handler, error_callback)
+        self._alert_listener = PlexAlertListener(self._plex, self._alert_handler, error_callback)
         logger.info("Starting alert listener")
         self._alert_listener.start()
 
@@ -177,8 +180,12 @@ def get_plex_instance_of_user(self, user_id: Union[int, str]):
         if len(matching_users) == 0:
             logger.error(f"Unable to find user with id '{user_id}'")
             return None
-        user_token = matching_users[0].get_token(self._plex.machineIdentifier)
-        user_plex = UnprivilegedPlexServer(self._plex_url, user_token)
+        user = matching_users[0]
+        user_token = self.cache.get_instance_user_token(user.id)
+        if user_token is None:
+            user_token = user.get_token(self.unique_id)
+            self.cache.set_instance_user_token(user.id, user_token)
+        user_plex = UnprivilegedPlexServer(self._plex_url, user_token, session=self._session)
         if not user_plex.connected:
             logger.error(f"Connection to the Plex server failed for user '{matching_users[0].name}'")
             return None

plex_auto_languages/plex_server_cache.py

@@ -33,6 +33,7 @@ def __init__(self, plex: PlexServer):
         self.recent_activities = {}  # (user_id, item_id): timestamp
         # Users cache
         self._instance_users = []
+        self._instance_user_tokens = {}
         self._instance_users_valid_until = datetime.fromtimestamp(0)
         # Library cache
         self.episode_parts = {}
@@ -87,6 +88,16 @@ def get_instance_users(self, check_validity=True):
     def set_instance_users(self, instance_users):
         self._instance_users = copy.deepcopy(instance_users)
         self._instance_users_valid_until = datetime.now() + timedelta(hours=12)
+        for user in self._instance_users:
+            if str(user.id) in self._instance_user_tokens:
+                continue
+            self._instance_user_tokens[str(user.id)] = user.get_token(self._plex.unique_id)
+
+    def get_instance_user_token(self, user_id):
+        return self._instance_user_tokens.get(str(user_id), None)
+
+    def set_instance_user_token(self, user_id, token):
+        self._instance_user_tokens[str(user_id)] = token
 
     def _get_cache_file_path(self):
         data_dir = self._plex.config.get("data_dir")

tests/test_plex_server.py

@@ -137,6 +137,12 @@ def test_get_plex_instance_of_user(plex):
     assert plex != new_plex
     assert isinstance(new_plex, UnprivilegedPlexServer)
 
+    plex.cache._instance_user_tokens.clear()
+    other_user_id = plex.get_all_user_ids()[1]
+    new_plex = plex.get_plex_instance_of_user(other_user_id)
+    assert plex != new_plex
+    assert isinstance(new_plex, UnprivilegedPlexServer)
+
     new_plex = plex.get_plex_instance_of_user("invalid_user_id")
     assert new_plex is None
 

tests/test_plex_server_cache.py

@@ -6,6 +6,15 @@
 from plex_auto_languages.plex_server_cache import PlexServerCache
 
 
+class FakeUser():
+
+    def __init__(self, user_id):
+        self.id = user_id
+
+    def get_token(self, machine_identifier):
+        return "token"
+
+
 def test_episode_parts(plex):
     assert len(plex.cache.episode_parts) == 46
 
@@ -52,8 +61,20 @@ def test_instance_users(plex):
     assert plex.cache.get_instance_users() is None
     assert plex.cache.get_instance_users(check_validity=False) == []
 
-    plex.cache.set_instance_users(["user1", "user2"])
+    plex.cache.set_instance_users([FakeUser("user1"), FakeUser("user2")])
     assert len(plex.cache.get_instance_users()) == 2
+    user1_token = plex.cache.get_instance_user_token("user1")
+    assert user1_token is not None
+
+    plex.cache.set_instance_users([FakeUser("user1")])
+    assert len(plex.cache.get_instance_users()) == 1
+    assert plex.cache.get_instance_user_token("user1") == user1_token
+
+
+def test_instance_user_tokens(plex):
+    assert plex.cache.get_instance_user_token("fake_user_id") is None
+    plex.cache.set_instance_user_token("user_id", "token")
+    assert plex.cache.get_instance_user_token("user_id") == "token"
 
 
 def test_refresh(plex):