[websockets] Add `WebSocketFragmentFD` (#399)

Necessary for #392

Co-authored by: nao20010128nao, pukkandan

14 files changed, 140 insertions, 18 deletions

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 70c43f208..a9fa01d54 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -103,7 +103,7 @@ jobs:
     - name: Upgrade pip and enable wheel support
       run: python -m pip install --upgrade pip setuptools wheel
     - name: Install Requirements
-      run: pip install pyinstaller mutagen pycryptodome
+      run: pip install pyinstaller mutagen pycryptodome websockets
     - name: Bump version
       id: bump_version
       run: python devscripts/update-version.py
@@ -147,7 +147,7 @@ jobs:
     - name: Upgrade pip and enable wheel support
       run: python -m pip install --upgrade pip setuptools wheel
     - name: Install Requirements
-      run: pip install pyinstaller mutagen pycryptodome
+      run: pip install pyinstaller mutagen pycryptodome websockets
     - name: Bump version
       id: bump_version
       run: python devscripts/update-version.py
diff --git a/README.md b/README.md
index e46edad98..5fc2db3a7 100644
--- a/README.md
+++ b/README.md
@@ -182,6 +182,7 @@ While all the other dependancies are optional, `ffmpeg` and `ffprobe` are highly
 * [**sponskrub**](https://github.com/faissaloo/SponSkrub) - For using the [sponskrub options](#sponskrub-sponsorblock-options). Licenced under [GPLv3+](https://github.com/faissaloo/SponSkrub/blob/master/LICENCE.md)
 * [**mutagen**](https://github.com/quodlibet/mutagen) - For embedding thumbnail in certain formats. Licenced under [GPLv2+](https://github.com/quodlibet/mutagen/blob/master/COPYING)
 * [**pycryptodome**](https://github.com/Legrandin/pycryptodome) - For decrypting various data. Licenced under [BSD2](https://github.com/Legrandin/pycryptodome/blob/master/LICENSE.rst)
+* [**websockets**](https://github.com/aaugustin/websockets) - For downloading over websocket. Licenced under [BSD3](https://github.com/aaugustin/websockets/blob/main/LICENSE)
 * [**AtomicParsley**](https://github.com/wez/atomicparsley) - For embedding thumbnail in mp4/m4a if mutagen is not present. Licenced under [GPLv2+](https://github.com/wez/atomicparsley/blob/master/COPYING)
 * [**rtmpdump**](http://rtmpdump.mplayerhq.hu) - For downloading `rtmp` streams. ffmpeg will be used as a fallback. Licenced under [GPLv2+](http://rtmpdump.mplayerhq.hu)
 * [**mplayer**](http://mplayerhq.hu/design7/info.html) or [**mpv**](https://mpv.io) - For downloading `rstp` streams. ffmpeg will be used as a fallback. Licenced under [GPLv2+](https://github.com/mpv-player/mpv/blob/master/Copyright)
@@ -190,14 +191,14 @@ While all the other dependancies are optional, `ffmpeg` and `ffprobe` are highly
 
 To use or redistribute the dependencies, you must agree to their respective licensing terms.
 
-Note that the windows releases are already built with the python interpreter, mutagen and pycryptodome included.
+Note that the windows releases are already built with the python interpreter, mutagen, pycryptodome and websockets included.
 
 ### COMPILE
 
 **For Windows**:
-To build the Windows executable, you must have pyinstaller (and optionally mutagen and pycryptodome)
+To build the Windows executable, you must have pyinstaller (and optionally mutagen, pycryptodome, websockets)
 
-    python3 -m pip install --upgrade pyinstaller mutagen pycryptodome
+    python3 -m pip install --upgrade pyinstaller mutagen pycryptodome websockets
 
 Once you have all the necessary dependencies installed, just run `py pyinst.py`. The executable will be built for the same architecture (32/64 bit) as the python used to build it.
 
@@ -1141,7 +1142,7 @@ You can change the criteria for being considered the `best` by using `-S` (`--fo
  - `lang`: Language preference as given by the extractor
  - `quality`: The quality of the format as given by the extractor
  - `source`: Preference of the source as given by the extractor
- - `proto`: Protocol used for download (`https`/`ftps` > `http`/`ftp` > `m3u8_native` > `m3u8` > `http_dash_segments` > other > `mms`/`rtsp` > unknown > `f4f`/`f4m`)
+ - `proto`: Protocol used for download (`https`/`ftps` > `http`/`ftp` > `m3u8_native`/`m3u8` > `http_dash_segments`> `websocket_frag` > other > `mms`/`rtsp` > unknown > `f4f`/`f4m`)
  - `vcodec`: Video Codec (`av01` > `vp9.2` > `vp9` > `h265` > `h264` > `vp8` > `h263` > `theora` > other > unknown)
  - `acodec`: Audio Codec (`opus` > `vorbis` > `aac` > `mp4a` > `mp3` > `ac3` > `dts` > other > unknown)
  - `codec`: Equivalent to `vcodec,acodec`
diff --git a/pyinst.py b/pyinst.py
index 0d8ff73c3..eac97e52d 100644
--- a/pyinst.py
+++ b/pyinst.py
@@ -6,6 +6,7 @@ import sys
 # import os
 import platform
 
+from PyInstaller.utils.hooks import collect_submodules
 from PyInstaller.utils.win32.versioninfo import (
     VarStruct, VarFileInfo, StringStruct, StringTable,
     StringFileInfo, FixedFileInfo, VSVersionInfo, SetVersion,
@@ -66,16 +67,15 @@ VERSION_FILE = VSVersionInfo(
     ]
 )
 
+dependancies = ['Crypto', 'mutagen'] + collect_submodules('websockets')
+excluded_modules = ['test', 'ytdlp_plugins', 'youtube-dl', 'youtube-dlc']
+
 PyInstaller.__main__.run([
     '--name=yt-dlp%s' % _x86,
     '--onefile',
     '--icon=devscripts/cloud.ico',
-    '--exclude-module=youtube_dl',
-    '--exclude-module=youtube_dlc',
-    '--exclude-module=test',
-    '--exclude-module=ytdlp_plugins',
-    '--hidden-import=mutagen',
-    '--hidden-import=Crypto',
+    *[f'--exclude-module={module}' for module in excluded_modules],
+    *[f'--hidden-import={module}' for module in dependancies],
     '--upx-exclude=vcruntime140.dll',
     'yt_dlp/__main__.py',
 ])
diff --git a/requirements.txt b/requirements.txt
index 97a6859ef..6a982fa36 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,3 @@
 mutagen
 pycryptodome
+websockets
diff --git a/setup.py b/setup.py
index 8f74c06c1..d54806f15 100644
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,7 @@ LONG_DESCRIPTION = '\n\n'.join((
     '**PS**: Some links in this document will not work since this is a copy of the README.md from Github',
     open('README.md', 'r', encoding='utf-8').read()))
 
-REQUIREMENTS = ['mutagen', 'pycryptodome']
+REQUIREMENTS = ['mutagen', 'pycryptodome', 'websockets']
 
 if sys.argv[1:2] == ['py2exe']:
     raise NotImplementedError('py2exe is not currently supported; instead, use "pyinst.py" to build with pyinstaller')
diff --git a/yt_dlp/YoutubeDL.py b/yt_dlp/YoutubeDL.py
index b4ac1f00a..aa93b6d1d 100644
--- a/yt_dlp/YoutubeDL.py
+++ b/yt_dlp/YoutubeDL.py
@@ -127,13 +127,14 @@ from .downloader import (
 )
 from .downloader.rtmp import rtmpdump_version
 from .postprocessor import (
+    get_postprocessor,
+    FFmpegFixupDurationPP,
     FFmpegFixupM3u8PP,
     FFmpegFixupM4aPP,
     FFmpegFixupStretchedPP,
+    FFmpegFixupTimestampPP,
     FFmpegMergerPP,
     FFmpegPostProcessor,
-    # FFmpegSubtitlesConvertorPP,
-    get_postprocessor,
     MoveFilesAfterDownloadPP,
 )
 from .version import __version__
@@ -2723,6 +2724,8 @@ class YoutubeDL(object):
                     downloader = (get_suitable_downloader(info_dict, self.params).__name__
                                   if 'protocol' in info_dict else None)
                     ffmpeg_fixup(downloader == 'HlsFD', 'malformed AAC bitstream detected', FFmpegFixupM3u8PP)
+                    ffmpeg_fixup(downloader == 'WebSocketFragmentFD', 'malformed timestamps detected', FFmpegFixupTimestampPP)
+                    ffmpeg_fixup(downloader == 'WebSocketFragmentFD', 'malformed duration detected', FFmpegFixupDurationPP)
 
                 fixup()
                 try:
diff --git a/yt_dlp/compat.py b/yt_dlp/compat.py
index 863bd2287..cffaa74a6 100644
--- a/yt_dlp/compat.py
+++ b/yt_dlp/compat.py
@@ -3030,6 +3030,21 @@ except AttributeError:
     compat_Match = type(re.compile('').match(''))
 
 
+import asyncio
+try:
+    compat_asyncio_run = asyncio.run
+except AttributeError:
+    def compat_asyncio_run(coro):
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+        loop.run_until_complete(coro)
+
+    asyncio.run = compat_asyncio_run
+
+
 __all__ = [
     'compat_HTMLParseError',
     'compat_HTMLParser',
@@ -3037,6 +3052,7 @@ __all__ = [
     'compat_Match',
     'compat_Pattern',
     'compat_Struct',
+    'compat_asyncio_run',
     'compat_b64decode',
     'compat_basestring',
     'compat_chr',
diff --git a/yt_dlp/downloader/__init__.py b/yt_dlp/downloader/__init__.py
index 82d7623f6..e469b512d 100644
--- a/yt_dlp/downloader/__init__.py
+++ b/yt_dlp/downloader/__init__.py
@@ -24,6 +24,7 @@ from .rtsp import RtspFD
 from .ism import IsmFD
 from .mhtml import MhtmlFD
 from .niconico import NiconicoDmcFD
+from .websocket import WebSocketFragmentFD
 from .youtube_live_chat import YoutubeLiveChatReplayFD
 from .external import (
     get_external_downloader,
@@ -42,6 +43,7 @@ PROTOCOL_MAP = {
     'ism': IsmFD,
     'mhtml': MhtmlFD,
     'niconico_dmc': NiconicoDmcFD,
+    'websocket_frag': WebSocketFragmentFD,
     'youtube_live_chat_replay': YoutubeLiveChatReplayFD,
 }
 
@@ -52,6 +54,7 @@ def shorten_protocol_name(proto, simplify=False):
         'rtmp_ffmpeg': 'rtmp_f',
         'http_dash_segments': 'dash',
         'niconico_dmc': 'dmc',
+        'websocket_frag': 'WSfrag',
     }
     if simplify:
         short_protocol_names.update({
diff --git a/yt_dlp/downloader/external.py b/yt_dlp/downloader/external.py
index 8a69b4847..28b1d4e2b 100644
--- a/yt_dlp/downloader/external.py
+++ b/yt_dlp/downloader/external.py
@@ -347,6 +347,10 @@ class FFmpegFD(ExternalFD):
         # TODO: Fix path for ffmpeg
         return FFmpegPostProcessor().available
 
+    def on_process_started(self, proc, stdin):
+        """ Override this in subclasses  """
+        pass
+
     def _call_downloader(self, tmpfilename, info_dict):
         urls = [f['url'] for f in info_dict.get('requested_formats', [])] or [info_dict['url']]
         ffpp = FFmpegPostProcessor(downloader=self)
@@ -474,6 +478,8 @@ class FFmpegFD(ExternalFD):
         self._debug_cmd(args)
 
         proc = subprocess.Popen(args, stdin=subprocess.PIPE, env=env)
+        if url in ('-', 'pipe:'):
+            self.on_process_started(proc, proc.stdin)
         try:
             retval = proc.wait()
         except BaseException as e:
@@ -482,7 +488,7 @@ class FFmpegFD(ExternalFD):
             # produces a file that is playable (this is mostly useful for live
             # streams). Note that Windows is not affected and produces playable
             # files (see https://github.com/ytdl-org/youtube-dl/issues/8300).
-            if isinstance(e, KeyboardInterrupt) and sys.platform != 'win32':
+            if isinstance(e, KeyboardInterrupt) and sys.platform != 'win32' and url not in ('-', 'pipe:'):
                 process_communicate_or_kill(proc, b'q')
             else:
                 proc.kill()
diff --git a/yt_dlp/downloader/websocket.py b/yt_dlp/downloader/websocket.py
new file mode 100644
index 000000000..088222046
--- /dev/null
+++ b/yt_dlp/downloader/websocket.py
@@ -0,0 +1,59 @@
+import os
+import signal
+import asyncio
+import threading
+
+try:
+    import websockets
+    has_websockets = True
+except ImportError:
+    has_websockets = False
+
+from .common import FileDownloader
+from .external import FFmpegFD
+
+
+class FFmpegSinkFD(FileDownloader):
+    """ A sink to ffmpeg for downloading fragments in any form """
+
+    def real_download(self, filename, info_dict):
+        info_copy = info_dict.copy()
+        info_copy['url'] = '-'
+
+        async def call_conn(proc, stdin):
+            try:
+                await self.real_connection(stdin, info_dict)
+            except (BrokenPipeError, OSError):
+                pass
+            finally:
+                try:
+                    stdin.flush()
+                    stdin.close()
+                except OSError:
+                    pass
+                os.kill(os.getpid(), signal.SIGINT)
+
+        class FFmpegStdinFD(FFmpegFD):
+            @classmethod
+            def get_basename(cls):
+                return FFmpegFD.get_basename()
+
+            def on_process_started(self, proc, stdin):
+                thread = threading.Thread(target=asyncio.run, daemon=True, args=(call_conn(proc, stdin), ))
+                thread.start()
+
+        return FFmpegStdinFD(self.ydl, self.params or {}).download(filename, info_copy)
+
+    async def real_connection(self, sink, info_dict):
+        """ Override this in subclasses """
+        raise NotImplementedError('This method must be implemented by subclasses')
+
+
+class WebSocketFragmentFD(FFmpegSinkFD):
+    async def real_connection(self, sink, info_dict):
+        async with websockets.connect(info_dict['url'], extra_headers=info_dict.get('http_headers', {})) as ws:
+            while True:
+                recv = await ws.recv()
+                if isinstance(recv, str):
+                    recv = recv.encode('utf8')
+                sink.write(recv)
diff --git a/yt_dlp/extractor/common.py b/yt_dlp/extractor/common.py
index b14cf0fc9..d210ec02f 100644
--- a/yt_dlp/extractor/common.py
+++ b/yt_dlp/extractor/common.py
@@ -1487,7 +1487,7 @@ class InfoExtractor(object):
             'acodec': {'type': 'ordered', 'regex': True,
                        'order': ['opus', 'vorbis', 'aac', 'mp?4a?', 'mp3', 'e?a?c-?3', 'dts', '', None, 'none']},
             'proto': {'type': 'ordered', 'regex': True, 'field': 'protocol',
-                      'order': ['(ht|f)tps', '(ht|f)tp$', 'm3u8.+', 'm3u8', '.*dash', '', 'mms|rtsp', 'none', 'f4']},
+                      'order': ['(ht|f)tps', '(ht|f)tp$', 'm3u8.+', '.*dash', 'ws|websocket', '', 'mms|rtsp', 'none', 'f4']},
             'vext': {'type': 'ordered', 'field': 'video_ext',
                      'order': ('mp4', 'webm', 'flv', '', 'none'),
                      'order_free': ('webm', 'mp4', 'flv', '', 'none')},
diff --git a/yt_dlp/options.py b/yt_dlp/options.py
index 20211a764..535178627 100644
--- a/yt_dlp/options.py
+++ b/yt_dlp/options.py
@@ -1165,7 +1165,7 @@ def parseOpts(overrideArguments=None):
             'to give the argument to the specified postprocessor/executable. Supported PP are: '
             'Merger, ExtractAudio, SplitChapters, Metadata, EmbedSubtitle, EmbedThumbnail, '
             'SubtitlesConvertor, ThumbnailsConvertor, VideoRemuxer, VideoConvertor, '
-            'SponSkrub, FixupStretched, FixupM4a and FixupM3u8. '
+            'SponSkrub, FixupStretched, FixupM4a, FixupM3u8, FixupTimestamp and FixupDuration. '
             'The supported executables are: AtomicParsley, FFmpeg, FFprobe, and SponSkrub. '
             'You can also specify "PP+EXE:ARGS" to give the arguments to the specified executable '
             'only when being used by the specified postprocessor. Additionally, for ffmpeg/ffprobe, '
diff --git a/yt_dlp/postprocessor/__init__.py b/yt_dlp/postprocessor/__init__.py
index d9e369d4d..98cbe8665 100644
--- a/yt_dlp/postprocessor/__init__.py
+++ b/yt_dlp/postprocessor/__init__.py
@@ -5,7 +5,9 @@ from .ffmpeg import (
     FFmpegPostProcessor,
     FFmpegEmbedSubtitlePP,
     FFmpegExtractAudioPP,
+    FFmpegFixupDurationPP,
     FFmpegFixupStretchedPP,
+    FFmpegFixupTimestampPP,
     FFmpegFixupM3u8PP,
     FFmpegFixupM4aPP,
     FFmpegMergerPP,
@@ -35,9 +37,11 @@ __all__ = [
     'FFmpegEmbedSubtitlePP',
     'FFmpegExtractAudioPP',
     'FFmpegSplitChaptersPP',
+    'FFmpegFixupDurationPP',
     'FFmpegFixupM3u8PP',
     'FFmpegFixupM4aPP',
     'FFmpegFixupStretchedPP',
+    'FFmpegFixupTimestampPP',
     'FFmpegMergerPP',
     'FFmpegMetadataPP',
     'FFmpegSubtitlesConvertorPP',
diff --git a/yt_dlp/postprocessor/ffmpeg.py b/yt_dlp/postprocessor/ffmpeg.py
index 4685288a7..83714358e 100644
--- a/yt_dlp/postprocessor/ffmpeg.py
+++ b/yt_dlp/postprocessor/ffmpeg.py
@@ -700,6 +700,35 @@ class FFmpegFixupM3u8PP(FFmpegFixupPostProcessor):
         return [], info
 
 
+class FFmpegFixupTimestampPP(FFmpegFixupPostProcessor):
+
+    def __init__(self, downloader=None, trim=0.001):
+        # "trim" should be used when the video contains unintended packets
+        super(FFmpegFixupTimestampPP, self).__init__(downloader)
+        assert isinstance(trim, (int, float))
+        self.trim = str(trim)
+
+    @PostProcessor._restrict_to(images=False)
+    def run(self, info):
+        required_version = '4.4'
+        if is_outdated_version(self._versions[self.basename], required_version):
+            self.report_warning(
+                'A re-encode is needed to fix timestamps in older versions of ffmpeg. '
+                f'Please install ffmpeg {required_version} or later to fixup without re-encoding')
+            opts = ['-vf', 'setpts=PTS-STARTPTS']
+        else:
+            opts = ['-c', 'copy', '-bsf', 'setts=ts=TS-STARTPTS']
+        self._fixup('Fixing frame timestamp', info['filepath'], opts + ['-map', '0', '-dn', '-ss', self.trim])
+        return [], info
+
+
+class FFmpegFixupDurationPP(FFmpegFixupPostProcessor):
+    @PostProcessor._restrict_to(images=False)
+    def run(self, info):
+        self._fixup('Fixing video duration', info['filepath'], ['-c', 'copy', '-map', '0', '-dn'])
+        return [], info
+
+
 class FFmpegSubtitlesConvertorPP(FFmpegPostProcessor):
     SUPPORTED_EXTS = ('srt', 'vtt', 'ass', 'lrc')