From 55d115027b5df061f7109bba267e8e07bb723f75 Mon Sep 17 00:00:00 2001 From: RemiRigal Date: Tue, 9 Aug 2022 14:54:04 +0200 Subject: [PATCH] Add Plex token caching for users --- plex_auto_languages/plex_alert_listener.py | 21 ++++++++++++++++++ plex_auto_languages/plex_server.py | 25 ++++++++++++++-------- plex_auto_languages/plex_server_cache.py | 11 ++++++++++ tests/test_plex_server.py | 6 ++++++ tests/test_plex_server_cache.py | 23 +++++++++++++++++++- 5 files changed, 76 insertions(+), 10 deletions(-) create mode 100644 plex_auto_languages/plex_alert_listener.py diff --git a/plex_auto_languages/plex_alert_listener.py b/plex_auto_languages/plex_alert_listener.py new file mode 100644 index 0000000..1b22e7d --- /dev/null +++ b/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) diff --git a/plex_auto_languages/plex_server.py b/plex_auto_languages/plex_server.py index 1c9a33f..a6ecd30 100644 --- a/plex_auto_languages/plex_server.py +++ b/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 diff --git a/plex_auto_languages/plex_server_cache.py b/plex_auto_languages/plex_server_cache.py index a82eebb..a0260e8 100644 --- a/plex_auto_languages/plex_server_cache.py +++ b/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") diff --git a/tests/test_plex_server.py b/tests/test_plex_server.py index 7e00c23..0b57f1e 100644 --- a/tests/test_plex_server.py +++ b/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 diff --git a/tests/test_plex_server_cache.py b/tests/test_plex_server_cache.py index 0f25766..4c13cc3 100644 --- a/tests/test_plex_server_cache.py +++ b/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):