bambuddy/backend/tests/integration/test_api_key_id_reuse.py

"""A deleted API key leaves nothing that the next key given its id would take.

SQLite gives a deleted key's id to the next key created. Open WebSockets opened
with the key, and the camera and WebSocket tokens it minted, are tied to
the key by id: they took the new key's printers.
"""

from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from fastapi import WebSocketDisconnect
from httpx import AsyncClient
from sqlalchemy import select, text

from backend.app.core.websocket import ws_manager
from backend.app.models.api_key import APIKey
from backend.app.models.auth_ephemeral import AuthEphemeralToken

pytestmark = [pytest.mark.integration, pytest.mark.asyncio]


@pytest.fixture
async def admin(async_client: AsyncClient) -> dict:
    await async_client.post(
        "/api/v1/auth/setup",
        json={"auth_enabled": True, "admin_username": "keyadmin", "admin_password": "AdminPass1!"},
    )
    login = await async_client.post("/api/v1/auth/login", json={"username": "keyadmin", "password": "AdminPass1!"})
    return {"Authorization": f"Bearer {login.json()['access_token']}"}


async def _key(db, user_id: int | None = None) -> int:
    key = APIKey(name="k", key_hash="h", key_prefix="bb_k", printer_ids=[1], user_id=user_id)
    db.add(key)
    await db.commit()
    return key.id


async def _tokens(db, key_id: int) -> None:
    expires = datetime.now(timezone.utc) + timedelta(hours=1)
    db.add_all(
        [
            AuthEphemeralToken(
                token=f"{kind}-{key_id}", token_type=kind, username="", api_key_id=key_id, expires_at=expires
            )
            for kind in ("camera_stream", "websocket")  # media tokens never record the key
        ]
    )
    await db.commit()


async def _token_keys(db) -> list:
    rows = await db.execute(
        select(AuthEphemeralToken.api_key_id)
        .where(AuthEphemeralToken.api_key_id.is_not(None))
        .execution_options(populate_existing=True)
    )
    return sorted(rows.scalars().all())


def _socket(api_key_id: int | None) -> MagicMock:
    socket = MagicMock()
    socket.state = SimpleNamespace(
        bambuddy_scope_principal=("", api_key_id), bambuddy_principal_user_id=None, bambuddy_principal=""
    )
    socket.accept = AsyncMock()
    socket.close = AsyncMock()
    socket.send_json = AsyncMock()
    socket.receive_json = AsyncMock(side_effect=WebSocketDisconnect())
    return socket


async def test_deleting_a_key_closes_its_sockets_and_drops_its_tokens(async_client: AsyncClient, admin, db_session):
    gone, kept = await _key(db_session), await _key(db_session)
    await _tokens(db_session, gone)
    await _tokens(db_session, kept)
    theirs, other = _socket(gone), _socket(kept)
    ws_manager.active_connections.extend([theirs, other])
    try:
        response = await async_client.delete(f"/api/v1/api-keys/{gone}", headers=admin)

        assert response.status_code == 200, response.text
        assert await _token_keys(db_session) == [kept] * 2
        theirs.close.assert_awaited_once_with(code=4401)
        assert theirs not in ws_manager.active_connections
        other.close.assert_not_awaited()
    finally:
        for socket in (theirs, other):
            if socket in ws_manager.active_connections:
                ws_manager.active_connections.remove(socket)


async def test_deleting_a_user_does_the_same_for_their_keys(async_client: AsyncClient, admin, db_session):
    created = await async_client.post(
        "/api/v1/users/", headers=admin, json={"username": "keyowner", "password": "Userpass1!", "role": "user"}
    )
    owner = created.json()["id"]
    key = await _key(db_session, user_id=owner)
    await _tokens(db_session, key)
    theirs = _socket(key)
    ws_manager.active_connections.append(theirs)
    try:
        response = await async_client.delete(f"/api/v1/users/{owner}", headers=admin)

        assert response.status_code == 204, response.text
        assert await _token_keys(db_session) == []
        theirs.close.assert_awaited_once_with(code=4401)
    finally:
        if theirs in ws_manager.active_connections:
            ws_manager.active_connections.remove(theirs)


