From f20b8f135687156f32ca36d7ceb0c34d934a2b74 Mon Sep 17 00:00:00 2001 From: ApfelTeeSaft <91074565+ApfelTeeSaft@users.noreply.github.com> Date: Mon, 8 Jun 2026 12:02:00 +0200 Subject: [PATCH] ho lee fuk --- .gitignore | 7 + README.md | 144 +++++++++++ examples/basic_playout.py | 34 +++ examples/dynamic_insertion.py | 31 +++ examples/reconnection.py | 28 +++ pyproject.toml | 20 ++ restream_playout/__init__.py | 18 ++ restream_playout/engine.py | 252 +++++++++++++++++++ restream_playout/events.py | 38 +++ restream_playout/exceptions.py | 19 ++ restream_playout/filler.py | 152 ++++++++++++ restream_playout/flv.py | 98 ++++++++ restream_playout/models.py | 68 ++++++ restream_playout/mp4.py | 358 +++++++++++++++++++++++++++ restream_playout/queue.py | 50 ++++ restream_playout/rtmp.py | 425 +++++++++++++++++++++++++++++++++ restream_playout/session.py | 126 ++++++++++ tests/__init__.py | 1 + tests/mp4_builder.py | 86 +++++++ tests/test_engine.py | 72 ++++++ tests/test_filler_flv.py | 35 +++ tests/test_mp4.py | 29 +++ tests/test_purity.py | 29 +++ tests/test_queue.py | 27 +++ tests/test_rtmp.py | 80 +++++++ tests/test_session.py | 73 ++++++ 26 files changed, 2300 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 examples/basic_playout.py create mode 100644 examples/dynamic_insertion.py create mode 100644 examples/reconnection.py create mode 100644 pyproject.toml create mode 100644 restream_playout/__init__.py create mode 100644 restream_playout/engine.py create mode 100644 restream_playout/events.py create mode 100644 restream_playout/exceptions.py create mode 100644 restream_playout/filler.py create mode 100644 restream_playout/flv.py create mode 100644 restream_playout/models.py create mode 100644 restream_playout/mp4.py create mode 100644 restream_playout/queue.py create mode 100644 restream_playout/rtmp.py create mode 100644 restream_playout/session.py create mode 100644 tests/__init__.py create mode 100644 tests/mp4_builder.py create mode 100644 tests/test_engine.py create mode 100644 tests/test_filler_flv.py create mode 100644 tests/test_mp4.py create mode 100644 tests/test_purity.py create mode 100644 tests/test_queue.py create mode 100644 tests/test_rtmp.py create mode 100644 tests/test_session.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a43e741 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +__pycache__/ +*.py[cod] +.pytest_cache/ +build/ +dist/ +*.egg-info/ + diff --git a/README.md b/README.md new file mode 100644 index 0000000..cd727e8 --- /dev/null +++ b/README.md @@ -0,0 +1,144 @@ +# restream_playout + +`restream_playout` is a self-contained, pure-Python RTMP playout library. It +passes H.264 video and AAC audio from MP4 files directly to an RTMP publishing +endpoint and emits black video with silent audio whenever the queue is empty. +The RTMP publishing session stays open while playout changes between filler and +queued clips. + +It does not invoke or depend on FFmpeg, ffprobe, GStreamer, VLC, OBS, MPV, +external commands, subprocesses, or native codec libraries. It does not +transcode. + +## Requirements + +- Python 3.11 or newer +- A plain `rtmp://` publishing endpoint +- Input MP4 files with exactly one `avc1` H.264 video track and one `mp4a` AAC + audio track +- Non-fragmented MP4 with four-byte AVC NAL lengths and a zero-based timeline + +Files that cannot be passed through correctly are rejected by `enqueue()`. +Fragmented MP4, edit lists, additional tracks, unsupported codecs, and +encrypted media are outside the pass-through contract. + +## Install + +```bash +python -m pip install . +``` + +The package has no runtime dependencies outside the Python standard library. + +## Basic Playout + +```python +import logging + +from restream_playout import PlayoutEngine + +logging.basicConfig(level=logging.INFO) + +engine = PlayoutEngine( + rtmp_url="rtmp://live.restream.io/live/STREAM_KEY", +) + +engine.on_clip_started(lambda path: logging.info("started %s", path)) +engine.on_clip_finished(lambda path: logging.info("finished %s", path)) +engine.on_queue_empty(lambda: logging.info("queue empty; filler active")) +engine.on_stream_connected(lambda: logging.info("connected")) +engine.on_stream_disconnected(lambda: logging.warning("disconnected; reconnecting")) + +engine.enqueue_many(["intro.mp4", "video1.mp4"]) +engine.start() + +try: + input("Press Enter to stop\n") +finally: + engine.stop() +``` + +`start()` creates one background thread containing the asyncio playout loop. +`enqueue()`, `enqueue_many()`, `clear()`, and the methods on `engine.queue` are +safe to call from other threads. + +## Dynamic Insertion + +Calling `enqueue()` while filler is active switches to the new clip at the next +filler packet boundary without closing the RTMP connection: + +```python +engine.start() +engine.enqueue("scheduled.mp4") + +# Later, from an API handler, scheduler, or another application thread: +engine.enqueue("breaking-news.mp4") +``` + +Queued clips play in FIFO order. Dynamic insertion does not interrupt a clip +that is already playing. + +## Reconnection + +`StreamSession` monitors the RTMP reader and sends periodic RTMP ping requests. +On read, write, or heartbeat failure, the engine reconnects with capped +exponential backoff. After reconnecting it resends metadata plus H.264 and AAC +sequence headers before retrying media delivery. + +```python +engine = PlayoutEngine( + "rtmp://live.restream.io/live/STREAM_KEY", + reconnect_min_delay=1.0, + reconnect_max_delay=30.0, + heartbeat_interval=10.0, +) + +engine.on_stream_disconnected(lambda: alert("RTMP disconnected")) +engine.on_stream_connected(lambda: alert("RTMP publishing")) +engine.start() +``` + +The media timeline remains monotonic across clips, filler transitions, and +reconnections. If network recovery takes longer than `max_lag_seconds`, wall +clock pacing is rebased so buffered packets are not emitted in a burst. + +## Architecture + +- `engine.py`: synchronous public API, background asyncio loop, pacing, filler + transitions, event dispatch, and codec bootstrap handling. +- `session.py`: RTMP connection lifecycle, heartbeat, reconnect backoff, and + connection generation tracking. +- `rtmp.py`: RTMP handshake, AMF0 commands, chunk writer/parser, publish flow, + acknowledgements, and ping control messages. +- `mp4.py`: ISO BMFF box and sample-table parser for H.264/AAC pass-through. +- `flv.py`: AMF0 metadata plus AVC/AAC FLV payload and FLV tag generation. +- `filler.py`: Python-generated baseline H.264 black IDR access unit and + continuously scheduled AAC silence. +- `queue.py` and `events.py`: thread-safe queue and callback dispatch. + +The filler video is a valid 16x16 constrained-baseline H.264 stream. It is +intentionally tiny because a pure-Python I-PCM encoder is used; this keeps idle +CPU and network use low. Each source clip supplies its own codec sequence +headers when playout switches. + +## Operations Notes + +- Use a dedicated Restream stream key and protect it like a password. +- Preflight media before scheduling it by calling `parse_mp4(path)`. +- Keep source frame rates, resolutions, H.264 profiles, and AAC layouts + consistent for the smoothest player-side transitions. +- Configure logging handlers appropriate for the service environment. +- Run the process under a service manager so process-level failures are + restarted. The library handles network/session recovery, not process + supervision. + +## Tests + +```bash +python -m unittest discover -v +``` + +The suite includes a synthetic H.264/AAC MP4 demux test and an in-process RTMP +server test covering handshake, command negotiation, publishing, and message +delivery. + diff --git a/examples/basic_playout.py b/examples/basic_playout.py new file mode 100644 index 0000000..9c45abc --- /dev/null +++ b/examples/basic_playout.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import logging +import signal +import threading + +from restream_playout import PlayoutEngine + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s") + stopped = threading.Event() + engine = PlayoutEngine("rtmp://live.restream.io/live/STREAM_KEY") + + engine.on_clip_started(lambda path: logging.info("clip started: %s", path)) + engine.on_clip_finished(lambda path: logging.info("clip finished: %s", path)) + engine.on_queue_empty(lambda: logging.info("filler active")) + engine.on_stream_connected(lambda: logging.info("stream connected")) + engine.on_stream_disconnected(lambda: logging.warning("stream disconnected")) + + signal.signal(signal.SIGINT, lambda *_: stopped.set()) + signal.signal(signal.SIGTERM, lambda *_: stopped.set()) + + engine.enqueue_many(["intro.mp4", "video1.mp4"]) + engine.start() + try: + stopped.wait() + finally: + engine.stop() + + +if __name__ == "__main__": + main() + diff --git a/examples/dynamic_insertion.py b/examples/dynamic_insertion.py new file mode 100644 index 0000000..46986a4 --- /dev/null +++ b/examples/dynamic_insertion.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import logging +import threading +import time + +from restream_playout import PlayoutEngine + + +def insert_later(engine: PlayoutEngine) -> None: + time.sleep(30) + engine.enqueue("new_clip.mp4") + + +def main() -> None: + logging.basicConfig(level=logging.INFO) + engine = PlayoutEngine("rtmp://live.restream.io/live/STREAM_KEY") + engine.start() + + insertion = threading.Thread(target=insert_later, args=(engine,)) + insertion.start() + try: + insertion.join() + input("New clip queued. Press Enter to stop.\n") + finally: + engine.stop() + + +if __name__ == "__main__": + main() + diff --git a/examples/reconnection.py b/examples/reconnection.py new file mode 100644 index 0000000..6dbe8c1 --- /dev/null +++ b/examples/reconnection.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +import logging + +from restream_playout import PlayoutEngine + + +def main() -> None: + logging.basicConfig(level=logging.INFO) + engine = PlayoutEngine( + "rtmp://live.restream.io/live/STREAM_KEY", + reconnect_min_delay=1, + reconnect_max_delay=30, + heartbeat_interval=10, + ) + engine.on_stream_connected(lambda: logging.info("publishing is healthy")) + engine.on_stream_disconnected(lambda: logging.warning("connection lost; automatic recovery active")) + engine.enqueue("program.mp4") + engine.start() + try: + input("Press Enter to stop.\n") + finally: + engine.stop() + + +if __name__ == "__main__": + main() + diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..89ffe6a --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,20 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "restream-playout" +version = "0.1.0" +description = "Pure-Python H.264/AAC MP4 playout to RTMP" +readme = "README.md" +requires-python = ">=3.11" +license = { text = "do whatever you want" } +authors = [{ name = "apfelteesaft" }] +dependencies = [] + +[tool.setuptools.packages.find] +include = ["restream_playout*"] + +[tool.pytest.ini_options] +testpaths = ["tests"] + diff --git a/restream_playout/__init__.py b/restream_playout/__init__.py new file mode 100644 index 0000000..826be4b --- /dev/null +++ b/restream_playout/__init__.py @@ -0,0 +1,18 @@ +from .engine import PlayoutEngine +from .exceptions import EngineStateError, InvalidMediaError, PlayoutError, RTMPError, StreamDisconnected +from .mp4 import parse_mp4 +from .queue import VideoQueue + +__all__ = [ + "EngineStateError", + "InvalidMediaError", + "PlayoutEngine", + "PlayoutError", + "RTMPError", + "StreamDisconnected", + "VideoQueue", + "parse_mp4", +] + +__version__ = "0.1.0" + diff --git a/restream_playout/engine.py b/restream_playout/engine.py new file mode 100644 index 0000000..bd28c7d --- /dev/null +++ b/restream_playout/engine.py @@ -0,0 +1,252 @@ +from __future__ import annotations + +import asyncio +import logging +import threading +import time +from collections.abc import Callable, Iterable +from pathlib import Path +from typing import Any + +from .events import EventDispatcher +from .exceptions import EngineStateError, StreamDisconnected +from .filler import FillerSource +from .flv import aac_sequence_header, avc_sequence_header, media_payload, metadata_payload +from .models import CodecConfig, MediaKind, MediaPacket, ParsedClip +from .mp4 import parse_mp4 +from .queue import VideoQueue +from .rtmp import RTMP_AUDIO, RTMP_DATA_AMF0, RTMP_VIDEO, RTMPUrl +from .session import StreamSession + + +class PlayoutEngine: + def __init__( + self, + rtmp_url: str, + *, + logger: logging.Logger | None = None, + reconnect_min_delay: float = 1.0, + reconnect_max_delay: float = 30.0, + heartbeat_interval: float = 10.0, + max_lag_seconds: float = 1.0, + ) -> None: + RTMPUrl.parse(rtmp_url) + self.rtmp_url = rtmp_url + self.logger = logger or logging.getLogger("restream_playout") + self.queue = VideoQueue() + self.events = EventDispatcher(self.logger) + self.reconnect_min_delay = reconnect_min_delay + self.reconnect_max_delay = reconnect_max_delay + self.heartbeat_interval = heartbeat_interval + self.max_lag_seconds = max_lag_seconds + self._clips: dict[Path, ParsedClip] = {} + self._clips_lock = threading.RLock() + self._thread: threading.Thread | None = None + self._thread_ready = threading.Event() + self._stop_requested = threading.Event() + self._loop: asyncio.AbstractEventLoop | None = None + self._async_stop: asyncio.Event | None = None + self._session: StreamSession | None = None + self._timeline_ms = 0 + self._wall_origin = 0.0 + self._active_config: CodecConfig | None = None + self._bootstrapped_generation = -1 + + def start(self) -> "PlayoutEngine": + if self._thread is not None and self._thread.is_alive(): + return self + self._stop_requested.clear() + self._thread_ready.clear() + self._thread = threading.Thread(target=self._thread_main, name="restream-playout", daemon=False) + self._thread.start() + if not self._thread_ready.wait(5): + raise EngineStateError("playout thread did not start") + return self + + def stop(self, timeout: float = 15.0) -> None: + thread = self._thread + if thread is None: + return + self._stop_requested.set() + self.queue.wake() + if self._loop is not None and self._async_stop is not None: + self._loop.call_soon_threadsafe(self._async_stop.set) + if thread is threading.current_thread(): + return + thread.join(timeout) + if thread.is_alive(): + raise EngineStateError("playout thread did not stop before timeout") + self._thread = None + + def enqueue(self, path: str | Path) -> ParsedClip: + clip = parse_mp4(path) + with self._clips_lock: + self._clips[clip.path] = clip + self.queue.enqueue(clip.path) + return clip + + def enqueue_many(self, paths: Iterable[str | Path]) -> list[ParsedClip]: + clips = [parse_mp4(path) for path in paths] + with self._clips_lock: + self._clips.update((clip.path, clip) for clip in clips) + self.queue.enqueue_many(clip.path for clip in clips) + return clips + + def clear(self) -> None: + self.queue.clear() + + def on_clip_started(self, callback: Callable[[Path], Any]) -> Callable[[Path], Any]: + return self.events.on("clip_started", callback) + + def on_clip_finished(self, callback: Callable[[Path], Any]) -> Callable[[Path], Any]: + return self.events.on("clip_finished", callback) + + def on_queue_empty(self, callback: Callable[[], Any]) -> Callable[[], Any]: + return self.events.on("queue_empty", callback) + + def on_stream_connected(self, callback: Callable[[], Any]) -> Callable[[], Any]: + return self.events.on("stream_connected", callback) + + def on_stream_disconnected(self, callback: Callable[[], Any]) -> Callable[[], Any]: + return self.events.on("stream_disconnected", callback) + + def _thread_main(self) -> None: + try: + asyncio.run(self._run()) + except Exception: + self.logger.exception("playout engine stopped unexpectedly") + + async def _run(self) -> None: + self._loop = asyncio.get_running_loop() + self._async_stop = asyncio.Event() + self._session = StreamSession( + self.rtmp_url, + stop_event=self._async_stop, + state_callback=self._on_session_state, + logger=self.logger, + reconnect_min_delay=self.reconnect_min_delay, + reconnect_max_delay=self.reconnect_max_delay, + heartbeat_interval=self.heartbeat_interval, + ) + self._timeline_ms = 0 + self._wall_origin = time.monotonic() + self._active_config = None + self._bootstrapped_generation = -1 + self._thread_ready.set() + empty_announced = False + try: + while not self._stop_requested.is_set(): + path = self.queue.pop() + if path is not None: + empty_announced = False + resolved_path = path.expanduser().resolve() + with self._clips_lock: + clip = self._clips.get(resolved_path) + if clip is None: + try: + clip = parse_mp4(resolved_path) + except Exception: + self.logger.exception("queued clip rejected", extra={"path": str(path)}) + continue + with self._clips_lock: + self._clips[resolved_path] = clip + await self._play_clip(clip) + continue + if not empty_announced: + await self.events.emit("queue_empty") + empty_announced = True + await self._play_filler() + finally: + if self._session is not None: + await self._session.close() + + async def _on_session_state(self, connected: bool) -> None: + await self.events.emit("stream_connected" if connected else "stream_disconnected") + + async def _play_clip(self, clip: ParsedClip) -> None: + self.logger.info("clip started", extra={"path": str(clip.path)}) + await self.events.emit("clip_started", clip.path) + self._activate_config(clip.config) + base = self._timeline_ms + try: + for packet in clip.packets(): + if self._stop_requested.is_set(): + return + timestamp = base + packet.dts_ms + await self._wait_until(timestamp) + if self._stop_requested.is_set(): + return + await self._publish_packet(packet, timestamp) + self._timeline_ms = max(self._timeline_ms, timestamp + packet.duration_ms) + except Exception: + self.logger.exception("clip playback failed", extra={"path": str(clip.path)}) + else: + self._timeline_ms = max(self._timeline_ms, base + clip.duration_ms) + self.logger.info("clip finished", extra={"path": str(clip.path)}) + await self.events.emit("clip_finished", clip.path) + + async def _play_filler(self) -> None: + filler = FillerSource() + self._activate_config(filler.config) + base = self._timeline_ms + while not self._stop_requested.is_set() and self.queue.peek() is None: + packet = filler.next_packet() + timestamp = base + packet.dts_ms + interrupted = await self._wait_until(timestamp, interrupt_for_queue=True) + if interrupted or self._stop_requested.is_set(): + return + await self._publish_packet(packet, timestamp) + self._timeline_ms = max(self._timeline_ms, timestamp + packet.duration_ms) + + def _activate_config(self, config: CodecConfig) -> None: + if config != self._active_config: + self._active_config = config + self._bootstrapped_generation = -1 + + async def _wait_until(self, timestamp_ms: int, *, interrupt_for_queue: bool = False) -> bool: + target = self._wall_origin + timestamp_ms / 1000 + now = time.monotonic() + if now - target > self.max_lag_seconds: + self._wall_origin = now - timestamp_ms / 1000 + target = now + delay = max(0.0, target - now) + if not delay: + return False + if interrupt_for_queue: + return await asyncio.to_thread(self.queue.wait, delay) + try: + await asyncio.wait_for(self._async_stop.wait(), delay) + except TimeoutError: + pass + return False + + async def _bootstrap(self, generation: int, timestamp: int) -> None: + if self._session is None or self._active_config is None: + raise EngineStateError("playout session is not initialized") + config = self._active_config + await self._session.send_message(RTMP_DATA_AMF0, timestamp, metadata_payload(config), chunk_stream_id=5) + await self._session.send_message(RTMP_VIDEO, timestamp, avc_sequence_header(config), chunk_stream_id=6) + await self._session.send_message(RTMP_AUDIO, timestamp, aac_sequence_header(config), chunk_stream_id=4) + self._bootstrapped_generation = generation + + async def _publish_packet(self, packet: MediaPacket, timestamp: int) -> None: + if self._session is None: + raise EngineStateError("playout session is not initialized") + while not self._stop_requested.is_set(): + try: + generation = await self._session.ensure_connected() + if generation != self._bootstrapped_generation: + await self._bootstrap(generation, timestamp) + message_type = RTMP_VIDEO if packet.kind is MediaKind.VIDEO else RTMP_AUDIO + chunk_stream_id = 6 if message_type == RTMP_VIDEO else 4 + channels = self._active_config.audio_channels if self._active_config else 2 + await self._session.send_message( + message_type, + timestamp, + media_payload(packet, channels), + chunk_stream_id=chunk_stream_id, + ) + return + except StreamDisconnected: + self._bootstrapped_generation = -1 + return diff --git a/restream_playout/events.py b/restream_playout/events.py new file mode 100644 index 0000000..3f774b0 --- /dev/null +++ b/restream_playout/events.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +import asyncio +import inspect +import logging +import threading +from collections import defaultdict +from collections.abc import Callable +from typing import Any + +Callback = Callable[..., Any] + + +class EventDispatcher: + def __init__(self, logger: logging.Logger | None = None) -> None: + self._callbacks: dict[str, list[Callback]] = defaultdict(list) + self._lock = threading.RLock() + self._logger = logger or logging.getLogger(__name__) + + def on(self, event: str, callback: Callback) -> Callback: + with self._lock: + self._callbacks[event].append(callback) + return callback + + async def emit(self, event: str, *args: Any) -> None: + with self._lock: + callbacks = tuple(self._callbacks[event]) + for callback in callbacks: + try: + result = callback(*args) + if inspect.isawaitable(result): + await result + except Exception: + self._logger.exception("event callback failed", extra={"event": event}) + + def emit_from_thread(self, loop: asyncio.AbstractEventLoop, event: str, *args: Any) -> None: + asyncio.run_coroutine_threadsafe(self.emit(event, *args), loop) + diff --git a/restream_playout/exceptions.py b/restream_playout/exceptions.py new file mode 100644 index 0000000..5b85da2 --- /dev/null +++ b/restream_playout/exceptions.py @@ -0,0 +1,19 @@ +class PlayoutError(Exception): + """Base exception for restream_playout.""" + + +class InvalidMediaError(PlayoutError): + """Raised when an input file cannot be passed through safely.""" + + +class RTMPError(PlayoutError): + """Raised for RTMP protocol or server errors.""" + + +class StreamDisconnected(RTMPError): + """Raised when the active RTMP connection is no longer usable.""" + + +class EngineStateError(PlayoutError): + """Raised when an engine operation is invalid for its current state.""" + diff --git a/restream_playout/filler.py b/restream_playout/filler.py new file mode 100644 index 0000000..8298524 --- /dev/null +++ b/restream_playout/filler.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from .models import CodecConfig, MediaKind, MediaPacket + + +class _Bits: + def __init__(self) -> None: + self.bits: list[int] = [] + + def bit(self, value: int) -> None: + self.bits.append(value & 1) + + def uint(self, value: int, width: int) -> None: + self.bits.extend((value >> shift) & 1 for shift in range(width - 1, -1, -1)) + + def ue(self, value: int) -> None: + code = value + 1 + zeros = code.bit_length() - 1 + self.bits.extend([0] * zeros) + self.uint(code, zeros + 1) + + def se(self, value: int) -> None: + self.ue(-2 * value if value <= 0 else 2 * value - 1) + + def align(self) -> None: + while len(self.bits) % 8: + self.bit(0) + + def trailing(self) -> None: + self.bit(1) + self.align() + + def bytes(self) -> bytes: + self.align() + return bytes(sum(self.bits[index + bit] << (7 - bit) for bit in range(8)) for index in range(0, len(self.bits), 8)) + + +def _escape_rbsp(data: bytes) -> bytes: + output = bytearray() + zeros = 0 + for byte in data: + if zeros >= 2 and byte <= 3: + output.append(3) + zeros = 0 + output.append(byte) + zeros = zeros + 1 if byte == 0 else 0 + return bytes(output) + + +def _sps() -> bytes: + bits = _Bits() + bits.uint(66, 8) + bits.uint(0xC0, 8) + bits.uint(10, 8) + bits.ue(0) + bits.ue(0) + bits.ue(0) + bits.ue(0) + bits.ue(1) + bits.bit(0) + bits.ue(0) + bits.ue(0) + bits.bit(1) + bits.bit(1) + bits.bit(0) + bits.bit(0) + bits.trailing() + return b"\x67" + _escape_rbsp(bits.bytes()) + + +def _pps() -> bytes: + bits = _Bits() + bits.ue(0) + bits.ue(0) + bits.bit(0) + bits.bit(0) + bits.ue(0) + bits.ue(0) + bits.ue(0) + bits.bit(0) + bits.uint(0, 2) + bits.se(0) + bits.se(0) + bits.se(0) + bits.bit(1) + bits.bit(0) + bits.bit(0) + bits.trailing() + return b"\x68" + _escape_rbsp(bits.bytes()) + + +def _black_idr() -> bytes: + bits = _Bits() + bits.ue(0) + bits.ue(2) + bits.ue(0) + bits.uint(0, 4) + bits.ue(0) + bits.uint(0, 4) + bits.bit(0) + bits.bit(0) + bits.se(0) + bits.ue(1) + bits.ue(25) + bits.align() + for _ in range(256): + bits.uint(16, 8) + for _ in range(128): + bits.uint(128, 8) + bits.trailing() + return b"\x65" + _escape_rbsp(bits.bytes()) + + +SPS = _sps() +PPS = _pps() +IDR = _black_idr() +AVCC = ( + bytes((1, SPS[1], SPS[2], SPS[3], 0xFF, 0xE1)) + + len(SPS).to_bytes(2, "big") + + SPS + + b"\x01" + + len(PPS).to_bytes(2, "big") + + PPS +) +AVC_SAMPLE = len(IDR).to_bytes(4, "big") + IDR +AAC_SILENCE = bytes.fromhex("211004608c1c") + + +@dataclass(slots=True) +class FillerSource: + frame_rate: int = 30 + audio_sample_rate: int = 44100 + _video_index: int = 0 + _audio_index: int = 0 + + @property + def config(self) -> CodecConfig: + return CodecConfig(AVCC, b"\x12\x10", 16, 16, float(self.frame_rate), self.audio_sample_rate, 2) + + def next_packet(self) -> MediaPacket: + video_dts = round(self._video_index * 1000 / self.frame_rate) + audio_dts = round(self._audio_index * 1024 * 1000 / self.audio_sample_rate) + if video_dts <= audio_dts: + next_dts = round((self._video_index + 1) * 1000 / self.frame_rate) + self._video_index += 1 + return MediaPacket(MediaKind.VIDEO, AVC_SAMPLE, video_dts, video_dts, max(1, next_dts - video_dts), True) + next_dts = round((self._audio_index + 1) * 1024 * 1000 / self.audio_sample_rate) + self._audio_index += 1 + return MediaPacket(MediaKind.AUDIO, AAC_SILENCE, audio_dts, audio_dts, max(1, next_dts - audio_dts), True) + diff --git a/restream_playout/flv.py b/restream_playout/flv.py new file mode 100644 index 0000000..44f8c84 --- /dev/null +++ b/restream_playout/flv.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +import struct +from collections.abc import Mapping +from typing import Any + +from .models import CodecConfig, MediaKind, MediaPacket + + +def _u24(value: int) -> bytes: + return value.to_bytes(3, "big", signed=False) + + +def _s24(value: int) -> bytes: + if not -(1 << 23) <= value < (1 << 23): + raise ValueError("signed 24-bit value out of range") + return (value & 0xFFFFFF).to_bytes(3, "big") + + +def amf0_encode(value: Any) -> bytes: + if value is None: + return b"\x05" + if isinstance(value, bool): + return b"\x01" + bytes((int(value),)) + if isinstance(value, (int, float)): + return b"\x00" + struct.pack(">d", float(value)) + if isinstance(value, str): + encoded = value.encode("utf-8") + if len(encoded) > 0xFFFF: + return b"\x0c" + len(encoded).to_bytes(4, "big") + encoded + return b"\x02" + len(encoded).to_bytes(2, "big") + encoded + if isinstance(value, Mapping): + body = bytearray(b"\x03") + for key, item in value.items(): + encoded_key = str(key).encode("utf-8") + body += len(encoded_key).to_bytes(2, "big") + encoded_key + body += amf0_encode(item) + return bytes(body + b"\x00\x00\x09") + raise TypeError(f"unsupported AMF0 value: {type(value).__name__}") + + +def amf0_ecma_array(values: Mapping[str, Any]) -> bytes: + body = bytearray(b"\x08" + len(values).to_bytes(4, "big")) + for key, value in values.items(): + encoded_key = key.encode("utf-8") + body += len(encoded_key).to_bytes(2, "big") + encoded_key + body += amf0_encode(value) + return bytes(body + b"\x00\x00\x09") + + +def metadata_payload(config: CodecConfig) -> bytes: + values = { + "width": config.width, + "height": config.height, + "framerate": config.frame_rate, + "videocodecid": 7, + "audiosamplerate": config.audio_sample_rate, + "audiosamplesize": 16, + "stereo": config.audio_channels > 1, + "audiocodecid": 10, + "encoder": "restream_playout", + } + return amf0_encode("@setDataFrame") + amf0_encode("onMetaData") + amf0_ecma_array(values) + + +def avc_sequence_header(config: CodecConfig) -> bytes: + return b"\x17\x00\x00\x00\x00" + config.avc_decoder_config + + +def aac_sequence_header(config: CodecConfig) -> bytes: + header = 0xAF if config.audio_channels > 1 else 0xAE + return bytes((header, 0)) + config.audio_specific_config + + +def media_payload(packet: MediaPacket, audio_channels: int = 2) -> bytes: + if packet.kind is MediaKind.VIDEO: + frame = 0x17 if packet.keyframe else 0x27 + return bytes((frame, 1)) + _s24(packet.pts_ms - packet.dts_ms) + packet.data + header = 0xAF if audio_channels > 1 else 0xAE + return bytes((header, 1)) + packet.data + + +def flv_header(has_audio: bool = True, has_video: bool = True) -> bytes: + flags = (4 if has_audio else 0) | (1 if has_video else 0) + return b"FLV\x01" + bytes((flags,)) + b"\x00\x00\x00\x09\x00\x00\x00\x00" + + +def flv_tag(tag_type: int, timestamp_ms: int, payload: bytes) -> bytes: + timestamp = timestamp_ms & 0xFFFFFFFF + header = ( + bytes((tag_type,)) + + _u24(len(payload)) + + _u24(timestamp & 0xFFFFFF) + + bytes(((timestamp >> 24) & 0xFF,)) + + b"\x00\x00\x00" + ) + return header + payload + (len(header) + len(payload)).to_bytes(4, "big") + diff --git a/restream_playout/models.py b/restream_playout/models.py new file mode 100644 index 0000000..e626c26 --- /dev/null +++ b/restream_playout/models.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Iterator + + +class MediaKind(str, Enum): + AUDIO = "audio" + VIDEO = "video" + + +@dataclass(frozen=True, slots=True) +class MediaPacket: + kind: MediaKind + data: bytes + dts_ms: int + pts_ms: int + duration_ms: int + keyframe: bool = False + + +@dataclass(frozen=True, slots=True) +class CodecConfig: + avc_decoder_config: bytes + audio_specific_config: bytes + width: int + height: int + frame_rate: float + audio_sample_rate: int + audio_channels: int + + +@dataclass(frozen=True, slots=True) +class Sample: + kind: MediaKind + offset: int + size: int + dts_ms: int + pts_ms: int + duration_ms: int + keyframe: bool + + +@dataclass(frozen=True, slots=True) +class ParsedClip: + path: Path + config: CodecConfig + samples: tuple[Sample, ...] + duration_ms: int + + def packets(self) -> Iterator[MediaPacket]: + with self.path.open("rb") as source: + for sample in self.samples: + source.seek(sample.offset) + data = source.read(sample.size) + if len(data) != sample.size: + raise OSError(f"short read from {self.path}") + yield MediaPacket( + kind=sample.kind, + data=data, + dts_ms=sample.dts_ms, + pts_ms=sample.pts_ms, + duration_ms=sample.duration_ms, + keyframe=sample.keyframe, + ) + diff --git a/restream_playout/mp4.py b/restream_playout/mp4.py new file mode 100644 index 0000000..8a3fc22 --- /dev/null +++ b/restream_playout/mp4.py @@ -0,0 +1,358 @@ +from __future__ import annotations + +import struct +from dataclasses import dataclass +from pathlib import Path +from typing import Iterator + +from .exceptions import InvalidMediaError +from .models import CodecConfig, MediaKind, ParsedClip, Sample + + +@dataclass(frozen=True, slots=True) +class _Box: + type: bytes + payload_start: int + end: int + + +@dataclass(slots=True) +class _Track: + kind: MediaKind + timescale: int + config: bytes + width: int = 0 + height: int = 0 + sample_rate: int = 0 + channels: int = 0 + samples: list[Sample] | None = None + + +def _boxes(data: bytes, start: int = 0, end: int | None = None) -> Iterator[_Box]: + end = len(data) if end is None else end + cursor = start + while cursor + 8 <= end: + size = int.from_bytes(data[cursor : cursor + 4], "big") + box_type = data[cursor + 4 : cursor + 8] + header = 8 + if size == 1: + if cursor + 16 > end: + raise InvalidMediaError("truncated extended MP4 box") + size = int.from_bytes(data[cursor + 8 : cursor + 16], "big") + header = 16 + elif size == 0: + size = end - cursor + if size < header or cursor + size > end: + raise InvalidMediaError(f"invalid MP4 box {box_type!r}") + yield _Box(box_type, cursor + header, cursor + size) + cursor += size + + +def _child(data: bytes, parent: _Box, box_type: bytes) -> _Box: + for box in _boxes(data, parent.payload_start, parent.end): + if box.type == box_type: + return box + raise InvalidMediaError(f"missing MP4 box {box_type.decode('ascii', 'replace')}") + + +def _optional_child(data: bytes, parent: _Box, box_type: bytes) -> _Box | None: + return next((box for box in _boxes(data, parent.payload_start, parent.end) if box.type == box_type), None) + + +def _expand_timing(payload: bytes) -> list[int]: + count = int.from_bytes(payload[4:8], "big") + cursor = 8 + values: list[int] = [] + for _ in range(count): + sample_count, delta = struct.unpack_from(">II", payload, cursor) + values.extend([delta] * sample_count) + cursor += 8 + return values + + +def _composition_offsets(payload: bytes, sample_count: int) -> list[int]: + version = payload[0] + count = int.from_bytes(payload[4:8], "big") + cursor = 8 + values: list[int] = [] + for _ in range(count): + run, raw_offset = struct.unpack_from(">II", payload, cursor) + offset = raw_offset if version == 0 else struct.unpack(">i", raw_offset.to_bytes(4, "big"))[0] + values.extend([offset] * run) + cursor += 8 + if len(values) != sample_count: + raise InvalidMediaError("ctts sample count mismatch") + return values + + +def _descriptor_length(data: bytes, cursor: int) -> tuple[int, int]: + value = 0 + for _ in range(4): + if cursor >= len(data): + raise InvalidMediaError("truncated MPEG-4 descriptor") + byte = data[cursor] + cursor += 1 + value = (value << 7) | (byte & 0x7F) + if not byte & 0x80: + return value, cursor + raise InvalidMediaError("invalid MPEG-4 descriptor length") + + +def _audio_specific_config(esds: bytes) -> bytes: + for cursor in range(4, len(esds)): + if esds[cursor] != 0x05: + continue + try: + length, body_start = _descriptor_length(esds, cursor + 1) + except InvalidMediaError: + continue + body_end = body_start + length + if 2 <= length <= 64 and body_end <= len(esds): + config = esds[body_start:body_end] + audio_object_type = config[0] >> 3 + if 1 <= audio_object_type <= 31: + return config + raise InvalidMediaError("AAC AudioSpecificConfig not found in esds") + + +def _parse_stsd(data: bytes, stsd: _Box, kind: MediaKind) -> tuple[bytes, int, int]: + entries = list(_boxes(data, stsd.payload_start + 8, stsd.end)) + if len(entries) != 1: + raise InvalidMediaError("exactly one sample description is required") + entry = entries[0] + if kind is MediaKind.VIDEO: + if entry.type != b"avc1": + raise InvalidMediaError("video must use avc1/H.264") + payload = data[entry.payload_start : entry.end] + if len(payload) < 78: + raise InvalidMediaError("truncated avc1 sample entry") + width, height = struct.unpack_from(">HH", payload, 24) + avcc = next((box for box in _boxes(data, entry.payload_start + 78, entry.end) if box.type == b"avcC"), None) + if avcc is None: + raise InvalidMediaError("avcC configuration is required") + config = data[avcc.payload_start : avcc.end] + if len(config) < 7 or (config[4] & 3) != 3: + raise InvalidMediaError("only four-byte AVC NAL lengths are supported") + _validate_avcc(config) + return config, width, height + if entry.type != b"mp4a": + raise InvalidMediaError("audio must use mp4a/AAC") + payload = data[entry.payload_start : entry.end] + if len(payload) < 28: + raise InvalidMediaError("truncated mp4a sample entry") + channels = int.from_bytes(payload[16:18], "big") + sample_rate = int.from_bytes(payload[24:28], "big") >> 16 + esds = next((box for box in _boxes(data, entry.payload_start + 28, entry.end) if box.type == b"esds"), None) + if esds is None: + raise InvalidMediaError("esds configuration is required") + return _audio_specific_config(data[esds.payload_start : esds.end]), sample_rate, channels + + +def _validate_avcc(config: bytes) -> None: + if config[0] != 1: + raise InvalidMediaError("unsupported avcC version") + cursor = 6 + sps_count = config[5] & 0x1F + if not sps_count: + raise InvalidMediaError("avcC must contain an SPS") + for _ in range(sps_count): + if cursor + 2 > len(config): + raise InvalidMediaError("truncated avcC SPS length") + length = int.from_bytes(config[cursor : cursor + 2], "big") + cursor += 2 + length + if not length or cursor > len(config): + raise InvalidMediaError("truncated avcC SPS") + if cursor >= len(config): + raise InvalidMediaError("avcC must contain a PPS") + pps_count = config[cursor] + cursor += 1 + if not pps_count: + raise InvalidMediaError("avcC must contain a PPS") + for _ in range(pps_count): + if cursor + 2 > len(config): + raise InvalidMediaError("truncated avcC PPS length") + length = int.from_bytes(config[cursor : cursor + 2], "big") + cursor += 2 + length + if not length or cursor > len(config): + raise InvalidMediaError("truncated avcC PPS") + + +def _sample_sizes(payload: bytes) -> list[int]: + fixed_size = int.from_bytes(payload[4:8], "big") + count = int.from_bytes(payload[8:12], "big") + if fixed_size: + return [fixed_size] * count + if len(payload) < 12 + count * 4: + raise InvalidMediaError("truncated stsz") + return list(struct.unpack_from(f">{count}I", payload, 12)) + + +def _sample_offsets(data: bytes, stbl: _Box, sizes: list[int]) -> list[int]: + stsc_box = _child(data, stbl, b"stsc") + stsc_payload = data[stsc_box.payload_start : stsc_box.end] + entry_count = int.from_bytes(stsc_payload[4:8], "big") + stsc = [struct.unpack_from(">III", stsc_payload, 8 + i * 12) for i in range(entry_count)] + if not stsc: + raise InvalidMediaError("stsc must contain at least one entry") + offset_box = _optional_child(data, stbl, b"stco") or _optional_child(data, stbl, b"co64") + if offset_box is None: + raise InvalidMediaError("missing chunk offsets") + payload = data[offset_box.payload_start : offset_box.end] + count = int.from_bytes(payload[4:8], "big") + width = 4 if offset_box.type == b"stco" else 8 + offsets = [int.from_bytes(payload[8 + i * width : 8 + (i + 1) * width], "big") for i in range(count)] + result: list[int] = [] + sample_index = 0 + stsc_index = 0 + for chunk_number, chunk_offset in enumerate(offsets, 1): + while stsc_index + 1 < len(stsc) and chunk_number >= stsc[stsc_index + 1][0]: + stsc_index += 1 + samples_per_chunk = stsc[stsc_index][1] + offset = chunk_offset + for _ in range(samples_per_chunk): + if sample_index >= len(sizes): + raise InvalidMediaError("chunk table contains too many samples") + result.append(offset) + offset += sizes[sample_index] + sample_index += 1 + if sample_index != len(sizes): + raise InvalidMediaError("chunk table does not cover all samples") + return result + + +def _parse_track(data: bytes, trak: _Box) -> _Track: + if _optional_child(data, trak, b"edts") is not None: + raise InvalidMediaError("MP4 edit lists are not supported; remux the source with a zero-based timeline") + mdia = _child(data, trak, b"mdia") + hdlr = data[_child(data, mdia, b"hdlr").payload_start : _child(data, mdia, b"hdlr").end] + handler = hdlr[8:12] + if handler == b"vide": + kind = MediaKind.VIDEO + elif handler == b"soun": + kind = MediaKind.AUDIO + else: + raise InvalidMediaError("only video and audio tracks are supported") + mdhd = data[_child(data, mdia, b"mdhd").payload_start : _child(data, mdia, b"mdhd").end] + timescale_offset = 20 if mdhd[0] == 1 else 12 + timescale = int.from_bytes(mdhd[timescale_offset : timescale_offset + 4], "big") + if not timescale: + raise InvalidMediaError("track timescale must be non-zero") + stbl = _child(data, _child(data, mdia, b"minf"), b"stbl") + config, first, second = _parse_stsd(data, _child(data, stbl, b"stsd"), kind) + sizes_box = _child(data, stbl, b"stsz") + sizes = _sample_sizes(data[sizes_box.payload_start : sizes_box.end]) + offsets = _sample_offsets(data, stbl, sizes) + stts_box = _child(data, stbl, b"stts") + durations = _expand_timing(data[stts_box.payload_start : stts_box.end]) + if len(durations) != len(sizes): + raise InvalidMediaError("stts sample count mismatch") + ctts_box = _optional_child(data, stbl, b"ctts") + composition = ( + _composition_offsets(data[ctts_box.payload_start : ctts_box.end], len(sizes)) + if ctts_box + else [0] * len(sizes) + ) + stss_box = _optional_child(data, stbl, b"stss") + sync_samples = None + if stss_box: + payload = data[stss_box.payload_start : stss_box.end] + count = int.from_bytes(payload[4:8], "big") + sync_samples = set(struct.unpack_from(f">{count}I", payload, 8)) + dts = 0 + samples: list[Sample] = [] + for index, (offset, size, duration, cts) in enumerate(zip(offsets, sizes, durations, composition), 1): + samples.append( + Sample( + kind=kind, + offset=offset, + size=size, + dts_ms=round(dts * 1000 / timescale), + pts_ms=round((dts + cts) * 1000 / timescale), + duration_ms=max(1, round(duration * 1000 / timescale)), + keyframe=kind is MediaKind.AUDIO or sync_samples is None or index in sync_samples, + ) + ) + dts += duration + return _Track( + kind=kind, + timescale=timescale, + config=config, + width=first if kind is MediaKind.VIDEO else 0, + height=second if kind is MediaKind.VIDEO else 0, + sample_rate=first if kind is MediaKind.AUDIO else 0, + channels=second if kind is MediaKind.AUDIO else 0, + samples=samples, + ) + + +def parse_mp4(path: str | Path) -> ParsedClip: + media_path = Path(path).expanduser().resolve() + if not media_path.is_file(): + raise InvalidMediaError(f"media file does not exist: {media_path}") + file_size = media_path.stat().st_size + with media_path.open("rb") as source: + moov_data = None + while True: + box_start = source.tell() + header = source.read(8) + if not header: + break + if len(header) != 8: + raise InvalidMediaError("truncated top-level MP4 box") + size = int.from_bytes(header[:4], "big") + box_type = header[4:] + header_size = 8 + if size == 1: + extended = source.read(8) + if len(extended) != 8: + raise InvalidMediaError("truncated extended MP4 box") + size = int.from_bytes(extended, "big") + header_size = 16 + elif size == 0: + size = file_size - box_start + if size < header_size: + raise InvalidMediaError("invalid top-level MP4 box size") + if box_start + size > file_size: + raise InvalidMediaError("top-level MP4 box extends beyond the file") + if box_type == b"moov": + moov_data = source.read(size - header_size) + if len(moov_data) != size - header_size: + raise InvalidMediaError("truncated moov box") + break + source.seek(size - header_size, 1) + if moov_data is None: + raise InvalidMediaError("MP4 moov box not found") + root = _Box(b"moov", 0, len(moov_data)) + if _optional_child(moov_data, root, b"mvex") is not None: + raise InvalidMediaError("fragmented MP4 is not supported") + tracks = [_parse_track(moov_data, box) for box in _boxes(moov_data) if box.type == b"trak"] + videos = [track for track in tracks if track.kind is MediaKind.VIDEO] + audios = [track for track in tracks if track.kind is MediaKind.AUDIO] + if len(videos) != 1 or len(audios) != 1: + raise InvalidMediaError("MP4 must contain exactly one H.264 video track and one AAC audio track") + video, audio = videos[0], audios[0] + if not video.samples or not audio.samples: + raise InvalidMediaError("audio and video tracks must contain samples") + if video.width <= 0 or video.height <= 0: + raise InvalidMediaError("video dimensions must be positive") + if audio.sample_rate <= 0 or audio.channels not in (1, 2): + raise InvalidMediaError("AAC audio must be mono or stereo with a positive sample rate") + for sample in (*video.samples, *audio.samples): + if sample.size <= 0 or sample.offset < 0 or sample.offset + sample.size > file_size: + raise InvalidMediaError("sample table points outside the MP4 file") + if not -(1 << 23) <= sample.pts_ms - sample.dts_ms < (1 << 23): + raise InvalidMediaError("composition offset cannot be represented by FLV") + samples = sorted((*(video.samples or []), *(audio.samples or [])), key=lambda sample: (sample.dts_ms, sample.kind.value)) + video_duration = max((sample.dts_ms + sample.duration_ms for sample in video.samples or []), default=0) + audio_duration = max((sample.dts_ms + sample.duration_ms for sample in audio.samples or []), default=0) + frame_rate = len(video.samples or []) * 1000 / video_duration if video_duration else 0.0 + config = CodecConfig( + avc_decoder_config=video.config, + audio_specific_config=audio.config, + width=video.width, + height=video.height, + frame_rate=frame_rate, + audio_sample_rate=audio.sample_rate, + audio_channels=audio.channels, + ) + return ParsedClip(media_path, config, tuple(samples), max(video_duration, audio_duration)) diff --git a/restream_playout/queue.py b/restream_playout/queue.py new file mode 100644 index 0000000..2e196a0 --- /dev/null +++ b/restream_playout/queue.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import threading +from collections import deque +from collections.abc import Iterable +from pathlib import Path + + +class VideoQueue: + def __init__(self) -> None: + self._items: deque[Path] = deque() + self._condition = threading.Condition() + + def enqueue(self, path: str | Path) -> None: + with self._condition: + self._items.append(Path(path)) + self._condition.notify_all() + + def enqueue_many(self, paths: Iterable[str | Path]) -> None: + with self._condition: + self._items.extend(Path(path) for path in paths) + self._condition.notify_all() + + def clear(self) -> None: + with self._condition: + self._items.clear() + + def size(self) -> int: + with self._condition: + return len(self._items) + + def peek(self) -> Path | None: + with self._condition: + return self._items[0] if self._items else None + + def pop(self) -> Path | None: + with self._condition: + return self._items.popleft() if self._items else None + + def wait(self, timeout: float | None = None) -> bool: + with self._condition: + if self._items: + return True + self._condition.wait(timeout) + return bool(self._items) + + def wake(self) -> None: + with self._condition: + self._condition.notify_all() + diff --git a/restream_playout/rtmp.py b/restream_playout/rtmp.py new file mode 100644 index 0000000..abd8680 --- /dev/null +++ b/restream_playout/rtmp.py @@ -0,0 +1,425 @@ +from __future__ import annotations + +import asyncio +import os +import struct +import time +from dataclasses import dataclass, field +from typing import Any +from urllib.parse import urlsplit + +from .exceptions import RTMPError, StreamDisconnected +from .flv import amf0_encode + +RTMP_SET_CHUNK_SIZE = 1 +RTMP_ABORT = 2 +RTMP_ACKNOWLEDGEMENT = 3 +RTMP_USER_CONTROL = 4 +RTMP_WINDOW_ACK_SIZE = 5 +RTMP_SET_PEER_BANDWIDTH = 6 +RTMP_AUDIO = 8 +RTMP_VIDEO = 9 +RTMP_DATA_AMF0 = 18 +RTMP_COMMAND_AMF0 = 20 + + +@dataclass(frozen=True, slots=True) +class RTMPUrl: + host: str + port: int + app: str + stream: str + tc_url: str + + @classmethod + def parse(cls, value: str) -> "RTMPUrl": + parsed = urlsplit(value) + if parsed.scheme.lower() != "rtmp": + raise RTMPError("only rtmp:// URLs are supported") + segments = [segment for segment in parsed.path.split("/") if segment] + if not parsed.hostname or len(segments) < 2: + raise RTMPError("RTMP URL must include a host, application, and stream key") + app = segments[0] + stream = "/".join(segments[1:]) + if parsed.query: + stream += "?" + parsed.query + port = parsed.port or 1935 + tc_url = f"rtmp://{parsed.hostname}:{port}/{app}" + return cls(parsed.hostname, port, app, stream, tc_url) + + +@dataclass(frozen=True, slots=True) +class RTMPMessage: + timestamp: int + message_type: int + stream_id: int + payload: bytes + + +@dataclass(slots=True) +class _ChunkState: + timestamp: int = 0 + timestamp_delta: int = 0 + message_length: int = 0 + message_type: int = 0 + stream_id: int = 0 + extended: bool = False + payload: bytearray = field(default_factory=bytearray) + + +class _AMFReader: + def __init__(self, data: bytes) -> None: + self.data = data + self.cursor = 0 + + def _take(self, length: int) -> bytes: + end = self.cursor + length + if end > len(self.data): + raise RTMPError("truncated AMF0 value") + result = self.data[self.cursor:end] + self.cursor = end + return result + + def value(self) -> Any: + marker = self._take(1)[0] + if marker == 0: + return struct.unpack(">d", self._take(8))[0] + if marker == 1: + return bool(self._take(1)[0]) + if marker == 2: + return self._take(int.from_bytes(self._take(2), "big")).decode("utf-8", "replace") + if marker in (5, 6): + return None + if marker == 3: + return self._object() + if marker == 8: + self._take(4) + return self._object() + if marker == 10: + return [self.value() for _ in range(int.from_bytes(self._take(4), "big"))] + if marker == 12: + return self._take(int.from_bytes(self._take(4), "big")).decode("utf-8", "replace") + raise RTMPError(f"unsupported AMF0 marker {marker}") + + def _object(self) -> dict[str, Any]: + result: dict[str, Any] = {} + while True: + length = int.from_bytes(self._take(2), "big") + if length == 0 and self._take(1) == b"\x09": + return result + key = self._take(length).decode("utf-8", "replace") + result[key] = self.value() + + +def amf0_decode_all(data: bytes) -> list[Any]: + reader = _AMFReader(data) + values = [] + while reader.cursor < len(data): + values.append(reader.value()) + return values + + +class RTMPClient: + def __init__(self, url: str, *, connect_timeout: float = 15.0, chunk_size: int = 4096) -> None: + self.url = RTMPUrl.parse(url) + self.connect_timeout = connect_timeout + self.out_chunk_size = chunk_size + self.in_chunk_size = 128 + self.stream_id = 0 + self._reader: asyncio.StreamReader | None = None + self._writer: asyncio.StreamWriter | None = None + self._write_lock = asyncio.Lock() + self._chunks: dict[int, _ChunkState] = {} + self._bytes_read = 0 + self._last_ack = 0 + self._ack_window = 0 + self.last_pong = time.monotonic() + + @property + def connected(self) -> bool: + return self._writer is not None and not self._writer.is_closing() + + async def connect(self) -> None: + try: + self._reader, self._writer = await asyncio.wait_for( + asyncio.open_connection(self.url.host, self.url.port), + timeout=self.connect_timeout, + ) + await self._handshake() + await self.send_message(RTMP_SET_CHUNK_SIZE, 0, self.out_chunk_size.to_bytes(4, "big"), chunk_stream_id=2) + connect_object = { + "app": self.url.app, + "type": "nonprivate", + "flashVer": "FMLE/3.0 (compatible; restream_playout)", + "tcUrl": self.url.tc_url, + "fpad": False, + "capabilities": 15, + "audioCodecs": 0x0FFF, + "videoCodecs": 0x00FF, + "videoFunction": 1, + "objectEncoding": 0, + } + await self._command("connect", 1, connect_object) + await self._wait_for_result(1) + await self._command("createStream", 2, None) + result = await self._wait_for_result(2) + self.stream_id = int(result[-1]) + await self._command("publish", 0, None, self.url.stream, "live", stream_id=self.stream_id) + await self._wait_for_publish_start() + except Exception: + await self.close() + raise + + async def _handshake(self) -> None: + if self._reader is None or self._writer is None: + raise StreamDisconnected("socket is not open") + c1 = int(time.time()).to_bytes(4, "big") + b"\x00\x00\x00\x00" + os.urandom(1528) + self._writer.write(b"\x03" + c1) + await self._writer.drain() + response = await asyncio.wait_for(self._reader.readexactly(3073), timeout=self.connect_timeout) + if response[0] != 3: + raise RTMPError(f"unsupported RTMP version {response[0]}") + self._writer.write(response[1:1537]) + await self._writer.drain() + + async def _command( + self, + name: str, + transaction_id: int, + command_object: Any, + *arguments: Any, + stream_id: int = 0, + ) -> None: + payload = b"".join(amf0_encode(value) for value in (name, transaction_id, command_object, *arguments)) + await self.send_message(RTMP_COMMAND_AMF0, 0, payload, stream_id=stream_id, chunk_stream_id=3) + + async def _wait_for_result(self, transaction_id: int) -> list[Any]: + while True: + message = await asyncio.wait_for(self.read_message(), timeout=self.connect_timeout) + if message.message_type != RTMP_COMMAND_AMF0: + continue + values = amf0_decode_all(message.payload) + if len(values) >= 2 and values[0] == "_error" and int(values[1]) == transaction_id: + raise RTMPError(f"RTMP command failed: {values}") + if len(values) >= 2 and values[0] == "_result" and int(values[1]) == transaction_id: + return values + + async def _wait_for_publish_start(self) -> None: + while True: + message = await asyncio.wait_for(self.read_message(), timeout=self.connect_timeout) + if message.message_type != RTMP_COMMAND_AMF0: + continue + values = amf0_decode_all(message.payload) + if values and values[0] == "onStatus": + status = next((value for value in values if isinstance(value, dict) and "code" in value), {}) + code = status.get("code", "") + if code == "NetStream.Publish.Start": + return + if code.startswith("NetStream.Publish.") and code != "NetStream.Publish.Start": + raise RTMPError(f"publish rejected: {code}") + + async def send_message( + self, + message_type: int, + timestamp: int, + payload: bytes, + *, + stream_id: int | None = None, + chunk_stream_id: int = 5, + ) -> None: + writer = self._writer + if writer is None or writer.is_closing(): + raise StreamDisconnected("RTMP socket is closed") + if len(payload) > 0xFFFFFF: + raise RTMPError("an RTMP message cannot exceed 16,777,215 bytes") + stream_id = self.stream_id if stream_id is None else stream_id + wire_timestamp = max(0, timestamp) & 0xFFFFFFFF + timestamp_field = min(wire_timestamp, 0xFFFFFF) + basic = self._encode_basic_header(0, chunk_stream_id) + header = ( + basic + + timestamp_field.to_bytes(3, "big") + + len(payload).to_bytes(3, "big") + + bytes((message_type,)) + + stream_id.to_bytes(4, "little") + ) + extended = wire_timestamp.to_bytes(4, "big") if wire_timestamp >= 0xFFFFFF else b"" + chunks = bytearray() + for offset in range(0, len(payload), self.out_chunk_size): + if offset == 0: + chunks += header + extended + else: + chunks += self._encode_basic_header(3, chunk_stream_id) + extended + chunks += payload[offset : offset + self.out_chunk_size] + if not payload: + chunks += header + extended + try: + async with self._write_lock: + writer.write(chunks) + await writer.drain() + except (ConnectionError, asyncio.IncompleteReadError, OSError) as exc: + raise StreamDisconnected("RTMP write failed") from exc + + @staticmethod + def _encode_basic_header(fmt: int, chunk_stream_id: int) -> bytes: + if not 0 <= fmt <= 3 or not 2 <= chunk_stream_id <= 65599: + raise RTMPError("invalid RTMP chunk header") + if chunk_stream_id <= 63: + return bytes(((fmt << 6) | chunk_stream_id,)) + value = chunk_stream_id - 64 + if value <= 0xFF: + return bytes((fmt << 6, value)) + return bytes(((fmt << 6) | 1, value & 0xFF, value >> 8)) + + async def _readexactly(self, length: int) -> bytes: + if self._reader is None: + raise StreamDisconnected("RTMP socket is closed") + try: + data = await self._reader.readexactly(length) + except (ConnectionError, asyncio.IncompleteReadError, OSError) as exc: + raise StreamDisconnected("RTMP read failed") from exc + self._bytes_read += length + return data + + async def _basic_header(self) -> tuple[int, int]: + first = (await self._readexactly(1))[0] + fmt, chunk_stream_id = first >> 6, first & 0x3F + if chunk_stream_id == 0: + chunk_stream_id = 64 + (await self._readexactly(1))[0] + elif chunk_stream_id == 1: + low, high = await self._readexactly(2) + chunk_stream_id = 64 + low + high * 256 + return fmt, chunk_stream_id + + async def _read_raw_message(self) -> RTMPMessage: + while True: + fmt, csid = await self._basic_header() + state = self._chunks.setdefault(csid, _ChunkState()) + continuing = bool(state.payload) and len(state.payload) < state.message_length + if fmt == 0: + header = await self._readexactly(11) + raw_timestamp = int.from_bytes(header[0:3], "big") + state.timestamp = raw_timestamp + state.timestamp_delta = 0 + state.message_length = int.from_bytes(header[3:6], "big") + state.message_type = header[6] + state.stream_id = int.from_bytes(header[7:11], "little") + state.extended = raw_timestamp == 0xFFFFFF + state.payload = bytearray() + if state.extended: + state.timestamp = int.from_bytes(await self._readexactly(4), "big") + elif fmt == 1: + header = await self._readexactly(7) + raw_delta = int.from_bytes(header[0:3], "big") + state.timestamp_delta = raw_delta + state.message_length = int.from_bytes(header[3:6], "big") + state.message_type = header[6] + state.extended = raw_delta == 0xFFFFFF + if state.extended: + state.timestamp_delta = int.from_bytes(await self._readexactly(4), "big") + state.timestamp += state.timestamp_delta + state.payload = bytearray() + elif fmt == 2: + raw_delta = int.from_bytes(await self._readexactly(3), "big") + state.timestamp_delta = raw_delta + state.extended = raw_delta == 0xFFFFFF + if state.extended: + state.timestamp_delta = int.from_bytes(await self._readexactly(4), "big") + state.timestamp += state.timestamp_delta + state.payload = bytearray() + elif fmt == 3: + if not state.message_length: + raise RTMPError("RTMP type-3 chunk has no preceding header") + if state.extended: + await self._readexactly(4) + if not continuing: + state.timestamp += state.timestamp_delta + state.payload = bytearray() + remaining = state.message_length - len(state.payload) + state.payload += await self._readexactly(min(self.in_chunk_size, remaining)) + if len(state.payload) == state.message_length: + return RTMPMessage(state.timestamp, state.message_type, state.stream_id, bytes(state.payload)) + + async def read_message(self) -> RTMPMessage: + while True: + message = await self._read_raw_message() + if message.message_type == RTMP_SET_CHUNK_SIZE and len(message.payload) >= 4: + self.in_chunk_size = int.from_bytes(message.payload[:4], "big") & 0x7FFFFFFF + await self._ack_if_needed() + continue + if message.message_type == RTMP_ABORT and len(message.payload) >= 4: + self._chunks.pop(int.from_bytes(message.payload[:4], "big"), None) + await self._ack_if_needed() + continue + if message.message_type == RTMP_WINDOW_ACK_SIZE and len(message.payload) >= 4: + self._ack_window = int.from_bytes(message.payload[:4], "big") + await self._ack_if_needed() + continue + if message.message_type == RTMP_SET_PEER_BANDWIDTH and len(message.payload) >= 4: + self._ack_window = int.from_bytes(message.payload[:4], "big") + await self._ack_if_needed() + continue + if message.message_type == RTMP_USER_CONTROL and len(message.payload) >= 6: + event = int.from_bytes(message.payload[:2], "big") + if event == 6: + await self.send_message( + RTMP_USER_CONTROL, + 0, + b"\x00\x07" + message.payload[2:6], + stream_id=0, + chunk_stream_id=2, + ) + elif event == 7: + self.last_pong = time.monotonic() + await self._ack_if_needed() + continue + await self._ack_if_needed() + return message + + async def _ack_if_needed(self) -> None: + if self._ack_window and self._bytes_read - self._last_ack >= self._ack_window: + self._last_ack = self._bytes_read + await self.send_message( + RTMP_ACKNOWLEDGEMENT, + 0, + (self._bytes_read & 0xFFFFFFFF).to_bytes(4, "big"), + stream_id=0, + chunk_stream_id=2, + ) + + async def ping(self) -> None: + timestamp = int(time.monotonic() * 1000) & 0xFFFFFFFF + await self.send_message( + RTMP_USER_CONTROL, + 0, + b"\x00\x06" + timestamp.to_bytes(4, "big"), + stream_id=0, + chunk_stream_id=2, + ) + + async def unpublish(self) -> None: + if not self.connected or not self.stream_id: + return + await self._command("FCUnpublish", 0, None, self.url.stream) + await self._command("deleteStream", 0, None, self.stream_id) + + async def monitor(self) -> None: + while True: + message = await self.read_message() + if message.message_type == RTMP_COMMAND_AMF0: + values = amf0_decode_all(message.payload) + if values and values[0] == "onStatus": + status = next((value for value in values if isinstance(value, dict) and "code" in value), {}) + code = status.get("code", "") + if code in {"NetStream.Unpublish.Success", "NetStream.Publish.BadName", "NetStream.Failed"}: + raise StreamDisconnected(f"RTMP publisher stopped: {code}") + + async def close(self) -> None: + writer, self._writer = self._writer, None + self._reader = None + if writer is not None: + writer.close() + try: + await writer.wait_closed() + except (ConnectionError, OSError): + pass diff --git a/restream_playout/session.py b/restream_playout/session.py new file mode 100644 index 0000000..bc46da8 --- /dev/null +++ b/restream_playout/session.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import asyncio +import logging +import random +import time +from collections.abc import Awaitable, Callable + +from .exceptions import StreamDisconnected +from .rtmp import RTMPClient, RTMPUrl + +StateCallback = Callable[[bool], Awaitable[None]] + + +class StreamSession: + def __init__( + self, + rtmp_url: str, + *, + stop_event: asyncio.Event, + state_callback: StateCallback, + logger: logging.Logger | None = None, + reconnect_min_delay: float = 1.0, + reconnect_max_delay: float = 30.0, + heartbeat_interval: float = 10.0, + ) -> None: + RTMPUrl.parse(rtmp_url) + self.rtmp_url = rtmp_url + self.stop_event = stop_event + self.state_callback = state_callback + self.logger = logger or logging.getLogger(__name__) + self.reconnect_min_delay = reconnect_min_delay + self.reconnect_max_delay = reconnect_max_delay + self.heartbeat_interval = heartbeat_interval + self.generation = 0 + self._client: RTMPClient | None = None + self._connect_lock = asyncio.Lock() + self._monitor_task: asyncio.Task[None] | None = None + self._heartbeat_task: asyncio.Task[None] | None = None + + @property + def connected(self) -> bool: + return self._client is not None and self._client.connected and self._monitor_task is not None and not self._monitor_task.done() + + async def ensure_connected(self) -> int: + async with self._connect_lock: + if self.connected: + return self.generation + delay = self.reconnect_min_delay + while not self.stop_event.is_set(): + client = RTMPClient(self.rtmp_url) + try: + self.logger.info("connecting to RTMP endpoint") + await client.connect() + self._client = client + self.generation += 1 + self._monitor_task = asyncio.create_task(self._monitor(client), name="rtmp-monitor") + self._heartbeat_task = asyncio.create_task(self._heartbeat(client), name="rtmp-heartbeat") + self.logger.info("RTMP stream connected", extra={"generation": self.generation}) + await self.state_callback(True) + return self.generation + except Exception as exc: + await client.close() + self.logger.warning("RTMP connection failed; retrying", extra={"delay": delay, "error": str(exc)}) + try: + await asyncio.wait_for(self.stop_event.wait(), delay + random.random() * delay * 0.2) + except TimeoutError: + pass + delay = min(self.reconnect_max_delay, delay * 2) + raise StreamDisconnected("session stopped") + + async def send_message(self, message_type: int, timestamp: int, payload: bytes, *, chunk_stream_id: int = 5) -> None: + client = self._client + if not self.connected or client is None: + raise StreamDisconnected("RTMP session is disconnected") + try: + await client.send_message(message_type, timestamp, payload, chunk_stream_id=chunk_stream_id) + except Exception as exc: + await self._mark_disconnected(client) + raise StreamDisconnected("RTMP send failed") from exc + + async def _monitor(self, client: RTMPClient) -> None: + try: + await client.monitor() + except asyncio.CancelledError: + raise + except Exception as exc: + self.logger.warning("RTMP monitor detected disconnect", extra={"error": str(exc)}) + finally: + await self._mark_disconnected(client) + + async def _heartbeat(self, client: RTMPClient) -> None: + try: + while True: + await asyncio.sleep(self.heartbeat_interval) + if time.monotonic() - client.last_pong > self.heartbeat_interval * 3: + raise StreamDisconnected("RTMP ping response timed out") + await client.ping() + except asyncio.CancelledError: + raise + except Exception as exc: + self.logger.warning("RTMP heartbeat failed", extra={"error": str(exc)}) + await self._mark_disconnected(client) + + async def _mark_disconnected(self, client: RTMPClient) -> None: + if self._client is not client: + return + self._client = None + current = asyncio.current_task() + for task in (self._monitor_task, self._heartbeat_task): + if task is not None and task is not current: + task.cancel() + self._monitor_task = None + self._heartbeat_task = None + await client.close() + self.logger.info("RTMP stream disconnected") + await self.state_callback(False) + + async def close(self) -> None: + client = self._client + if client is not None: + try: + await client.unpublish() + except Exception: + pass + await self._mark_disconnected(client) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/mp4_builder.py b/tests/mp4_builder.py new file mode 100644 index 0000000..dc98b50 --- /dev/null +++ b/tests/mp4_builder.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import struct +from pathlib import Path + +from restream_playout.filler import AAC_SILENCE, AVCC, AVC_SAMPLE + + +def box(kind: bytes, payload: bytes) -> bytes: + return (len(payload) + 8).to_bytes(4, "big") + kind + payload + + +def full_box(kind: bytes, payload: bytes, version: int = 0, flags: int = 0) -> bytes: + return box(kind, bytes((version,)) + flags.to_bytes(3, "big") + payload) + + +def descriptor(tag: int, payload: bytes) -> bytes: + if len(payload) >= 128: + raise ValueError("test descriptor is too large") + return bytes((tag, len(payload))) + payload + + +def stbl(sample_entry: bytes, sample: bytes, chunk_offset: int, duration: int, *, sync: bool) -> bytes: + stsd = full_box(b"stsd", b"\x00\x00\x00\x01" + sample_entry) + stts = full_box(b"stts", b"\x00\x00\x00\x01" + struct.pack(">II", 1, duration)) + stsc = full_box(b"stsc", b"\x00\x00\x00\x01" + struct.pack(">III", 1, 1, 1)) + stsz = full_box(b"stsz", b"\x00\x00\x00\x00\x00\x00\x00\x01" + len(sample).to_bytes(4, "big")) + stco = full_box(b"stco", b"\x00\x00\x00\x01" + chunk_offset.to_bytes(4, "big")) + stss = full_box(b"stss", b"\x00\x00\x00\x01\x00\x00\x00\x01") if sync else b"" + return box(b"stbl", stsd + stts + stsc + stsz + stco + stss) + + +def video_entry() -> bytes: + payload = ( + b"\x00" * 6 + + b"\x00\x01" + + b"\x00" * 16 + + struct.pack(">HH", 16, 16) + + b"\x00\x48\x00\x00" * 2 + + b"\x00\x00\x00\x00" + + b"\x00\x01" + + b"\x00" * 32 + + b"\x00\x18\xff\xff" + + box(b"avcC", AVCC) + ) + return box(b"avc1", payload) + + +def audio_entry() -> bytes: + asc = b"\x12\x10" + decoder_config = b"\x40\x15\x00\x00\x00\x00\x00\x00\x00\x00\x00" + descriptor(0x05, asc) + es_descriptor = b"\x00\x01\x00" + descriptor(0x04, decoder_config) + descriptor(0x06, b"\x02") + esds = full_box(b"esds", descriptor(0x03, es_descriptor)) + payload = ( + b"\x00" * 6 + + b"\x00\x01" + + b"\x00" * 8 + + b"\x00\x02" + + b"\x00\x10" + + b"\x00\x00\x00\x00" + + (44100 << 16).to_bytes(4, "big") + + esds + ) + return box(b"mp4a", payload) + + +def track(kind: bytes, timescale: int, table: bytes) -> bytes: + mdhd = full_box( + b"mdhd", + b"\x00" * 8 + timescale.to_bytes(4, "big") + timescale.to_bytes(4, "big") + b"\x00\x00\x00\x00", + ) + hdlr = full_box(b"hdlr", b"\x00\x00\x00\x00" + kind + b"\x00" * 12) + minf = box(b"minf", table) + return box(b"trak", box(b"mdia", mdhd + hdlr + minf)) + + +def write_test_mp4(path: Path) -> None: + ftyp = box(b"ftyp", b"isom\x00\x00\x02\x00isomavc1") + video_offset = len(ftyp) + 8 + audio_offset = video_offset + len(AVC_SAMPLE) + mdat = box(b"mdat", AVC_SAMPLE + AAC_SILENCE) + video_table = stbl(video_entry(), AVC_SAMPLE, video_offset, 3000, sync=True) + audio_table = stbl(audio_entry(), AAC_SILENCE, audio_offset, 1024, sync=False) + moov = box(b"moov", track(b"vide", 90000, video_table) + track(b"soun", 44100, audio_table)) + path.write_bytes(ftyp + mdat + moov) + diff --git a/tests/test_engine.py b/tests/test_engine.py new file mode 100644 index 0000000..1ce12cd --- /dev/null +++ b/tests/test_engine.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import asyncio +import tempfile +import threading +import unittest +from pathlib import Path +from unittest.mock import patch + +from restream_playout import PlayoutEngine +from restream_playout.rtmp import RTMP_AUDIO, RTMP_DATA_AMF0, RTMP_VIDEO + +from .mp4_builder import write_test_mp4 + + +class _FakeSession: + messages: list[tuple[int, int, bytes]] = [] + + def __init__(self, rtmp_url: str, *, stop_event: asyncio.Event, state_callback: object, **kwargs: object) -> None: + self.stop_event = stop_event + self.state_callback = state_callback + self.generation = 1 + self._announced = False + + async def ensure_connected(self) -> int: + if not self._announced: + self._announced = True + await self.state_callback(True) + return self.generation + + async def send_message( + self, + message_type: int, + timestamp: int, + payload: bytes, + *, + chunk_stream_id: int = 5, + ) -> None: + self.messages.append((message_type, timestamp, payload)) + + async def close(self) -> None: + if self._announced: + await self.state_callback(False) + + +class EngineTests(unittest.TestCase): + def test_dynamic_insertion_switches_from_filler_and_finishes_clip(self) -> None: + with tempfile.TemporaryDirectory() as directory: + clip_path = Path(directory, "dynamic.mp4") + write_test_mp4(clip_path) + started = threading.Event() + finished = threading.Event() + connected = threading.Event() + _FakeSession.messages = [] + + with patch("restream_playout.engine.StreamSession", _FakeSession): + engine = PlayoutEngine("rtmp://localhost/live/key") + engine.on_clip_started(lambda path: started.set()) + engine.on_clip_finished(lambda path: finished.set()) + engine.on_stream_connected(lambda: connected.set()) + engine.start() + self.assertTrue(connected.wait(2)) + engine.enqueue(clip_path) + self.assertTrue(started.wait(2)) + self.assertTrue(finished.wait(2)) + engine.stop() + + message_types = {message_type for message_type, _, _ in _FakeSession.messages} + self.assertEqual(message_types, {RTMP_AUDIO, RTMP_DATA_AMF0, RTMP_VIDEO}) + timestamps = [timestamp for _, timestamp, _ in _FakeSession.messages] + self.assertEqual(timestamps, sorted(timestamps)) + diff --git a/tests/test_filler_flv.py b/tests/test_filler_flv.py new file mode 100644 index 0000000..d59291f --- /dev/null +++ b/tests/test_filler_flv.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import unittest + +from restream_playout.filler import FillerSource, IDR, PPS, SPS +from restream_playout.flv import aac_sequence_header, avc_sequence_header, flv_tag, media_payload, metadata_payload +from restream_playout.models import MediaKind + + +class FillerAndFLVTests(unittest.TestCase): + def test_filler_has_valid_avcc_and_monotonic_packets(self) -> None: + filler = FillerSource() + config = filler.config + self.assertEqual(config.avc_decoder_config[0], 1) + self.assertIn(SPS, config.avc_decoder_config) + self.assertIn(PPS, config.avc_decoder_config) + self.assertEqual(IDR[0] & 0x1F, 5) + packets = [filler.next_packet() for _ in range(20)] + self.assertEqual([packet.dts_ms for packet in packets], sorted(packet.dts_ms for packet in packets)) + self.assertIn(MediaKind.VIDEO, {packet.kind for packet in packets}) + self.assertIn(MediaKind.AUDIO, {packet.kind for packet in packets}) + + def test_flv_payloads_and_tag(self) -> None: + filler = FillerSource() + config = filler.config + video = next(packet for packet in iter(filler.next_packet, None) if packet.kind is MediaKind.VIDEO) + self.assertEqual(avc_sequence_header(config)[:5], b"\x17\x00\x00\x00\x00") + self.assertEqual(aac_sequence_header(config)[:2], b"\xaf\x00") + self.assertEqual(media_payload(video)[:2], b"\x17\x01") + metadata = metadata_payload(config) + self.assertIn(b"onMetaData", metadata) + tag = flv_tag(9, 1234, media_payload(video)) + self.assertEqual(tag[0], 9) + self.assertEqual(int.from_bytes(tag[-4:], "big"), len(tag) - 4) + diff --git a/tests/test_mp4.py b/tests/test_mp4.py new file mode 100644 index 0000000..a5a7914 --- /dev/null +++ b/tests/test_mp4.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +import tempfile +import unittest +from pathlib import Path + +from restream_playout.exceptions import InvalidMediaError +from restream_playout.models import MediaKind +from restream_playout.mp4 import parse_mp4 + +from .mp4_builder import write_test_mp4 + + +class MP4Tests(unittest.TestCase): + def test_parse_h264_aac_mp4(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory, "clip.mp4") + write_test_mp4(path) + clip = parse_mp4(path) + self.assertEqual((clip.config.width, clip.config.height), (16, 16)) + self.assertEqual(clip.config.audio_sample_rate, 44100) + packets = list(clip.packets()) + self.assertEqual({packet.kind for packet in packets}, {MediaKind.VIDEO, MediaKind.AUDIO}) + self.assertTrue(next(packet for packet in packets if packet.kind is MediaKind.VIDEO).keyframe) + + def test_rejects_missing_file(self) -> None: + with self.assertRaises(InvalidMediaError): + parse_mp4("does-not-exist.mp4") + diff --git a/tests/test_purity.py b/tests/test_purity.py new file mode 100644 index 0000000..42a89e5 --- /dev/null +++ b/tests/test_purity.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +import ast +import unittest +from pathlib import Path + + +class PurityTests(unittest.TestCase): + def test_package_has_no_process_launching_calls(self) -> None: + package = Path(__file__).parents[1] / "restream_playout" + forbidden_imports = {"subprocess"} + forbidden_calls = {"os.system", "os.popen", "subprocess.run", "subprocess.Popen", "subprocess.call"} + violations: list[str] = [] + for path in package.glob("*.py"): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for name in node.names: + if name.name in forbidden_imports: + violations.append(f"{path.name}: imports {name.name}") + elif isinstance(node, ast.ImportFrom) and node.module in forbidden_imports: + violations.append(f"{path.name}: imports from {node.module}") + elif isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + if isinstance(node.func.value, ast.Name): + call = f"{node.func.value.id}.{node.func.attr}" + if call in forbidden_calls: + violations.append(f"{path.name}: calls {call}") + self.assertEqual(violations, []) + diff --git a/tests/test_queue.py b/tests/test_queue.py new file mode 100644 index 0000000..21648bf --- /dev/null +++ b/tests/test_queue.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +import threading +import time +import unittest + +from restream_playout.queue import VideoQueue + + +class QueueTests(unittest.TestCase): + def test_fifo_and_wait(self) -> None: + queue = VideoQueue() + + def producer() -> None: + time.sleep(0.02) + queue.enqueue_many(["a.mp4", "b.mp4"]) + + thread = threading.Thread(target=producer) + thread.start() + self.assertTrue(queue.wait(1)) + self.assertEqual(str(queue.peek()), "a.mp4") + self.assertEqual(str(queue.pop()), "a.mp4") + self.assertEqual(queue.size(), 1) + queue.clear() + self.assertEqual(queue.size(), 0) + thread.join() + diff --git a/tests/test_rtmp.py b/tests/test_rtmp.py new file mode 100644 index 0000000..be8fb3e --- /dev/null +++ b/tests/test_rtmp.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import asyncio +import unittest + +from restream_playout.flv import amf0_encode +from restream_playout.rtmp import RTMP_COMMAND_AMF0, RTMP_DATA_AMF0, RTMPClient, amf0_decode_all + + +def command(*values: object) -> bytes: + return b"".join(amf0_encode(value) for value in values) + + +class RTMPTests(unittest.IsolatedAsyncioTestCase): + async def test_connect_publish_and_send(self) -> None: + media_received = asyncio.Event() + errors: list[BaseException] = [] + + async def server_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + peer = RTMPClient("rtmp://localhost/live/key") + try: + handshake = await reader.readexactly(1537) + c1 = handshake[1:] + s1 = b"\x00" * 1536 + writer.write(b"\x03" + s1 + c1) + await writer.drain() + await reader.readexactly(1536) + peer._reader = reader + peer._writer = writer + connect_message = await peer.read_message() + self.assertEqual(amf0_decode_all(connect_message.payload)[0], "connect") + await peer.send_message( + RTMP_COMMAND_AMF0, + 0, + command("_result", 1, {"fmsVer": "test"}, {"code": "NetConnection.Connect.Success"}), + stream_id=0, + chunk_stream_id=3, + ) + create_message = await peer.read_message() + self.assertEqual(amf0_decode_all(create_message.payload)[0], "createStream") + await peer.send_message( + RTMP_COMMAND_AMF0, + 0, + command("_result", 2, None, 1), + stream_id=0, + chunk_stream_id=3, + ) + publish_message = await peer.read_message() + self.assertEqual(amf0_decode_all(publish_message.payload)[0], "publish") + await peer.send_message( + RTMP_COMMAND_AMF0, + 0, + command("onStatus", 0, None, {"level": "status", "code": "NetStream.Publish.Start"}), + stream_id=1, + chunk_stream_id=3, + ) + message = await peer.read_message() + self.assertEqual(message.message_type, RTMP_DATA_AMF0) + self.assertEqual(message.payload, b"hello") + media_received.set() + except BaseException as exc: + errors.append(exc) + media_received.set() + finally: + await peer.close() + + server = await asyncio.start_server(server_handler, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + client = RTMPClient(f"rtmp://127.0.0.1:{port}/live/key") + try: + await client.connect() + await client.send_message(RTMP_DATA_AMF0, 0, b"hello") + await asyncio.wait_for(media_received.wait(), 2) + if errors: + raise errors[0] + finally: + await client.close() + server.close() + await server.wait_closed() + diff --git a/tests/test_session.py b/tests/test_session.py new file mode 100644 index 0000000..48aea5b --- /dev/null +++ b/tests/test_session.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import asyncio +import unittest +from unittest.mock import patch + +from restream_playout.exceptions import StreamDisconnected +from restream_playout.session import StreamSession + + +class _FakeClient: + attempts = 0 + + def __init__(self, url: str) -> None: + self.url = url + self.connected = False + self.last_pong = 0.0 + self.fail_send = False + self._closed = asyncio.Event() + + async def connect(self) -> None: + type(self).attempts += 1 + if type(self).attempts == 1: + raise OSError("simulated connect failure") + self.connected = True + self.last_pong = 10**12 + + async def monitor(self) -> None: + await self._closed.wait() + + async def ping(self) -> None: + pass + + async def send_message(self, *args: object, **kwargs: object) -> None: + if self.fail_send: + raise OSError("simulated send failure") + + async def close(self) -> None: + self.connected = False + self._closed.set() + + async def unpublish(self) -> None: + pass + + +class SessionTests(unittest.IsolatedAsyncioTestCase): + async def test_reconnects_and_increments_generation(self) -> None: + _FakeClient.attempts = 0 + states: list[bool] = [] + stop = asyncio.Event() + + async def state_callback(connected: bool) -> None: + states.append(connected) + + with patch("restream_playout.session.RTMPClient", _FakeClient): + session = StreamSession( + "rtmp://localhost/live/key", + stop_event=stop, + state_callback=state_callback, + reconnect_min_delay=0.001, + reconnect_max_delay=0.001, + heartbeat_interval=100, + ) + self.assertEqual(await session.ensure_connected(), 1) + client = session._client + self.assertIsNotNone(client) + client.fail_send = True + with self.assertRaises(StreamDisconnected): + await session.send_message(9, 0, b"frame") + self.assertEqual(await session.ensure_connected(), 2) + await session.close() + + self.assertEqual(states, [True, False, True, False])