""" webrtc_record.py — 录制走 AWS IVS Real-Time Stage 协议的 WebRTC 直播为 mp4 关键点：SDP offer/answer 的交换全程走同一条信令 WebSocket 完成 （不是单独发 HTTP WHIP POST，那样会被服务端拒绝返回 403 no token）。 用法示例 (用 uv 运行，无需手动建venv/装依赖，uv会自动处理): uv run webrtc_record.py # 交互式粘贴token uv run webrtc_record.py --output-dir ./recordings --name my_stream uv run webrtc_record.py --token "eyJhbGc..." --output-dir ./recordings uv run webrtc_record.py --json-file play_response.json --duration 1800 # 只录30分钟自动停止 """ import argparse import asyncio import base64 import json import logging import random import signal import time import uuid from pathlib import Path import websockets from aiortc import RTCPeerConnection, RTCSessionDescription from aiortc.contrib.media import MediaRecorder logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") log = logging.getLogger("webrtc_record") def decode_jwt_payload(token: str) -> dict: payload_b64 = token.split(".")[1] padded = payload_b64 + "=" * (-len(payload_b64) % 4) return json.loads(base64.urlsafe_b64decode(padded)) def make_trace_id() -> str: ts = format(int(time.time()), "08x") rand = "".join(random.choices("0123456789abcdef", k=24)) return f"1-{ts}-{rand}" def load_token(args) -> str: if args.token: return args.token.strip() if args.json_file: data = json.loads(Path(args.json_file).read_text(encoding="utf-8")) return data["PlayList"]["whip"][0]["token"] token = input("请粘贴 whip token（JWT字符串），回车确认: ").strip() if not token: raise SystemExit("没有输入任何内容，退出") return token async def wait_for_publisher_id(ws, timeout: float = 15.0) -> str: """在已连接的 ws 上等 STAGE_STATE 消息，找出 isPublishing=true 的 participantId""" deadline = asyncio.get_event_loop().time() + timeout while True: remaining = deadline - asyncio.get_event_loop().time() if remaining <= 0: raise RuntimeError("超时：没有在 STAGE_STATE 里找到正在推流的 participant，可能主播已下播") raw = await asyncio.wait_for(ws.recv(), timeout=remaining) msg = json.loads(raw) if msg.get("type") == "aws:ivs:STAGE_STATE": for p in msg["payload"].get("participants", []): if p.get("isPublishing"): return p["id"] async def sdp_offer_answer(ws, whip_url: str, participant_id: str, offer_sdp: str, trace_id: str, timeout: float = 15.0) -> str: """把本地offer通过WS发出去(aws:ivs:SDP_OFFER)，等匹配的aws:ivs:SDP_ANSWER返回answer sdp""" request_id = str(uuid.uuid4()) offer_msg = { "type": "aws:ivs:SDP_OFFER", "payload": { "whipUrl": f"{whip_url}/subscribe/{participant_id}", "requestId": request_id, "traceId": trace_id, "platform": "web", "sdk": "1.37.0", "whipVersion": "draft-ietf-wish-whip-04", "content": offer_sdp, "contentType": "application/sdp", "initialLayerPref": "highest", }, } await ws.send(json.dumps(offer_msg)) log.info("已通过 WebSocket 发送 SDP_OFFER (requestId=%s)，等待 SDP_ANSWER...", request_id) deadline = asyncio.get_event_loop().time() + timeout while True: remaining = deadline - asyncio.get_event_loop().time() if remaining <= 0: raise RuntimeError("等待 SDP_ANSWER 超时") raw = await asyncio.wait_for(ws.recv(), timeout=remaining) msg = json.loads(raw) if msg.get("type") != "aws:ivs:SDP_ANSWER": continue p = msg["payload"] if p.get("requestId") != request_id: continue if p.get("errorCode"): raise RuntimeError(f"SDP_ANSWER 返回错误: {p.get('errorCode')} {p.get('errorMessage')}") return p["content"] async def record(args): whip_token = load_token(args) payload = decode_jwt_payload(whip_token) whip_url = payload["whip_url"] stage_id = payload["topic"] ttl = payload["exp"] - payload["iat"] log.info("stage_id=%s whip_url=%s token有效期约 %ds", stage_id, whip_url, ttl) trace_id = make_trace_id() ws_url = ( f"wss://global.events.live-video.net/" f"?error_over_ws=1&topic={stage_id}/publishers&traceid={trace_id}" ) log.info("连接信令 WebSocket: %s", ws_url) async with websockets.connect(ws_url, subprotocols=[whip_token], origin=args.origin) as ws: participant_id = await wait_for_publisher_id(ws) log.info("找到正在推流的 participantId=%s", participant_id) pc = RTCPeerConnection() pc.addTransceiver("audio", direction="recvonly") pc.addTransceiver("video", direction="recvonly") output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) ts = time.strftime("%Y%m%d_%H%M%S") out_path = output_dir / f"{args.name or stage_id}_{ts}.mp4" recorder = MediaRecorder(str(out_path)) log.info("输出文件: %s", out_path) stop_event = asyncio.Event() @pc.on("track") def on_track(track): log.info("收到 track: %s", track.kind) recorder.addTrack(track) @pc.on("connectionstatechange") async def on_state_change(): log.info("PeerConnection 状态: %s", pc.connectionState) if pc.connectionState in ("failed", "closed", "disconnected"): stop_event.set() offer = await pc.createOffer() await pc.setLocalDescription(offer) answer_sdp = await sdp_offer_answer(ws, whip_url, participant_id, pc.localDescription.sdp, trace_id) await pc.setRemoteDescription(RTCSessionDescription(sdp=answer_sdp, type="answer")) log.info("SDP 协商完成，媒体应该马上开始流入") await recorder.start() log.info("开始录制，按 Ctrl+C 停止（后台运行时用 kill -TERM ）") loop = asyncio.get_running_loop() for sig in (signal.SIGINT, signal.SIGTERM): loop.add_signal_handler(sig, stop_event.set) # 后台持续消费这条WS上的消息，防止缓冲堆积，同时如果连接断了能及时触发停止 async def drain_ws(): try: async for _ in ws: pass except Exception: pass finally: stop_event.set() drain_task = asyncio.create_task(drain_ws()) try: if args.duration: await asyncio.wait_for(stop_event.wait(), timeout=args.duration) else: await stop_event.wait() except asyncio.TimeoutError: log.info("达到设定录制时长 %ds，自动停止", args.duration) finally: drain_task.cancel() log.info("正在停止录制并写入文件……请稍等，不要强制kill") await recorder.stop() await pc.close() log.info("完成: %s", out_path) def main(): parser = argparse.ArgumentParser(description="录制 AWS IVS 实时舞台 WebRTC 直播为 mp4") src = parser.add_mutually_exclusive_group(required=False) src.add_argument("--json-file", help="保存的 play 接口响应 JSON 文件路径") src.add_argument("--token", help="直接传入 whip token 字符串") parser.add_argument("--output-dir", default="./recordings", help="录制文件保存目录") parser.add_argument("--name", default=None, help="文件名前缀，默认用 stage_id") parser.add_argument("--origin", default="https://www.pandalive.co.kr", help="WS 握手用的 Origin 头") parser.add_argument("--duration", type=int, default=None, help="录制时长(秒)，不填则一直录到手动停止") args = parser.parse_args() asyncio.run(record(args)) if __name__ == "__main__": main()