async def _connect(key_id: int, connect=None) -> MagicMock:
    from backend.app.api.routes import websocket as ws_route
    from backend.app.core import database
    from backend.app.core.printer_scope import ALL_PRINTERS

    socket = _socket(None)
    socket.state = SimpleNamespace()
    with (
        patch.object(ws_route, "is_auth_enabled", AsyncMock(return_value=True)),
        patch.object(ws_route, "verify_websocket_token_principal", AsyncMock(return_value=("", key_id))),
        patch.object(ws_route, "principal_printer_scope", AsyncMock(return_value=ALL_PRINTERS)),
        patch.object(ws_route, "async_session", database.async_session),
        patch.object(ws_manager, "connect", connect or ws_manager.connect),
    ):
        await ws_route.websocket_endpoint(socket, token="t")
    return socket


async def test_a_websocket_of_a_deleted_key_is_refused(async_client: AsyncClient):
    socket = await _connect(9999)

    socket.close.assert_awaited_once_with(code=4401)
    socket.accept.assert_not_awaited()


async def test_a_key_deleted_while_its_websocket_joins_closes_it(async_client: AsyncClient, db_session):
    key = await _key(db_session)
    joined = ws_manager.connect

    async def join_while_deleted(websocket):
        await db_session.execute(text("DELETE FROM api_keys WHERE id = :id"), {"id": key})
        await db_session.commit()
        await ws_manager.close_api_keys([key])  # what the delete does after its commit
        await joined(websocket)

    socket = await _connect(key, join_while_deleted)

    socket.close.assert_awaited_once_with(code=4401)
    assert socket not in ws_manager.active_connections


async def test_the_startup_cleanup_drops_tokens_of_keys_that_are_gone(db_session):
    from backend.app.core.database import _drop_ephemeral_tokens_of_missing_owners

    kept = await _key(db_session)
    await _tokens(db_session, kept)
    await _tokens(db_session, kept + 50)  # deleted before deletes took them along

    await _drop_ephemeral_tokens_of_missing_owners(await db_session.connection())
    await db_session.commit()

    assert await _token_keys(db_session) == [kept] * 2


async def test_a_deleted_keys_message_rate_goes_too(async_client: AsyncClient, admin, db_session):
    """The next key given the id would start out throttled."""
    from backend.app.api.routes import notifications

    key = await _key(db_session)
    notifications._app_message_times[f"key:{key}"].extend([1.0] * notifications.APP_MESSAGE_LIMIT)
    try:
        response = await async_client.delete(f"/api/v1/api-keys/{key}", headers=admin)

        assert response.status_code == 200, response.text
        assert f"key:{key}" not in notifications._app_message_times
    finally:
        notifications._app_message_times.pop(f"key:{key}", None)


@pytest.mark.parametrize("replaced", [False, True])
async def test_a_scope_refresh_closes_a_socket_whose_key_is_gone(async_client: AsyncClient, db_session, replaced):
    """Rotated from the command line, or an id that came free later: the socket
    was re-scoped by id, taking the printers of whichever key had it next."""
    from backend.app.core import database

    await async_client.post(
        "/api/v1/auth/setup",
        json={"auth_enabled": True, "admin_username": "keyadmin2", "admin_password": "AdminPass1!"},
    )
    key = await _key(db_session)
    prefix = (await db_session.execute(select(APIKey.key_prefix).where(APIKey.id == key))).scalar_one()
    socket = _socket(key)
    socket.state.bambuddy_api_key_prefix = prefix
    ws_manager.active_connections.append(socket)
    await db_session.execute(text("DELETE FROM api_keys WHERE id = :id"), {"id": key})
    await db_session.commit()
    if replaced:  # a new key given the freed id
        db_session.add(APIKey(id=key, name="new", key_hash="h2", key_prefix="bb_new", printer_ids=None))
        await db_session.commit()
    try:
        with patch("backend.app.core.database.async_session", database.async_session):
            await ws_manager.refresh_printer_scopes()

        socket.close.assert_awaited_once_with(code=4401)
        assert socket not in ws_manager.active_connections
    finally:
        if socket in ws_manager.active_connections:
            ws_manager.active_connections.remove(socket)


async def test_a_deleted_users_message_rate_goes_too(async_client: AsyncClient, admin):
    """The next user given the id would start out throttled."""
    from backend.app.api.routes import notifications

    created = await async_client.post(
        "/api/v1/users/", headers=admin, json={"username": "sender", "password": "Userpass1!", "role": "user"}
    )
    user = created.json()["id"]
    notifications._app_message_times[f"user:{user}"].extend([1.0] * notifications.APP_MESSAGE_LIMIT)
    try:
        response = await async_client.delete(f"/api/v1/users/{user}", headers=admin)

        assert response.status_code == 204, response.text
        assert f"user:{user}" not in notifications._app_message_times
    finally:
        notifications._app_message_times.pop(f"user:{user}", None)