Merge pull request #438 from dmunozv04/improve-pairing

Improve pairing

Author: dmunozv04

requirements.txt

@@ -8,3 +8,6 @@ textual==8.2.3
 textual-slider==0.2.0
 xmltodict==1.0.4
 rich_click==1.9.7
+yarl==1.23.0
+zeroconf==0.148.0
+pychromecast==14.0.7

src/iSponsorBlockTV/api_helpers.py

@@ -1,13 +1,19 @@
 import html
 from hashlib import sha256
+from asyncio import FIRST_COMPLETED, Task, create_task, wait
+from typing import AsyncIterator, Collection, TypeVar
 
 from aiohttp import ClientSession
 from cache import AsyncLRU
+from pyytlounge.wrapper import api_base
 
-from . import constants, dial_client
+from . import chromecast_client, constants, dial_client
 from .conditional_ttl_cache import AsyncConditionalTTL
 
 
+_T = TypeVar("_T")
+
+
 def list_to_tuple(function):
     def wrapper(*args):
         args = [tuple(x) if isinstance(x, list) else x for x in args]
@@ -18,6 +24,33 @@ def wrapper(*args):
     return wrapper
 
 
+async def _await_next(iterator: AsyncIterator[_T]) -> _T:
+    return await iterator.__anext__()
+
+
+def _as_task(iterator: AsyncIterator[_T]) -> Task[_T]:
+    return create_task(_await_next(iterator))
+
+
+async def merge_iterators(iterators: Collection[AsyncIterator[_T]]) -> AsyncIterator[_T]:
+    pending_tasks = {_as_task(iterator): iterator for iterator in iterators}
+    try:
+        while pending_tasks:
+            done, _ = await wait(pending_tasks, return_when=FIRST_COMPLETED)
+            for task in done:
+                iterator = pending_tasks.pop(task)
+                try:
+                    yield task.result()
+                except StopAsyncIteration:
+                    continue
+                except Exception:
+                    continue
+                pending_tasks[_as_task(iterator)] = iterator
+    finally:
+        for task in pending_tasks:
+            task.cancel()
+
+
 # Class that handles all the api calls and their cache
 class ApiHelper:
     def __init__(self, config, web_session: ClientSession) -> None:
@@ -29,6 +62,37 @@ def __init__(self, config, web_session: ClientSession) -> None:
         self.num_devices = len(config.devices)
         self.minimum_skip_length = config.minimum_skip_length
 
+    @staticmethod
+    def _normalize_pairing_code(pairing_code):
+        return str(pairing_code).replace("-", "").replace(" ", "")
+
+    async def pair_with_code(self, pairing_code):
+        normalized_code = self._normalize_pairing_code(pairing_code)
+        pair_url = f"{api_base}/pairing/get_screen"
+        pair_data = {"pairing_code": normalized_code}
+
+        async with self.web_session.post(url=pair_url, data=pair_data) as response:
+            if response.status != 200:
+                return None
+            try:
+                pair_response = await response.json()
+            except BaseException:
+                return None
+
+        screen = pair_response.get("screen")
+        if not screen:
+            return None
+
+        screen_id = screen.get("screenId")
+        if not screen_id:
+            return None
+
+        return {
+            "screen_id": screen_id,
+            "name": screen.get("name"),
+            "lounge_token": screen.get("loungeToken"),
+        }
+
     # Not used anymore, maybe it can stay here a little longer
     @AsyncLRU(maxsize=10)
     async def get_vid_id(self, title, artist, api_key, web_session):
@@ -204,8 +268,21 @@ async def mark_viewed_segments(self, uuids):
                 params = {"UUID": i}
                 await self.web_session.post(url, params=params)
 
-    async def discover_youtube_devices_dial(self):
-        """Discovers YouTube devices using DIAL"""
-        dial_screens = await dial_client.discover(self.web_session)
-        # print(dial_screens)
-        return dial_screens
+    async def discover_youtube_devices_dial(self, active=True):
+        """Discovers YouTube devices using DIAL and Chromecast discovery.
+
+        Yields devices as they are discovered instead of waiting for the full
+        discovery cycle to finish.
+        """
+        seen_screen_ids = set()
+        async for device in merge_iterators(
+            [
+                dial_client.discover(self.web_session, self, active=active),
+                chromecast_client.discover(active=active),
+            ]
+        ):
+            screen_id = device.get("screen_id")
+            if not screen_id or screen_id in seen_screen_ids:
+                continue
+            seen_screen_ids.add(screen_id)
+            yield device

src/iSponsorBlockTV/chromecast_client.py

