supervisor/tests/api/test_proxy.py

"""Test Home Assistant proxy."""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable, Coroutine, Generator
from json import dumps
import logging
from typing import Any, cast
from unittest.mock import AsyncMock, patch

from aiohttp import ClientPayloadError, ClientWebSocketResponse, WSCloseCode, web
from aiohttp.http_websocket import WSMessage, WSMsgType
from aiohttp.test_utils import TestClient
import pytest

from supervisor.api.proxy import APIProxy
from supervisor.apps.app import App
from supervisor.const import ATTR_ACCESS_TOKEN
from supervisor.homeassistant.api import HomeAssistantAPI


def id_generator() -> Generator[int]:
    """Generate IDs for WS messages."""
    i = 0
    while True:
        yield (i := i + 1)


class MockHAClientWebSocket(ClientWebSocketResponse):
    """Protocol for a wrapped ClientWebSocketResponse."""

    client: TestClient
    send_json_auto_id: Callable[[dict[str, Any]], Coroutine[Any, Any, None]]


class MockHAServerWebSocket:
    """Mock of HA Websocket server."""

    closed: bool = False
    close_code: int | None = None

    def __init__(self) -> None:
        """Initialize object."""
        self.outgoing: asyncio.Queue[WSMessage] = asyncio.Queue()
        self.incoming: asyncio.Queue[WSMessage] = asyncio.Queue()
        self._id_generator = id_generator()

    async def receive(self) -> WSMessage:
        """Receive next message."""
        try:
            return await self.outgoing.get()
        except asyncio.QueueShutDown:
            return WSMessage(WSMsgType.CLOSED, None, None)

    def send_str(self, data: str) -> Awaitable[None]:
        """Incoming string message."""
        return self.incoming.put(WSMessage(WSMsgType.TEXT, data, None))

    def send_bytes(self, data: bytes) -> Awaitable[None]:
        """Incoming string message."""
        return self.incoming.put(WSMessage(WSMsgType.BINARY, data, None))

    def respond_json(self, data: dict[str, Any]) -> Awaitable[None]:
        """Respond with JSON."""
        return self.outgoing.put(
            WSMessage(
                WSMsgType.TEXT, dumps(data | {"id": next(self._id_generator)}), None
            )
        )

    def respond_bytes(self, data: bytes) -> Awaitable[None]:
        """Respond with binary."""
        return self.outgoing.put(WSMessage(WSMsgType.BINARY, data, None))

    async def close(self, code: int = WSCloseCode.OK) -> None:
        """Close connection."""
        self.closed = True
        self.outgoing.shutdown(immediate=True)
        self.close_code = code


WebSocketGenerator = Callable[..., Coroutine[Any, Any, MockHAClientWebSocket]]


@pytest.fixture(name="ha_ws_server")
async def fixture_ha_ws_server() -> MockHAServerWebSocket:
    """Mock HA WS server for testing."""
    with patch.object(
        APIProxy,
        "_websocket_client",
        return_value=(mock_server := MockHAServerWebSocket()),
    ):
        yield mock_server


@pytest.fixture(name="proxy_ws_client")
def fixture_proxy_ws_client(
    api_client: TestClient, ha_ws_server: MockHAServerWebSocket
) -> WebSocketGenerator:
    """Websocket client fixture connected to websocket server."""

    async def create_client(auth_token: str) -> MockHAClientWebSocket:
        """Create a websocket client."""
        websocket = await api_client.ws_connect("/core/websocket")
        auth_resp = await websocket.receive_json()
        assert auth_resp["type"] == "auth_required"
        await websocket.send_json({"type": "auth", "access_token": auth_token})

        auth_ok = await websocket.receive_json()
        assert auth_ok["type"] == "auth_ok"

        _id_generator = id_generator()

        def _send_json_auto_id(data: dict[str, Any]) -> Coroutine[Any, Any, None]:
            data["id"] = next(_id_generator)
            return websocket.send_json(data)

        # wrap in client
        wrapped_websocket = cast(MockHAClientWebSocket, websocket)
        wrapped_websocket.client = api_client
        wrapped_websocket.send_json_auto_id = _send_json_auto_id
        return wrapped_websocket

    return create_client


async def test_proxy_message(
    proxy_ws_client: WebSocketGenerator,
    ha_ws_server: MockHAServerWebSocket,
    install_app_ssh: App,
):
    """Test proxy a message to and from Home Assistant."""
    install_app_ssh.persist[ATTR_ACCESS_TOKEN] = "abc123"
    client: MockHAClientWebSocket = await proxy_ws_client(
        install_app_ssh.supervisor_token
    )

    await client.send_json_auto_id({"hello": "world"})
    proxied_msg = await ha_ws_server.incoming.get()
    assert proxied_msg.type == WSMsgType.TEXT
    assert proxied_msg.data == '{"hello": "world", "id": 1}'

    await ha_ws_server.respond_json({"world": "received"})
    assert await client.receive_json() == {"world": "received", "id": 1}

    assert await client.close()


