[cleanup] Misc

[yt-dlp.git] / yt_dlp / downloader / niconico.py

diff --git a/yt_dlp/downloader/niconico.py b/yt_dlp/downloader/niconico.py
index 256840d689a05e118237a2e2fa3cc8d69883bc9c..7d8575c2a4c6af45d29bcc24058c3c499515a019 100644 (file)
--- a/yt_dlp/downloader/niconico.py
+++ b/yt_dlp/downloader/niconico.py
@@ -1,22 +1,26 @@
-# coding: utf-8
-from __future__ import unicode_literals
-
+import json
 import threading
+import time
 
+from . import get_suitable_downloader
 from .common import FileDownloader
-from ..downloader import get_suitable_downloader
-from ..extractor.niconico import NiconicoIE
-from ..compat import compat_urllib_request
+from .external import FFmpegFD
+from ..utils import (
+    DownloadError,
+    WebSocketsWrapper,
+    sanitized_Request,
+    str_or_none,
+    try_get,
+)
 
 
 class NiconicoDmcFD(FileDownloader):
     """ Downloading niconico douga from DMC with heartbeat """
 
-    FD_NAME = 'niconico_dmc'
-
     def real_download(self, filename, info_dict):
-        self.to_screen('[%s] Downloading from DMC' % self.FD_NAME)
+        from ..extractor.niconico import NiconicoIE
 
+        self.to_screen('[%s] Downloading from DMC' % self.FD_NAME)
         ie = NiconicoIE(self.ydl)
         info_dict, heartbeat_info_dict = ie._get_heartbeat_info(info_dict)
 
@@ -29,9 +33,11 @@ def real_download(self, filename, info_dict):
         heartbeat_data = heartbeat_info_dict['data'].encode()
         heartbeat_interval = heartbeat_info_dict.get('interval', 30)
 
+        request = sanitized_Request(heartbeat_url, heartbeat_data)
+
         def heartbeat():
             try:
-                compat_urllib_request.urlopen(url=heartbeat_url, data=heartbeat_data)
+                self.ydl.urlopen(request).read()
             except Exception:
                 self.to_screen('[%s] Heartbeat failed' % self.FD_NAME)
 
@@ -52,4 +58,94 @@ def heartbeat():
                 with heartbeat_lock:
                     timer[0].cancel()
                     download_complete = True
-            return success
+        return success
+
+
+class NiconicoLiveFD(FileDownloader):
+    """ Downloads niconico live without being stopped """
+
+    def real_download(self, filename, info_dict):
+        video_id = info_dict['video_id']
+        ws_url = info_dict['url']
+        ws_extractor = info_dict['ws']
+        ws_origin_host = info_dict['origin']
+        cookies = info_dict.get('cookies')
+        live_quality = info_dict.get('live_quality', 'high')
+        live_latency = info_dict.get('live_latency', 'high')
+        dl = FFmpegFD(self.ydl, self.params or {})
+
+        new_info_dict = info_dict.copy()
+        new_info_dict.update({
+            'protocol': 'm3u8',
+        })
+
+        def communicate_ws(reconnect):
+            if reconnect:
+                ws = WebSocketsWrapper(ws_url, {
+                    'Cookies': str_or_none(cookies) or '',
+                    'Origin': f'https://{ws_origin_host}',
+                    'Accept': '*/*',
+                    'User-Agent': self.params['http_headers']['User-Agent'],
+                })
+                if self.ydl.params.get('verbose', False):
+                    self.to_screen('[debug] Sending startWatching request')
+                ws.send(json.dumps({
+                    'type': 'startWatching',
+                    'data': {
+                        'stream': {
+                            'quality': live_quality,
+                            'protocol': 'hls+fmp4',
+                            'latency': live_latency,
+                            'chasePlay': False
+                        },
+                        'room': {
+                            'protocol': 'webSocket',
+                            'commentable': True
+                        },
+                        'reconnect': True,
+                    }
+                }))
+            else:
+                ws = ws_extractor
+            with ws:
+                while True:
+                    recv = ws.recv()
+                    if not recv:
+                        continue
+                    data = json.loads(recv)
+                    if not data or not isinstance(data, dict):
+                        continue
+                    if data.get('type') == 'ping':
+                        # pong back
+                        ws.send(r'{"type":"pong"}')
+                        ws.send(r'{"type":"keepSeat"}')
+                    elif data.get('type') == 'disconnect':
+                        self.write_debug(data)
+                        return True
+                    elif data.get('type') == 'error':
+                        self.write_debug(data)
+                        message = try_get(data, lambda x: x['body']['code'], str) or recv
+                        return DownloadError(message)
+                    elif self.ydl.params.get('verbose', False):
+                        if len(recv) > 100:
+                            recv = recv[:100] + '...'
+                        self.to_screen('[debug] Server said: %s' % recv)
+
+        def ws_main():
+            reconnect = False
+            while True:
+                try:
+                    ret = communicate_ws(reconnect)
+                    if ret is True:
+                        return
+                except BaseException as e:
+                    self.to_screen('[%s] %s: Connection error occured, reconnecting after 10 seconds: %s' % ('niconico:live', video_id, str_or_none(e)))
+                    time.sleep(10)
+                    continue
+                finally:
+                    reconnect = True
+
+        thread = threading.Thread(target=ws_main, daemon=True)
+        thread.start()
+
+        return dl.download(filename, new_info_dict)