@@ -0,0 +1,131 @@
+"""Chromecast discovery using Cast protocol via pychromecast."""
+
+# skipcq: PYL-R0401
+import asyncio
+from typing import Any
+import logging
+
+import pychromecast
+from pychromecast.controllers.youtube import YouTubeController
+from pychromecast import Chromecast
+
+logger = logging.getLogger(__name__)
+
+
+def _build_device_from_cast(cast: Chromecast) -> dict[str, Any] | None:
+    try:
+        cast.wait(timeout=5)
+
+        yt_controller = YouTubeController()
+        cast.register_handler(yt_controller)
+        try:
+            yt_controller.update_screen_id()
+        except Exception:
+            logger.exception("YouTubeController.update_screen_id() failed")
+
+        # Give the controller a moment to process status messages.
+        cast.wait(timeout=2)
+        screen_id = getattr(yt_controller, "_screen_id", None)
+        if not screen_id:
+            logger.debug("No YouTube MDX screen_id for cast=%s", getattr(cast, "name", cast))
+            return None
+
+        try:
+            cast_name = cast.cast_info.friendly_name
+        except Exception:
+            try:
+                cast_name = cast.name
+            except Exception:
+                cast_name = "Chromecast"
+
+        return {
+            "screen_id": screen_id,
+            "name": cast_name,
+            "offset": 0,
+        }
+    except Exception:
+        logger.exception("Error while building device from cast")
+        return None
+    finally:
+        try:
+            cast.disconnect(timeout=1)
+        except Exception:
+            pass
+
+
+async def discover(active=True):
+    """Discover Chromecast devices and yield YouTube MDX-capable screens.
+
+    The signature mirrors the newer DIAL discovery generator so callers can
+    consume both sources with the same `async for` pattern.
+    """
+    loop = asyncio.get_running_loop()
+    discovered_casts: asyncio.Queue = asyncio.Queue()
+
+    def on_cast_discovered(cast) -> None:
+        print(f"Discovered cast: {cast}")
+        loop.call_soon_threadsafe(discovered_casts.put_nowait, cast)
+
+    browser = pychromecast.get_chromecasts(
+        timeout=5,
+        blocking=False,
+        callback=on_cast_discovered,
+    )
+
+    pending_tasks: set[asyncio.Task] = set()
+    seen_screen_ids: set[str] = set()
+    discovery_deadline = loop.time() + (5 if active else 15)
+
+    try:
+        while True:
+            if not active and pending_tasks:
+                pass
+            elif (
+                loop.time() >= discovery_deadline and discovered_casts.empty() and not pending_tasks
+            ):
+                break
+
+            try:
+                cast = await asyncio.wait_for(discovered_casts.get(), timeout=0.5)
+                pending_tasks.add(
+                    asyncio.create_task(asyncio.to_thread(_build_device_from_cast, cast))
+                )
+                discovery_deadline = max(discovery_deadline, loop.time() + 1)
+            except asyncio.TimeoutError:
+                pass
+
+            done_tasks = {task for task in pending_tasks if task.done()}
+            for task in done_tasks:
+                pending_tasks.discard(task)
+                device = None
+                try:
+                    device = task.result()
+                except Exception:
+                    device = None
+                if not device:
+                    continue
+                screen_id = device.get("screen_id")
+                if not screen_id or screen_id in seen_screen_ids:
+                    continue
+                seen_screen_ids.add(screen_id)
+                yield device
+    finally:
+        for task in pending_tasks:
+            task.cancel()
+        await asyncio.to_thread(browser.stop_discovery)
+
+
+if __name__ == "__main__":
+    import json
+
+    logging.basicConfig(
+        level=logging.DEBUG,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+    logger.debug("Starting chromecast discovery test...")
+
+    async def main():
+        async for device in discover(None):
+            print(json.dumps(device, indent=2))
+
+    asyncio.run(main())

src/iSponsorBlockTV/config_setup.py

@@ -2,7 +2,7 @@
 
 import aiohttp
 
-from . import api_helpers, ytlounge
+from . import api_helpers
 
 # Constants for user input prompts
 USE_PROXY_PROMPT = "Do you want to use system-wide proxy? (y/N)"
@@ -50,22 +50,18 @@ async def create_web_session(use_proxy):
     return aiohttp.ClientSession(trust_env=use_proxy)
 
 
-async def pair_device(web_session: aiohttp.ClientSession):
+async def pair_device(config, web_session: aiohttp.ClientSession, api_helper):
     try:
-        lounge_controller = ytlounge.YtLoungeApi()
-        await lounge_controller.change_web_session(web_session)
+        api_helper = api_helpers.ApiHelper(config, web_session)
         pairing_code = input(PAIRING_CODE_PROMPT)
-        pairing_code = int(
-            pairing_code.replace("-", "").replace(" ", "")
-        )  # remove dashes and spaces
         print("Pairing...")
-        paired = await lounge_controller.pair(pairing_code)
-        if not paired:
-            print("Failed to pair device")
+        paired_device = await api_helper.pair_with_code(pairing_code)
+        if not paired_device:
+            print("Failed to pair device. Try restarting the YouTube app")
             return
         device = {
-            "screen_id": lounge_controller.auth.screen_id,
-            "name": lounge_controller.screen_name,
+            "screen_id": paired_device["screen_id"],
+            "name": paired_device["name"],
         }
         print(f"Paired device: {device['name']}")
         return device
@@ -96,9 +92,10 @@ def main(config, debug: bool) -> None:
             del config["atvs"]
 
     devices = config.devices
+    api_helper = api_helpers.ApiHelper(config, web_session)
     choice = get_yn_input(ADD_MORE_DEVICES_PROMPT.format(num_devices=len(devices)))
     while choice == "y":
-        device = loop.run_until_complete(pair_device(web_session))
+        device = loop.run_until_complete(pair_device(config, web_session, api_helper))
         if device:
             devices.append(device)
         choice = get_yn_input(ADD_MORE_DEVICES_PROMPT.format(num_devices=len(devices)))