From d9a6507fe6da64c3a23d4665f265e80569adbbaa Mon Sep 17 00:00:00 2001
From: Mozi <29089388+pzhlkj6612@users.noreply.github.com>
Date: Tue, 12 Mar 2024 15:34:37 +0000
Subject: [PATCH] [ie/niconico] Support "--load-info-json" by saving WebSocket
 url

aka "--load-info".

Don't save a Response object to info JSON. Just create a new WebSocket
  connection during the download. Due to Niconico's logic, the manifest
  m3u8 url will be unusable soon if there is no active WebSocket
  connection, so the reconnection will give us a valid manifest m3u8,
  unless the WebSocket url has already expired.
---
 yt_dlp/downloader/niconico.py | 53 +++++++++++++++++------------------
 yt_dlp/extractor/niconico.py  | 15 ++++++----
 2 files changed, 34 insertions(+), 34 deletions(-)

diff --git a/yt_dlp/downloader/niconico.py b/yt_dlp/downloader/niconico.py
index ea00a9a8e..6ec2c8766 100644
--- a/yt_dlp/downloader/niconico.py
+++ b/yt_dlp/downloader/niconico.py
@@ -79,33 +79,33 @@ def _ws_context(self, info_dict):
 
         video_id = info_dict['id']
         live_latency = info_dict['downloader_options']['live_latency']
-        self.ws = info_dict['downloader_options']['ws']
+        ws_url = info_dict['downloader_options']['ws_url']
+
+        self.ws = None
 
         self.m3u8_lock = threading.Event()
-        self.m3u8_url = info_dict['manifest_url']
-        self.m3u8_lock.set()
+        self.m3u8_url = None
 
-        def communicate_ws(reconnect):
-            if reconnect:
-                self.ws = self.ydl.urlopen(Request(self.ws.url, headers=info_dict.get('http_headers')))
-                if self.ydl.params.get('verbose', False):
-                    self.to_screen('[debug] Sending startWatching request')
-                self.ws.send(json.dumps({
-                    'type': 'startWatching',
-                    'data': {
-                        'stream': {
-                            'quality': 'abr',
-                            'protocol': 'hls',
-                            'latency': live_latency,
-                            'chasePlay': False
-                        },
-                        'room': {
-                            'protocol': 'webSocket',
-                            'commentable': True
-                        },
-                        'reconnect': True,
-                    }
-                }))
+        def communicate_ws():
+            self.ws = self.ydl.urlopen(Request(ws_url, headers=info_dict.get('http_headers')))
+            if self.ydl.params.get('verbose', False):
+                self.to_screen('[debug] Sending startWatching request')
+            self.ws.send(json.dumps({
+                'type': 'startWatching',
+                'data': {
+                    'stream': {
+                        'quality': 'abr',
+                        'protocol': 'hls',
+                        'latency': live_latency,
+                        'chasePlay': False
+                    },
+                    'room': {
+                        'protocol': 'webSocket',
+                        'commentable': True
+                    },
+                    'reconnect': True,
+                }
+            }))
             with self.ws:
                 while True:
                     recv = self.ws.recv()
@@ -136,10 +136,9 @@ def communicate_ws(reconnect):
         stopped = threading.Event()
 
         def ws_main():
-            reconnect = False
             while not stopped.is_set():
                 try:
-                    communicate_ws(reconnect)
+                    communicate_ws()
                     break  # Disconnected
                 except BaseException as e:  # Including TransportError
                     if stopped.is_set():
@@ -150,8 +149,6 @@ def ws_main():
                     self.to_screen('[%s] %s: Connection error occured, reconnecting after %d seconds: %s' % ('niconico:live', video_id, self._WEBSOCKET_RECONNECT_DELAY, str_or_none(e)))
                     time.sleep(self._WEBSOCKET_RECONNECT_DELAY)
 
-                    reconnect = True
-
             self.m3u8_lock.set()  # Release possible locks
 
         thread = threading.Thread(target=ws_main, daemon=True)
diff --git a/yt_dlp/extractor/niconico.py b/yt_dlp/extractor/niconico.py
index cba25b532..5eee857f5 100644
--- a/yt_dlp/extractor/niconico.py
+++ b/yt_dlp/extractor/niconico.py
@@ -956,7 +956,10 @@ class NiconicoLiveIE(InfoExtractor):
 
     _KNOWN_LATENCY = ('high', 'low')
 
-    def _yield_formats(self, ws, video_id, latency, is_live):
+    def _yield_formats(self, ws_url, headers, latency, video_id, is_live):
+        ws = self._request_webpage(
+            Request(ws_url, headers=headers), video_id, note='Connecting to WebSocket server')
+
         self.write_debug('[debug] Sending HLS server request')
         ws.send(json.dumps({
             'type': 'startWatching',
@@ -998,6 +1001,8 @@ def _yield_formats(self, ws, video_id, latency, is_live):
                     recv = recv[:100] + '...'
                 self.write_debug('Server said: %s' % recv)
 
+        ws.close()
+
         formats = self._extract_m3u8_formats(m3u8_url, video_id, ext='mp4', live=is_live)
         for fmt, q in zip(formats, reversed(qualities[1:])):
             fmt.update({
@@ -1014,14 +1019,11 @@ def _real_extract(self, url):
         embedded_data = self._parse_json(unescapeHTML(self._search_regex(
             r'<script\s+id="embedded-data"\s*data-props="(.+?)"', webpage, 'embedded data')), video_id)
 
-        ws = None
         ws_url = traverse_obj(embedded_data, ('site', 'relive', 'webSocketUrl'))
         if ws_url:
             ws_url = update_url_query(ws_url, {
                 'frontend_id': traverse_obj(embedded_data, ('site', 'frontendId')) or '9',
             })
-            ws = self._request_webpage(
-                Request(ws_url, headers=headers), video_id, note='Connecting to WebSocket server')
         else:
             self.raise_no_formats('The live hasn\'t started yet or already ended.', expected=True)
 
@@ -1074,10 +1076,11 @@ def _real_extract(self, url):
             'timestamp': int_or_none(traverse_obj(embedded_data, ('program', 'openTime'))),
             'live_status': live_status,
             'thumbnails': thumbnails,
-            'formats': [*self._yield_formats(ws, video_id, latency, live_status == 'is_live')] if ws else None,
+            'formats': [*self._yield_formats(
+                ws_url, headers, latency, video_id, live_status == 'is_live')] if ws_url else None,
             'http_headers': headers,
             'downloader_options': {
                 'live_latency': latency,
-                'ws': ws,
+                'ws_url': ws_url,
             },
         }