import asyncio import subprocess import time from enum import Enum import aiohttp import ffmpeg import select from starlette.responses import JSONResponse, Response, FileResponse from config import MEDIAMTX_API, MEDIAMTX_RTSP from . import router from fastapi import Request, Depends from ..auth import decodeJWT, JWTBearer from ..util import redis, ergo, privilegedIps def pathParts(path): if len(path) == 2: channel = None user, token = path elif len(path) == 3: channel, user, token = path channel = "#" + channel else: return None, None, None return channel, user, token class TargetType(str, Enum): user = "u" channel = "c" @router.get("/mediamtx/streams/{channel}", dependencies=[Depends(JWTBearer(True))]) async def mediamtxChannelStreams(request: Request, channel: str): inChannel = request.state.jwt.get("channel", "").lower() == "#" + channel.lower() results = [] async for result in redis.scan_iter(f"stream #{channel} *"): _, channel, user, token = result.decode("utf8").split() if inChannel or token == "public": results.append({ "user": user, "token": token }) return JSONResponse(status_code=200, content=results) @router.get("/mediamtx/thumbnail/{user}/{key}") @router.get("/mediamtx/thumbnail/{channel}/{user}/{key}") async def mediamtxThumbnail(user: str, key: str, channel: str = None): blob = "".join([x for x in (channel, user, key) if x]) if "?" in blob: return Response(status_code=404) print(f"rtsp://{MEDIAMTX_RTSP}/" + "/".join([x for x in (channel, user, key) if x])) proc: subprocess.Popen = ffmpeg.input(f"rtsp://{MEDIAMTX_RTSP}/" + "/".join([x for x in (channel, user, key) if x]), ).output('pipe:', loglevel="quiet", vframes=1, format='image2', vcodec='mjpeg').run_async(pipe_stdout=True) timeout = time.time() + 8 data = b"" while 1: if select.select([proc.stdout], [], [], 0)[0]: data += proc.stdout.read(65535) if proc.poll() is not None: break else: if timeout < time.time(): proc.kill() return Response(status_code=503) await asyncio.sleep(0.1) return Response(content=data, media_type="image/jpeg", headers={ "Cache-Control": "public, max-age=60, stale-while-revalidate=300" }) @router.get("/mediamtx/public", dependencies=[]) async def mediamtxAllPublicStreams(): results = [] async for result in redis.scan_iter(f"stream *"): _, channel, user, token = result.decode("utf8").split() if token == "public": results.append({ "target": channel, "user": user, "token": token }) return JSONResponse(status_code=200, content=results) @router.get("/mediamtx/public/{targetType}/{target}", dependencies=[]) async def mediamtxEntityPublicStreams(targetType: TargetType, target: str): if targetType == TargetType.channel: target = "#" + target results = [] async for result in redis.scan_iter(f"stream {target} *"): _, channel, user, token = result.decode("utf8").split() print(result) if token == "public": results.append({ "user": user, "token": token }) return JSONResponse(status_code=200, content=results) @router.post("/mediamtx/auth", include_in_schema=False) async def mediamtxAuth(request: Request): if request.client.host not in privilegedIps: return False body = await request.json() # Just never expose RTSP :) if body["protocol"] == "rtsp": return JSONResponse(status_code=200, content={"success": True}) jwt = decodeJWT(body["query"][4:]) path = body["path"].split("/") reading = body["action"] == "read" channel, user, token = pathParts(path) if user is None: return JSONResponse(status_code=400, content={"error": "bad path"}) if " " in token: return JSONResponse(status_code=400, content={"error": "bad token"}) # the only time we don't care about JWT is if someone is watching a public stream if reading and token == "public": return JSONResponse(status_code=200, content={"success": True}) if len(jwt.keys()) == 0: return JSONResponse(status_code=403, content={"error": "bad jwt"}) # TODO: channel stream permissions # publishing if not reading: if user != jwt["account"]: return JSONResponse(status_code=403, content={"error": "nuh uh"}) if channel and jwt["channel"] != channel: return JSONResponse(status_code=403, content={"error": "nuh uh"}) return JSONResponse(status_code=200, content={"success": True}) @router.post("/mediamtx/add", include_in_schema=False) async def mediamtxAdd(request: Request): if request.client.host not in privilegedIps: return False body = await request.json() path = body["env"]["MTX_PATH"].split("/") parts = [x for x in pathParts(path) if x] await redis.set("stream " + " ".join(parts), parts[2], ex=60) if len(parts) == 3: await ergo.broadcastTo(parts[0], "STREAMSTART", parts[0], parts[1], parts[2]) @router.post("/mediamtx/del", include_in_schema=False) async def mediamtxDelete(request: Request): if request.client.host not in privilegedIps: return False body = await request.json() path = body["env"]["MTX_PATH"].split("/") parts = [x for x in pathParts(path) if x] await redis.delete("stream " + " ".join(parts)) if len(parts) == 3: await ergo.broadcastTo(parts[0], "STREAMEND", parts[0], parts[1], parts[2]) # Not an endpoint async def mediamtxPoll(): while 1: async with aiohttp.ClientSession() as sess: req = await sess.get(MEDIAMTX_API +"/v3/paths/list?itemsPerPage=9999999") streams = await req.json() for stream in streams["items"]: if stream["ready"]: path = stream["name"].split("/") parts = [x for x in pathParts(path) if x] await redis.set(f"stream " + " ".join(parts), parts[-1], ex=60) await asyncio.sleep(30)