3M/cef_3M/endpoints/mediamtx.py

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)