Add ffmpeg proxy HTTP view

This commit is contained in:
Michael Hansen 2024-07-29 16:13:41 -05:00
parent 1c03c83c0a
commit 6b5d4fcccb
4 changed files with 186 additions and 1 deletions

View file

@ -4,12 +4,17 @@ from __future__ import annotations
import asyncio
from functools import cached_property
from http import HTTPStatus
import logging
import re
from aiohttp import web
from aiohttp.abc import AbstractStreamWriter, BaseRequest
from haffmpeg.core import HAFFmpeg
from haffmpeg.tools import IMAGE_JPEG, FFVersion, ImageFrame
import voluptuous as vol
from homeassistant.components.http import HomeAssistantView
from homeassistant.const import (
ATTR_ENTITY_ID,
CONTENT_TYPE_MULTIPART,
@ -28,7 +33,7 @@ from homeassistant.helpers.typing import ConfigType
from homeassistant.loader import bind_hass
from homeassistant.util.signal_type import SignalType
DOMAIN = "ffmpeg"
from .const import DOMAIN
SERVICE_START = "start"
SERVICE_STOP = "stop"
@ -65,6 +70,8 @@ CONFIG_SCHEMA = vol.Schema(
SERVICE_FFMPEG_SCHEMA = vol.Schema({vol.Optional(ATTR_ENTITY_ID): cv.entity_ids})
_LOGGER = logging.getLogger(__name__)
async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
"""Set up the FFmpeg component."""
@ -98,6 +105,8 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
DOMAIN, SERVICE_RESTART, async_service_handle, schema=SERVICE_FFMPEG_SCHEMA
)
hass.http.register_view(FFmpegProxyView(manager))
hass.data[DATA_FFMPEG] = manager
return True
@ -259,3 +268,133 @@ class FFmpegBase[_HAFFmpegT: HAFFmpeg](Entity):
self.async_write_ha_state()
self.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_START, async_start_handle)
class FFmpegConvertResponse(web.StreamResponse):
"""HTTP streaming response that uses ffmpeg to convert audio from a URL."""
def __init__(
self,
manager: FFmpegManager,
url: str,
fmt: str,
rate: int | None,
channels: int | None,
chunk_size: int = 2048,
) -> None:
"""Initialize response.
Parameters
----------
manager: FFmpegManager
ffmpeg manager
url: str
URL of source audio stream
fmt: str
Target format of audio (flac, mp3, wav, etc.)
rate: int, optional
Target sample rate in hertz (None = same as source)
channels: int, optional
Target number of channels (None = same as source)
chunk_size: int
Number of bytes to read from ffmpeg process at a time
"""
super().__init__(status=200)
self.manager = manager
self.url = url
self.fmt = fmt
self.rate = rate
self.channels = channels
self.chunk_size = chunk_size
async def prepare(self, request: BaseRequest) -> AbstractStreamWriter | None:
"""Stream url through ffmpeg conversion and out to HTTP client."""
writer = await super().prepare(request)
assert writer is not None
command_args = ["-i", self.url, "-f", self.fmt]
if self.rate is not None:
# Sample rate
command_args.extend(["-ar", str(self.rate)])
if self.channels is not None:
# Number of channels
command_args.extend(["-ac", str(self.channels)])
# Output to stdout
command_args.append("pipe:")
proc = await asyncio.create_subprocess_exec(
self.manager.binary,
*command_args,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
)
assert proc.stdout is not None
assert proc.stderr is not None
try:
# Pull audio chunks from ffmpeg and pass them to the HTTP client
while chunk := await proc.stdout.read(self.chunk_size):
await writer.write(chunk)
# Try to gracefully stop
proc.terminate()
await proc.wait()
finally:
await writer.write_eof()
if proc.returncode != 0:
# Process did not exit successfully
stderr_text = ""
while line := await proc.stderr.readline():
stderr_text += line.decode()
_LOGGER.error(stderr_text)
return writer
class FFmpegProxyView(HomeAssistantView):
"""FFmpeg web view to convert audio and stream back to client."""
requires_auth = False
url = "/api/ffmpeg_proxy"
name = "api:ffmpeg_proxy"
def __init__(self, manager: FFmpegManager) -> None:
"""Initialize an ffmpeg view."""
self.manager = manager
async def get(self, request: web.Request) -> web.StreamResponse:
"""Start a get request."""
query = request.query
url = query.get("url")
fmt = query.get("format")
if (not url) or (not fmt):
return web.Response(
body="url and format are required", status=HTTPStatus.BAD_REQUEST
)
try:
if rate_str := query.get("rate"):
rate = int(rate_str)
else:
rate = None
if channels_str := query.get("channels"):
channels = int(channels_str)
else:
channels = None
except ValueError:
return web.Response(
body="Invalid rate or channels value", status=HTTPStatus.BAD_REQUEST
)
# Stream converted audio back to client
return FFmpegConvertResponse(
self.manager, url=url, fmt=fmt, rate=rate, channels=channels
)

View file

@ -0,0 +1,3 @@
"""Constants for ffmpeg integration."""
DOMAIN = "ffmpeg"

View file

@ -1,7 +1,9 @@
{
"domain": "ffmpeg",
"name": "FFmpeg",
"after_dependencies": ["media_source"],
"codeowners": [],
"dependencies": ["http"],
"documentation": "https://www.home-assistant.io/integrations/ffmpeg",
"requirements": ["ha-ffmpeg==3.2.0"]
}

View file

@ -1,6 +1,13 @@
"""The tests for Home Assistant ffmpeg."""
from http import HTTPStatus
import io
import tempfile
from unittest.mock import AsyncMock, MagicMock, Mock, call, patch
from urllib.request import pathname2url
import wave
import mutagen
from homeassistant.components import ffmpeg
from homeassistant.components.ffmpeg import (
@ -19,6 +26,7 @@ from homeassistant.core import HomeAssistant, callback
from homeassistant.setup import async_setup_component, setup_component
from tests.common import assert_setup_component, get_test_home_assistant
from tests.typing import ClientSessionGenerator
@callback
@ -294,3 +302,36 @@ async def test_ffmpeg_using_official_image(
manager = get_ffmpeg_manager(hass)
assert "ffmpeg" in manager.ffmpeg_stream_content_type
async def test_proxy_view(
hass: HomeAssistant,
hass_client: ClientSessionGenerator,
) -> None:
"""Test proxy HTTP view for converting audio."""
await async_setup_component(hass, ffmpeg.DOMAIN, {ffmpeg.DOMAIN: {}})
client = await hass_client()
with tempfile.NamedTemporaryFile(mode="wb+", suffix=".wav") as temp_file:
with wave.open(temp_file.name, "wb") as wav_file:
wav_file.setframerate(16000)
wav_file.setsampwidth(2)
wav_file.setnchannels(1)
wav_file.writeframes(bytes(16000 * 2)) # 1s
temp_file.seek(0)
wav_url = pathname2url(temp_file.name)
url = f"/api/ffmpeg_proxy?url={wav_url}&format=mp3&rate=22050&channels=2"
req = await client.get(url)
assert req.status == HTTPStatus.OK
mp3_data = await req.content.read()
# Verify conversion
with io.BytesIO(mp3_data) as mp3_io:
mp3_file = mutagen.File(mp3_io)
assert mp3_file.info.sample_rate == 22050
assert mp3_file.info.channels == 2
# About a second, but not exact
assert round(mp3_file.info.length, 0) == 1