128 lines
5.3 KiB
Python
128 lines
5.3 KiB
Python
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import functools
|
|
import http.server
|
|
from pathlib import Path
|
|
import shutil
|
|
import subprocess
|
|
import threading
|
|
|
|
from n_m3u8dl_py.cli import main
|
|
from n_m3u8dl_py.hls import HLSExtractor
|
|
from n_m3u8dl_py.http import HttpClient
|
|
|
|
|
|
def test_hls_master_and_media_playlist(tmp_path: Path) -> None:
|
|
(tmp_path / "video.m3u8").write_text("#EXTM3U\n#EXTINF:2,\na.ts\n#EXT-X-ENDLIST\n")
|
|
master = """#EXTM3U
|
|
#EXT-X-MEDIA:TYPE=AUDIO,GROUP-ID="aud",NAME="English",LANGUAGE="en",URI="audio.m3u8"
|
|
#EXT-X-STREAM-INF:BANDWIDTH=1000,RESOLUTION=640x360,AUDIO="aud"
|
|
video.m3u8
|
|
"""
|
|
extractor = HLSExtractor((tmp_path / "master.m3u8").as_uri(), HttpClient())
|
|
streams = extractor.extract(master)
|
|
assert len(streams) == 2
|
|
video = next(stream for stream in streams if stream.resolution == "640x360")
|
|
extractor.fetch_playlists([video])
|
|
assert video.playlist is not None
|
|
assert video.playlist.segments[0].url.endswith("/a.ts")
|
|
|
|
|
|
def test_local_hls_download_and_binary_merge(tmp_path: Path) -> None:
|
|
(tmp_path / "one.ts").write_bytes(b"first")
|
|
(tmp_path / "two.ts").write_bytes(b"second")
|
|
manifest = tmp_path / "test.m3u8"
|
|
manifest.write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXTINF:1,\ntwo.ts\n#EXT-X-ENDLIST\n")
|
|
output = tmp_path / "out"
|
|
assert main([str(manifest), "--auto-select", "--binary-merge", "--save-name", "joined", "--save-dir", str(output), "--tmp-dir", str(tmp_path / "tmp")]) == 0
|
|
assert (output / "joined.ts").read_bytes() == b"firstsecond"
|
|
|
|
|
|
def test_local_hls_ffmpeg_merge_remuxes_ts_to_mp4(tmp_path: Path) -> None:
|
|
(tmp_path / "one.ts").write_bytes(b"first")
|
|
manifest = tmp_path / "test.m3u8"
|
|
manifest.write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXT-X-ENDLIST\n")
|
|
ffmpeg = tmp_path / "fake-ffmpeg"
|
|
ffmpeg.write_text("#!/bin/sh\nfor argument do output=$argument; done\nprintf remuxed > \"$output\"\n", encoding="utf-8")
|
|
ffmpeg.chmod(0o755)
|
|
output = tmp_path / "out"
|
|
|
|
assert main([
|
|
str(manifest),
|
|
"--auto-select",
|
|
"--ffmpeg-binary-path",
|
|
str(ffmpeg),
|
|
"--save-name",
|
|
"episode",
|
|
"--save-dir",
|
|
str(output),
|
|
"--tmp-dir",
|
|
str(tmp_path / "tmp"),
|
|
]) == 0
|
|
|
|
assert (output / "episode.mp4").read_bytes() == b"remuxed"
|
|
assert not (output / "episode.ts").exists()
|
|
|
|
|
|
def test_local_hls_save_pattern(tmp_path: Path) -> None:
|
|
(tmp_path / "one.ts").write_bytes(b"first")
|
|
manifest = tmp_path / "test.m3u8"
|
|
manifest.write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXT-X-ENDLIST\n")
|
|
output = tmp_path / "out"
|
|
assert main([str(manifest), "--auto-select", "--binary-merge", "--save-name", "episode", "--save-pattern", "<SaveName>_<MediaType>.<Ext>", "--save-dir", str(output), "--tmp-dir", str(tmp_path / "tmp")]) == 0
|
|
assert (output / "episode_video.ts").read_bytes() == b"first"
|
|
|
|
|
|
def test_local_hls_file_uri_decodes_spaces(tmp_path: Path) -> None:
|
|
source = tmp_path / "directory with spaces"
|
|
source.mkdir()
|
|
(source / "one.ts").write_bytes(b"space-safe")
|
|
manifest = source / "test.m3u8"
|
|
manifest.write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXT-X-ENDLIST\n")
|
|
output = tmp_path / "out"
|
|
assert main([str(manifest), "--auto-select", "--binary-merge", "--save-name", "spaced", "--save-dir", str(output), "--tmp-dir", str(tmp_path / "tmp")]) == 0
|
|
assert (output / "spaced.ts").read_bytes() == b"space-safe"
|
|
|
|
|
|
def test_http_hls_download(tmp_path: Path) -> None:
|
|
(tmp_path / "one.ts").write_bytes(b"served-over-http")
|
|
(tmp_path / "test.m3u8").write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXT-X-ENDLIST\n")
|
|
output = tmp_path / "out"
|
|
with _http_server(tmp_path) as base_url:
|
|
assert main([f"{base_url}/test.m3u8", "--no-use-system-proxy", "--auto-select", "--binary-merge", "--save-name", "http", "--save-dir", str(output), "--tmp-dir", str(tmp_path / "tmp")]) == 0
|
|
assert (output / "http.ts").read_bytes() == b"served-over-http"
|
|
|
|
|
|
def test_local_hls_aes128_download(tmp_path: Path) -> None:
|
|
key = bytes(range(16))
|
|
iv = bytes(16)
|
|
plaintext = b"encrypted payload"
|
|
padding = 16 - len(plaintext) % 16
|
|
ciphertext = subprocess.run(
|
|
[shutil.which("openssl") or "openssl", "enc", "-aes-128-cbc", "-K", key.hex(), "-iv", iv.hex(), "-nopad"],
|
|
input=plaintext + bytes([padding]) * padding,
|
|
capture_output=True,
|
|
check=True,
|
|
).stdout
|
|
(tmp_path / "key.bin").write_bytes(key)
|
|
(tmp_path / "encrypted.ts").write_bytes(ciphertext)
|
|
manifest = tmp_path / "encrypted.m3u8"
|
|
manifest.write_text("#EXTM3U\n#EXT-X-KEY:METHOD=AES-128,URI=\"key.bin\",IV=0x00000000000000000000000000000000\n#EXTINF:1,\nencrypted.ts\n#EXT-X-ENDLIST\n")
|
|
output = tmp_path / "out"
|
|
assert main([str(manifest), "--auto-select", "--binary-merge", "--save-name", "plain", "--save-dir", str(output), "--tmp-dir", str(tmp_path / "tmp")]) == 0
|
|
assert (output / "plain.ts").read_bytes() == plaintext
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _http_server(directory: Path):
|
|
handler = functools.partial(http.server.SimpleHTTPRequestHandler, directory=str(directory))
|
|
server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), handler)
|
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
thread.start()
|
|
try:
|
|
yield f"http://127.0.0.1:{server.server_port}"
|
|
finally:
|
|
server.shutdown()
|
|
thread.join()
|