async def test_proxy_binary_message(
    proxy_ws_client: WebSocketGenerator,
    ha_ws_server: MockHAServerWebSocket,
    install_app_ssh: App,
):
    """Test proxy a binary message to and from Home Assistant."""
    install_app_ssh.persist[ATTR_ACCESS_TOKEN] = "abc123"
    client: MockHAClientWebSocket = await proxy_ws_client(
        install_app_ssh.supervisor_token
    )

    await client.send_bytes(b"hello world")
    proxied_msg = await ha_ws_server.incoming.get()
    assert proxied_msg.type == WSMsgType.BINARY
    assert proxied_msg.data == b"hello world"

    await ha_ws_server.respond_bytes(b"world received")
    assert await client.receive_bytes() == b"world received"

    assert await client.close()


async def test_proxy_large_message(
    proxy_ws_client: WebSocketGenerator,
    ha_ws_server: MockHAServerWebSocket,
    install_app_ssh: App,
):
    """Test too large message handled gracefully."""
    install_app_ssh.persist[ATTR_ACCESS_TOKEN] = "abc123"
    client: MockHAClientWebSocket = await proxy_ws_client(
        install_app_ssh.supervisor_token
    )

    # Test message over size limit of 4MB. Since aiohttp 3.14.1 the server
    # rejects the oversized frame from its header and resets the connection
    # before the full payload is sent, so the send itself may raise in
    # addition to the CLOSE frame. See aio-libs/aiohttp#12817.
    try:
        await client.send_bytes(bytearray(1024 * 1024 * 4))
    except ConnectionError:
        pass
    else:
        msg = await client.receive()
        assert msg.type == WSMsgType.CLOSE
        assert msg.data == WSCloseCode.MESSAGE_TOO_BIG

    assert ha_ws_server.closed


@pytest.mark.parametrize("auth_token", ["abc123", "bad"])
async def test_proxy_invalid_auth(
    api_client: TestClient, install_app_example: App, auth_token: str
):
    """Test invalid access token or app with no access."""
    install_app_example.persist[ATTR_ACCESS_TOKEN] = "abc123"
    websocket = await api_client.ws_connect("/core/websocket")
    auth_resp = await websocket.receive_json()
    assert auth_resp["type"] == "auth_required"
    await websocket.send_json({"type": "auth", "access_token": auth_token})

    auth_not_ok = await websocket.receive_json()
    assert auth_not_ok["type"] == "auth_invalid"
    assert auth_not_ok["message"] == "Invalid access"


async def test_proxy_auth_abort_log(
    api_client: TestClient,
    install_app_example: App,
    caplog: pytest.LogCaptureFixture,
):
    """Test WebSocket closed during authentication gets logged."""
    install_app_example.persist[ATTR_ACCESS_TOKEN] = "abc123"
    websocket = await api_client.ws_connect("/core/websocket")
    auth_resp = await websocket.receive_json()
    assert auth_resp["type"] == "auth_required"
    caplog.clear()
    with caplog.at_level(logging.ERROR):
        await websocket.close()
        assert (
            "Unexpected message during authentication for WebSocket API" in caplog.text
        )


async def test_websocket_transport_none(
    coresys,
    caplog: pytest.LogCaptureFixture,
):
    """Test WebSocket connection with transport None is handled gracefully."""
    # Get the API proxy instance from coresys
    api_proxy = APIProxy.__new__(APIProxy)
    api_proxy.coresys = coresys

    # Create a mock request with transport set to None to simulate connection loss
    mock_request = AsyncMock(spec=web.Request)
    mock_request.transport = None

    caplog.clear()
    with caplog.at_level(logging.WARNING):
        # This should raise HTTPBadRequest, not AssertionError
        with pytest.raises(web.HTTPBadRequest) as exc_info:
            await api_proxy.websocket(mock_request)

        # Verify the error reason
        assert exc_info.value.reason == "Connection closed"

        # Verify the warning was logged
        assert "WebSocket connection lost before upgrade" in caplog.text


