From e36d50c5dd35973c090f87df05d4e94963e8036c Mon Sep 17 00:00:00 2001 From: pukkandan Date: Mon, 21 Jun 2021 22:53:17 +0530 Subject: [PATCH] [websockets] Add `WebSocketFragmentFD` (#399) Necessary for #392 Co-authored by: nao20010128nao, pukkandan --- .github/workflows/build.yml | 4 +-- README.md | 9 ++--- pyinst.py | 12 +++---- requirements.txt | 1 + setup.py | 2 +- yt_dlp/YoutubeDL.py | 7 ++-- yt_dlp/compat.py | 16 +++++++++ yt_dlp/downloader/__init__.py | 3 ++ yt_dlp/downloader/external.py | 8 ++++- yt_dlp/downloader/websocket.py | 59 ++++++++++++++++++++++++++++++++ yt_dlp/extractor/common.py | 2 +- yt_dlp/options.py | 2 +- yt_dlp/postprocessor/__init__.py | 4 +++ yt_dlp/postprocessor/ffmpeg.py | 29 ++++++++++++++++ 14 files changed, 140 insertions(+), 18 deletions(-) create mode 100644 yt_dlp/downloader/websocket.py diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 70c43f208d..a9fa01d542 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 e46edad985..5fc2db3a74 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 0d8ff73c3d..eac97e52d9 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 97a6859ef8..6a982fa369 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ mutagen pycryptodome +websockets diff --git a/setup.py b/setup.py index 8f74c06c1a..d54806f151 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 b4ac1f00ad..aa93b6d1d6 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 863bd2287c..cffaa74a69 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 82d7623f62..e469b512d2 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 8a69b48470..28b1d4e2b4 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 0000000000..088222046b --- /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 b14cf0fc9b..d210ec02fc 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 20211a764f..5351786278 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 d9e369d4d4..98cbe86657 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 4685288a77..83714358e1 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')