mirror of
https://github.com/ApfelTeeSaft/restream_playout.git
synced 2026-08-26 19:33:32 +00:00
ho lee fuk
This commit is contained in:
@@ -0,0 +1,7 @@
|
|||||||
|
__pycache__/
|
||||||
|
*.py[cod]
|
||||||
|
.pytest_cache/
|
||||||
|
build/
|
||||||
|
dist/
|
||||||
|
*.egg-info/
|
||||||
|
|
||||||
@@ -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.
|
||||||
|
|
||||||
@@ -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()
|
||||||
|
|
||||||
@@ -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()
|
||||||
|
|
||||||
@@ -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()
|
||||||
|
|
||||||
@@ -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"]
|
||||||
|
|
||||||
@@ -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"
|
||||||
|
|
||||||
@@ -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
|
||||||
@@ -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)
|
||||||
|
|
||||||
@@ -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."""
|
||||||
|
|
||||||
@@ -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)
|
||||||
|
|
||||||
@@ -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")
|
||||||
|
|
||||||
@@ -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,
|
||||||
|
)
|
||||||
|
|
||||||
@@ -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))
|
||||||
@@ -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()
|
||||||
|
|
||||||
@@ -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
|
||||||
@@ -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)
|
||||||
@@ -0,0 +1 @@
|
|||||||
|
|
||||||
@@ -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)
|
||||||
|
|
||||||
@@ -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))
|
||||||
|
|
||||||
@@ -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)
|
||||||
|
|
||||||
@@ -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")
|
||||||
|
|
||||||
@@ -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, [])
|
||||||
|
|
||||||
@@ -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()
|
||||||
|
|
||||||
@@ -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()
|
||||||
|
|
||||||
@@ -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])
|
||||||
Reference in New Issue
Block a user