@pytest.mark.parametrize("path", ["", "mock_path"])
async def test_api_proxy_get_request(
    api_client: TestClient,
    install_app_example: App,
    request: pytest.FixtureRequest,
    path: str,
):
    """Test the API proxy request using patch for make_request."""
    install_app_example.persist[ATTR_ACCESS_TOKEN] = "abc123"
    install_app_example.data["homeassistant_api"] = True

    request.param = "local_example"

    with patch.object(HomeAssistantAPI, "make_request") as make_request:
        # Mock the response from make_request
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.content_type = "application/json"
        mock_response.read.return_value = b"mocked response"
        make_request.return_value.__aenter__.return_value = mock_response

        response = await api_client.get(
            f"/core/api/{path}", headers={"Authorization": "Bearer abc123"}
        )

        assert make_request.call_args[0][0] == "get"
        assert make_request.call_args[0][1] == f"api/{path}"

        assert response.status == 200
        assert await response.text() == "mocked response"
        assert response.content_type == "application/json"


@pytest.mark.parametrize(
    "path", ["hassio_auth", "hassio_auth/password_reset", "hassio/app"]
)
async def test_api_proxy_blocks_core_hassio_endpoints(
    api_client: TestClient,
    install_app_example: App,
    request: pytest.FixtureRequest,
    path: str,
):
    """Test the proxy refuses to forward Core's Supervisor-only hassio endpoints.

    These run as the Supervisor user on Core; an add-on must not reach them
    through the proxy even if the security middleware blacklist is bypassed.
    """
    install_app_example.persist[ATTR_ACCESS_TOKEN] = "abc123"
    install_app_example.data["homeassistant_api"] = True

    request.param = "local_example"

    with patch.object(HomeAssistantAPI, "make_request") as make_request:
        response = await api_client.post(
            f"/core/api/{path}",
            headers={"Authorization": "Bearer abc123"},
            json={"username": "owner", "password": "attacker"},
        )

        assert response.status == 403
        make_request.assert_not_called()


@pytest.mark.parametrize(
    "path", ["config/automation/config/test_id", "services/light/turn_on"]
)
async def test_api_proxy_post_request(
    api_client: TestClient,
    install_app_example: App,
    request: pytest.FixtureRequest,
    path: str,
):
    """Test the API proxy POST request."""
    install_app_example.persist[ATTR_ACCESS_TOKEN] = "abc123"
    install_app_example.data["homeassistant_api"] = True

    request.param = "local_example"

    with patch.object(HomeAssistantAPI, "make_request") as make_request:
        # Mock the response from make_request
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.content_type = "application/json"
        mock_response.read.return_value = b'{"result": "ok"}'
        make_request.return_value.__aenter__.return_value = mock_response

        response = await api_client.post(
            f"/core/api/{path}",
            headers={"Authorization": "Bearer abc123"},
            json={"test": "data"},
        )

        assert make_request.call_args[0][0] == "post"
        assert make_request.call_args[0][1] == f"api/{path}"

        assert response.status == 200
        assert await response.text() == '{"result": "ok"}'
        assert response.content_type == "application/json"


@pytest.mark.parametrize(
    "path", ["config/automation/config/test_id", "states/light.test"]
)
async def test_api_proxy_delete_request(
    api_client: TestClient,
    install_app_example: App,
    request: pytest.FixtureRequest,
    path: str,
):
    """Test the API proxy DELETE request."""
    install_app_example.persist[ATTR_ACCESS_TOKEN] = "abc123"
    install_app_example.data["homeassistant_api"] = True

    request.param = "local_example"

    with patch.object(HomeAssistantAPI, "make_request") as make_request:
        # Mock the response from make_request
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.content_type = "application/json"
        mock_response.read.return_value = b'{"result": "ok"}'
        make_request.return_value.__aenter__.return_value = mock_response

        response = await api_client.delete(
            f"/core/api/{path}", headers={"Authorization": "Bearer abc123"}
        )

        assert make_request.call_args[0][0] == "delete"
        assert make_request.call_args[0][1] == f"api/{path}"

        assert response.status == 200
        assert await response.text() == '{"result": "ok"}'
        assert response.content_type == "application/json"


async def test_api_proxy_multipart_content_type_preserved(
    api_client: TestClient,
    install_app_example: App,
):
    """Test multipart Content-Type keeps its boundary parameter when proxied."""
    install_app_example.persist[ATTR_ACCESS_TOKEN] = "abc123"
    install_app_example.data["homeassistant_api"] = True

    with patch.object(HomeAssistantAPI, "make_request") as make_request:
        # Mock the response from make_request
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.content_type = "application/json"
        mock_response.read.return_value = b'{"result": "ok"}'
        make_request.return_value.__aenter__.return_value = mock_response

        boundary = "d1b1a3a3f0a94a5c9e3a"
        body = (
            f"--{boundary}\r\n"
            'Content-Disposition: form-data; name="file"; filename="test.png"\r\n'
            "Content-Type: image/png\r\n\r\n"
            "fakepng\r\n"
            f"--{boundary}--\r\n"
        ).encode()

        response = await api_client.post(
            "/core/api/media_source/local_source/upload",
            headers={
                "Authorization": "Bearer abc123",
                "Content-Type": f"multipart/form-data; boundary={boundary}",
            },
            data=body,
        )

        # The boundary parameter must survive the proxy — without it Core
        # cannot parse the multipart body ("boundary missed for Content-Type").
        assert (
            make_request.call_args[1]["content_type"]
            == f"multipart/form-data; boundary={boundary}"
        )
        assert response.status == 200


