yum/3ner

A toon shader for Unity's BIRP.

git clone https://git.yummers.dev/yum/3ner

yumadd basic encryption. note that keys are publicly availablea91dd1d

master
branchesmasterall refs →
14.1 KiB452 linesraw
1#!/usr/bin/env python3
2from dotenv import load_dotenv
3import os
4import secrets
5import subprocess
6import threading
7import shutil
8import time
9from pathlib import Path
10from typing import Optional
11from flask import Flask, request
12import logging
13import atexit
14
15# Setup Flask and load environment variables
16app = Flask(__name__)
17application = app  # WSGI servers like gunicorn look for 'application'
18load_dotenv()
19
20# Configuration
21INGEST_PSK = os.environ.get('OBS_STREAM_KEY') or os.environ.get('STREAM_PSK')
22INGEST_RTMP_HOST = os.environ.get('INGEST_RTMP_HOST', '127.0.0.1')
23try:
24    INGEST_RTMP_PORT = int(os.environ.get('INGEST_RTMP_PORT', '1936'))
25except ValueError as exc:
26    raise ValueError('INGEST_RTMP_PORT must be an integer') from exc
27INGEST_THREAD_QUEUE_SIZE = int(os.environ.get('INGEST_THREAD_QUEUE_SIZE', '4096'))
28HLS_SEGMENT_TIME = float(os.environ.get('HLS_SEGMENT_TIME', '2'))
29HLS_PLAYLIST_SIZE = int(os.environ.get('HLS_PLAYLIST_SIZE', '3'))
30HLS_DELETE_THRESHOLD = int(os.environ.get('HLS_DELETE_THRESHOLD', '20'))
31HLS_PLAYLIST_TIMEOUT = float(os.environ.get('HLS_PLAYLIST_TIMEOUT', '5'))
32HLS_PLAYLIST_POLL_INTERVAL = float(os.environ.get('HLS_PLAYLIST_POLL_INTERVAL', '0.1'))
33BASE_DIR = Path(os.environ.get('STREAM_DIR', '/var/www/streams'))
34SERVER_DOMAIN = os.environ.get('SERVER_DOMAIN', 'yummers.b-cdn.net')
35STREAM_HEX = secrets.token_hex(16)
36STREAM_PATH = BASE_DIR / 'live' / STREAM_HEX
37HLS_ROUTE_PREFIX = f"/hls/{STREAM_HEX}"
38SESSION_KEY_NAME = 'session.key'
39SESSION_KEYINFO_NAME = 'session.keyinfo'
40SESSION_KEY_URI: Optional[str] = None
41# Media output settings tuned for VRChat playback
42AUDIO_BITRATE = '256k'
43AUDIO_CHANNELS = 2
44AUDIO_SAMPLE_RATE = 48000
45
46# Setup logging
47logging.basicConfig(
48    level=getattr(logging, os.environ.get('LOG_LEVEL', 'INFO')),
49    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
50)
51logger = logging.getLogger('obs_proxy')
52
53# Global state
54ffmpeg_process = None
55ffmpeg_worker_thread = None
56ffmpeg_stop_event = threading.Event()
57
58
59def _worker_running() -> bool:
60    """Return True if the ffmpeg worker thread is currently active."""
61    return ffmpeg_worker_thread is not None and ffmpeg_worker_thread.is_alive()
62
63# Validate configuration
64if not INGEST_PSK:
65    logger.error("OBS_STREAM_KEY/STREAM_PSK is not set")
66    exit(1)
67
68# Create required directories
69BASE_DIR.mkdir(parents=True, exist_ok=True)
70STREAM_PATH.mkdir(parents=True, exist_ok=True)
71
72
73def _session_key_path() -> Path:
74    return STREAM_PATH / SESSION_KEY_NAME
75
76
77def _session_keyinfo_path() -> Path:
78    return STREAM_PATH / SESSION_KEYINFO_NAME
79
80
81def _write_key_material() -> None:
82    """Generate and persist AES-128 key + keyinfo for the current session."""
83    global SESSION_KEY_URI
84
85    key_bytes = secrets.token_bytes(16)
86    iv_bytes = secrets.token_bytes(16)
87    key_path = _session_key_path()
88    key_path.write_bytes(key_bytes)
89
90    key_uri = f"https://{SERVER_DOMAIN}{HLS_ROUTE_PREFIX}/{SESSION_KEY_NAME}"
91    keyinfo_path = _session_keyinfo_path()
92    iv_hex = format(int.from_bytes(iv_bytes, 'big'), '032x')
93    keyinfo_path.write_text(
94        f"{key_uri}\n{key_path}\n{iv_hex}\n",
95        encoding="utf-8",
96    )
97    SESSION_KEY_URI = key_uri
98
99def reset_stream_path():
100    """Ensure the live stream directory is empty and ready."""
101    shutil.rmtree(STREAM_PATH, ignore_errors=True)
102    STREAM_PATH.mkdir(parents=True, exist_ok=True)
103    _write_key_material()
104
105
106def _safe_reset_stream_path(context: str) -> bool:
107    """Reset the stream directory and log any failure."""
108    try:
109        reset_stream_path()
110        return True
111    except Exception as exc:  # pragma: no cover - best effort logging
112        logger.error("Error resetting stream directory %s: %s", context, exc)
113        return False
114
115
116def _verify_stream_path_writable() -> None:
117    """Ensure the stream directory is writable before launching FFmpeg."""
118    test_file = STREAM_PATH / "write_test.txt"
119    try:
120        with open(test_file, "w", encoding="utf-8") as probe:
121            probe.write("ok")
122    except Exception as exc:  # pragma: no cover - best effort logging
123        logger.error("Could not write to stream directory: %s", exc)
124    finally:
125        if test_file.exists():
126            test_file.unlink()
127
128
129def _wait_for_playlist(process: subprocess.Popen[str], attempt: int):
130    """Block until the playlist appears or we hit the configured timeout."""
131    playlist_path = STREAM_PATH / 'stream.m3u8'
132    start = time.monotonic()
133
134    while True:
135        if playlist_path.exists():
136            elapsed = time.monotonic() - start
137            segments = sorted(seg.name for seg in playlist_path.parent.glob('segment-*.ts'))
138            logger.info(
139                "HLS playlist materialized after %.2fs at %s; segments=%s",
140                elapsed,
141                playlist_path,
142                segments,
143            )
144            return True, 'ready'
145
146        if ffmpeg_stop_event.is_set():
147            return False, 'stop_requested'
148
149        if process.poll() is not None:
150            return False, 'ffmpeg_exited'
151
152        elapsed = time.monotonic() - start
153        if elapsed >= HLS_PLAYLIST_TIMEOUT:
154            logger.warning(
155                "HLS playlist still missing after %.2fs on attempt %s; recycling FFmpeg",
156                elapsed,
157                attempt,
158            )
159            return False, 'timeout'
160
161        if ffmpeg_stop_event.wait(HLS_PLAYLIST_POLL_INTERVAL):
162            return False, 'stop_requested'
163
164
165def _start_pipe_logger(pipe, level):
166    """Drain an ffmpeg pipe on a background thread to avoid deadlocks."""
167
168    def pipe_logger():
169        with pipe:
170            for line in iter(pipe.readline, ''):
171                line = line.strip()
172                if line:
173                    logger.log(level, "ffmpeg: %s", line)
174
175    threading.Thread(target=pipe_logger, daemon=True).start()
176
177
178def _terminate_ffmpeg_process(process: subprocess.Popen[str], *, timeout: float = 5.0, log_errors: bool = True) -> None:
179    """Terminate an ffmpeg process, falling back to kill if needed."""
180    try:
181        process.terminate()
182        process.wait(timeout=timeout)
183    except Exception as exc:  # pragma: no cover - best effort logging
184        if log_errors:
185            logger.error(f"Error stopping FFmpeg: {exc}")
186        try:
187            process.kill()
188        except Exception:  # pragma: no cover
189            pass
190
191
192def _build_ffmpeg_command() -> list[str]:
193    """Construct the ffmpeg command line we execute for each attempt."""
194
195    keyinfo_path = _session_keyinfo_path()
196    if not keyinfo_path.exists():
197        _write_key_material()
198
199    keyinfo_path = _session_keyinfo_path()
200    return [
201        'ffmpeg',
202        '-nostdin',
203        '-hide_banner',
204        '-loglevel', os.environ.get('FFMPEG_LOGLEVEL', 'warning'),
205        '-fflags', '+genpts',
206        '-thread_queue_size', str(INGEST_THREAD_QUEUE_SIZE),
207        '-i', f'rtmp://{INGEST_RTMP_HOST}:{INGEST_RTMP_PORT}/live/{INGEST_PSK}',
208        '-map', '0:v:0?',
209        '-map', '0:a:0?',
210        '-c:v', 'copy',
211        '-c:a', 'aac',
212        '-b:a', AUDIO_BITRATE,
213        '-ac', str(AUDIO_CHANNELS),
214        '-ar', str(AUDIO_SAMPLE_RATE),
215        '-f', 'hls',
216        '-hls_time', str(HLS_SEGMENT_TIME),
217        '-hls_list_size', str(HLS_PLAYLIST_SIZE),
218        '-hls_flags', 'delete_segments+independent_segments',
219        '-hls_delete_threshold', str(HLS_DELETE_THRESHOLD),
220        '-hls_key_info_file', str(keyinfo_path),
221        '-hls_segment_filename', str(STREAM_PATH / 'segment-%05d.ts'),
222        str(STREAM_PATH / 'stream.m3u8'),
223    ]
224
225
226def _run_ffmpeg_once(attempt: int) -> bool:
227    """Start ffmpeg, wait for it to exit, and report whether it ran cleanly."""
228    global ffmpeg_process
229
230    logger.info("Starting FFmpeg for live stream (attempt %s)", attempt)
231
232    try:
233        process = subprocess.Popen(
234            _build_ffmpeg_command(),
235            stdout=subprocess.PIPE,
236            stderr=subprocess.PIPE,
237            text=True,
238            bufsize=1,
239        )
240    except Exception as exc:  # pragma: no cover - best effort logging
241        logger.error("Failed to start FFmpeg: %s", exc)
242        return False
243
244    ffmpeg_process = process
245    start_time = time.monotonic()
246
247    logger.info("FFmpeg process started with PID %s", process.pid)
248    logger.info('Stream active; waiting for playlist to appear')
249
250    _start_pipe_logger(process.stderr, logging.WARNING)
251    _start_pipe_logger(process.stdout, logging.DEBUG)
252
253    playlist_ready, playlist_reason = _wait_for_playlist(process, attempt)
254
255    if not playlist_ready and not ffmpeg_stop_event.is_set():
256        if playlist_reason == 'timeout' and process.poll() is None:
257            logger.info("Terminating FFmpeg attempt %s after playlist timeout", attempt)
258            _terminate_ffmpeg_process(process, log_errors=False)
259        elif playlist_reason == 'ffmpeg_exited':
260            logger.debug("FFmpeg exited before playlist became available on attempt %s", attempt)
261
262    exit_code = process.wait()
263    elapsed = time.monotonic() - start_time
264
265    ffmpeg_process = None
266
267    if ffmpeg_stop_event.is_set():
268        logger.info(
269            "FFmpeg stop requested; process exited with code %s after %.2fs",
270            exit_code,
271            elapsed,
272        )
273        return True
274
275    if not playlist_ready:
276        logger.error(
277            "FFmpeg attempt %s ended (exit %s) before playlist became available (reason=%s)",
278            attempt,
279            exit_code,
280            playlist_reason,
281        )
282        return False
283
284    if exit_code != 0:
285        logger.error("FFmpeg exited with code %s after %.2fs", exit_code, elapsed)
286        return False
287
288    logger.info("FFmpeg process completed successfully after %.2fs", elapsed)
289    return True
290
291
292def _ffmpeg_worker_loop() -> None:
293    """Keep spawning FFmpeg until it runs cleanly or a stop is requested."""
294    attempt = 0
295
296    while not ffmpeg_stop_event.is_set():
297        attempt += 1
298        if _run_ffmpeg_once(attempt):
299            break
300
301        if ffmpeg_stop_event.is_set():
302            break
303
304        _safe_reset_stream_path("between FFmpeg attempts")
305
306    logger.debug("FFmpeg worker exiting")
307
308
309def start_ffmpeg_process():
310    """Start FFmpeg to convert RTMP ingest into HLS."""
311    global ffmpeg_worker_thread
312
313    if _worker_running():
314        logger.warning("FFmpeg worker already running; skipping duplicate start")
315        return True
316
317    if _safe_reset_stream_path("before FFmpeg start"):
318        logger.info(f"Stream directory ready at {STREAM_PATH}")
319    _verify_stream_path_writable()
320
321    ffmpeg_stop_event.clear()
322
323    try:
324        ffmpeg_worker_thread = threading.Thread(
325            target=_ffmpeg_worker_loop,
326            daemon=True,
327            name="ffmpeg-worker",
328        )
329        ffmpeg_worker_thread.start()
330        return True
331    except Exception as exc:  # pragma: no cover - defensive logging
332        logger.error(f"Failed to start FFmpeg worker thread: {exc}")
333        ffmpeg_worker_thread = None
334        return False
335
336
337def cleanup_stream():
338    """Stop FFmpeg and purge any cached HLS segments."""
339    global ffmpeg_process, ffmpeg_worker_thread
340
341    ffmpeg_stop_event.set()
342
343    if ffmpeg_process:
344        _terminate_ffmpeg_process(ffmpeg_process)
345        ffmpeg_process = None
346
347    worker = ffmpeg_worker_thread
348    if worker and worker.is_alive():
349        worker.join(timeout=5)
350        if worker.is_alive():  # pragma: no cover - diagnostic
351            logger.warning("FFmpeg worker thread did not exit cleanly")
352    ffmpeg_worker_thread = None
353
354    _safe_reset_stream_path("during cleanup")
355
356    ffmpeg_stop_event.clear()
357
358
359# Routes
360@app.route('/rtmp_callbacks/on_publish', methods=['POST'])
361def on_publish():
362    """Callback when a stream starts"""
363    stream_key = request.form.get('name')
364    logger.info("on_publish received for key=%s", stream_key)
365
366    if not stream_key or stream_key != INGEST_PSK:
367        logger.warning("Unauthorized stream key attempted to publish: %s", stream_key)
368        return "Unauthorized", 403
369
370    if ffmpeg_process or _worker_running():
371        logger.info("Stream already active, recycling existing session")
372        cleanup_stream()
373
374    if not start_ffmpeg_process():
375        return "Failed to start stream", 500
376
377    playlist_path = STREAM_PATH / 'stream.m3u8'
378    if playlist_path.exists():
379        logger.info("Existing playlist found at %s", playlist_path)
380
381    return "OK"
382
383
384@app.route('/rtmp_callbacks/on_publish_done', methods=['POST'])
385def on_publish_done():
386    """Callback when a stream ends"""
387    stream_key = request.form.get('name')
388    logger.info("on_publish_done received for key=%s", stream_key)
389
390    if not stream_key or stream_key != INGEST_PSK:
391        logger.warning("on_publish_done received for unknown key: %s", stream_key)
392        return "Bad request", 400
393
394    if ffmpeg_process or _worker_running():
395        cleanup_stream()
396
397    logger.info("Stream publishing ended")
398    return "OK"
399
400
401@app.route('/health')
402def health_check():
403    """Health check endpoint"""
404    return {
405        "status": "healthy",
406        "streaming": ffmpeg_process is not None
407    }
408
409
410def print_instructions():
411    """Print usage instructions"""
412    obs_url = f"rtmps://{SERVER_DOMAIN}:1935/live"
413    hls_url = f"https://{SERVER_DOMAIN}{HLS_ROUTE_PREFIX}/stream.m3u8"
414
415    print("\n" + "="*80)
416    print(f"{'OBS TO VRCHAT STREAMING PROXY':^80}")
417    print("="*80)
418
419    print("\n[URLS]")
420    print(f"  OBS ingest: {obs_url}")
421    print(f"  HLS:        {hls_url}")
422    if SESSION_KEY_URI:
423        print(f"  HLS key:   {SESSION_KEY_URI}")
424
425    print("\n[STATUS]")
426    print(f"  Stream is {'ACTIVE' if ffmpeg_process else 'INACTIVE'}")
427    print(f"  Session ID: {STREAM_HEX}")
428    print(f"  On-disk path: {STREAM_PATH}")
429    print("="*80 + "\n")
430
431# Register cleanup once the module is imported so any WSGI server benefits.
432atexit.register(cleanup_stream)
433
434
435def main():
436    """Entry point for running with a production WSGI server."""
437    try:
438        from waitress import serve
439    except ImportError as exc:  # pragma: no cover - defensive guardrail
440        raise RuntimeError(
441            "Waitress is required to run this service. Install it with 'pip install waitress'."
442        ) from exc
443
444    print_instructions()
445    host = os.environ.get('HOST', '0.0.0.0')
446    port = int(os.environ.get('PORT', 5000))
447    logger.info("Starting Waitress on %s:%s", host, port)
448    serve(app, host=host, port=port)
449
450
451if __name__ == '__main__':
452    main()