async def test_api_proxy_mcp_headers_forwarded(
    api_client: TestClient,
    install_app_example: App,
):
    """Test that MCP headers are forwarded to Home Assistant."""
    install_app_example.persist[ATTR_ACCESS_TOKEN] = "abc123"
    install_app_example.data["homeassistant_api"] = True

    with patch.object(HomeAssistantAPI, "make_request") as make_request:
        # Mock the response from make_request
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.content_type = "application/json"
        mock_response.read.return_value = b"mocked response"
        mock_response.headers = {"Mcp-Session-Id": "test-session-123"}
        make_request.return_value.__aenter__.return_value = mock_response

        response = await api_client.get(
            "/core/api/mcp",
            headers={
                "Authorization": "Bearer abc123",
                "Accept": "text/event-stream",
                "Last-Event-ID": "5",
                "Mcp-Session-Id": "test-session-123",
            },
        )

        # Verify headers were forwarded in the request
        assert make_request.call_args[1]["headers"]["Accept"] == "text/event-stream"
        assert make_request.call_args[1]["headers"]["Last-Event-ID"] == "5"
        assert (
            make_request.call_args[1]["headers"]["Mcp-Session-Id"] == "test-session-123"
        )

        # Verify response headers are preserved
        assert response.status == 200
        assert response.headers.get("Mcp-Session-Id") == "test-session-123"


async def test_api_proxy_streaming_response(
    api_client: TestClient,
    install_app_example: App,
):
    """Test that streaming responses (text/event-stream) are handled properly."""
    install_app_example.persist[ATTR_ACCESS_TOKEN] = "abc123"
    install_app_example.data["homeassistant_api"] = True

    async def mock_content_iter():
        """Mock async iterator for streaming content."""
        yield b"data: event1\n\n"
        yield b"data: event2\n\n"
        yield b"data: event3\n\n"

    with patch.object(HomeAssistantAPI, "make_request") as make_request:
        # Mock the response from make_request
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.content_type = "text/event-stream"
        mock_response.headers = {
            "Cache-Control": "no-cache",
            "Mcp-Session-Id": "session-456",
        }
        mock_response.content = mock_content_iter()
        make_request.return_value.__aenter__.return_value = mock_response

        response = await api_client.get(
            "/core/api/mcp",
            headers={
                "Authorization": "Bearer abc123",
                "Accept": "text/event-stream",
            },
        )

        # Verify it's a streaming response
        assert response.status == 200
        assert response.content_type == "text/event-stream"
        assert response.headers.get("X-Accel-Buffering") == "no"
        assert response.headers.get("Mcp-Session-Id") == "session-456"

        # Read the streamed content
        content = await response.read()
        assert b"data: event1\n\n" in content
        assert b"data: event2\n\n" in content
        assert b"data: event3\n\n" in content


async def test_api_proxy_streaming_response_client_payload_error(
    api_client: TestClient,
    install_app_example: App,
):
    """Test that client payload errors during streaming are handled gracefully."""
    install_app_example.persist[ATTR_ACCESS_TOKEN] = "abc123"
    install_app_example.data["homeassistant_api"] = True

    async def mock_content_iter_error():
        yield b"data: event1\n\n"
        raise ClientPayloadError("boom")

    with patch.object(HomeAssistantAPI, "make_request") as make_request:
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.content_type = "text/event-stream"
        mock_response.headers = {
            "Cache-Control": "no-cache",
            "Mcp-Session-Id": "session-789",
        }
        mock_response.content = mock_content_iter_error()
        make_request.return_value.__aenter__.return_value = mock_response

        response = await api_client.get(
            "/core/api/mcp",
            headers={
                "Authorization": "Bearer abc123",
                "Accept": "text/event-stream",
            },
        )

        assert response.status == 200
        assert response.content_type == "text/event-stream"
        assert response.headers.get("X-Accel-Buffering") == "no"
        assert response.headers.get("Mcp-Session-Id") == "session-789"

        content = await response.read()
        assert b"data: event1\n\n" in content