core/homeassistant/components/portainer/services.py

"""Services for the Portainer integration."""

from datetime import timedelta

from pyportainer import (
    PortainerAuthenticationError,
    PortainerConnectionError,
    PortainerTimeoutError,
)
import voluptuous as vol

from homeassistant.const import ATTR_DEVICE_ID
from homeassistant.core import HomeAssistant, ServiceCall, callback
from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
from homeassistant.helpers import (
    config_validation as cv,
    device_registry as dr,
    service,
)

from .const import DOMAIN
from .coordinator import PortainerConfigEntry

ATTR_DATE_UNTIL = "until"
ATTR_DANGLING = "dangling"
ATTR_TIMEOUT = "timeout"
ATTR_PULL_IMAGE = "pull_image"
ATTR_CONTAINER_DEVICE_ID = "container_device_id"

SERVICE_PRUNE_IMAGES = "prune_images"
SERVICE_PRUNE_IMAGES_SCHEMA = vol.Schema(
    {
        vol.Required(ATTR_DEVICE_ID): cv.string,
        vol.Optional(ATTR_DATE_UNTIL): vol.All(
            cv.time_period, vol.Range(min=timedelta(minutes=1))
        ),
        vol.Optional(ATTR_DANGLING): cv.boolean,
    },
)

SERVICE_RECREATE_CONTAINER = "recreate_container"
SERVICE_RECREATE_CONTAINER_SCHEMA = vol.Schema(
    {
        vol.Required(ATTR_CONTAINER_DEVICE_ID): cv.string,
        vol.Optional(ATTR_TIMEOUT): vol.All(
            cv.time_period, vol.Range(min=timedelta(minutes=1))
        ),
        vol.Optional(ATTR_PULL_IMAGE): cv.boolean,
    }
)


@callback
def _async_get_device(call: ServiceCall, device_id: str) -> dr.DeviceEntry:
    """Get a device entry from a device ID."""
    device_reg = dr.async_get(call.hass)
    if (device := device_reg.async_get(device_id, include_child_devices=False)) is None:
        raise ServiceValidationError(
            translation_domain=DOMAIN,
            translation_key="invalid_target",
        )
    return device


@callback
def _async_get_entry_from_device(
    call: ServiceCall, device: dr.DeviceEntry
) -> PortainerConfigEntry:
    """Resolve and validate the Portainer config entry for a device."""
    for entry in call.hass.config_entries.async_entries(DOMAIN):
        if entry.entry_id in device.config_entries:
            return service.async_get_config_entry(call.hass, DOMAIN, entry.entry_id)

    raise ServiceValidationError(
        translation_domain=DOMAIN,
        translation_key="invalid_target",
    )


@callback
def _async_get_endpoint_id(
    device: dr.DeviceEntry,
    config_entry: PortainerConfigEntry,
) -> int:
    """Get the endpoint ID from a device entry."""
    coordinator = config_entry.runtime_data

    for data in coordinator.data.values():
        if (
            DOMAIN,
            f"{config_entry.entry_id}_{data.endpoint.id}",
        ) in device.identifiers:
            return data.endpoint.id

    raise ServiceValidationError(
        translation_domain=DOMAIN,
        translation_key="invalid_target",
    )


@callback
def _async_get_container_and_endpoint_ids(
    device: dr.DeviceEntry,
    config_entry: PortainerConfigEntry,
) -> tuple[int, str]:
    """Get the endpoint ID and container ID from a container device entry."""
    coordinator = config_entry.runtime_data

    for data in coordinator.data.values():
        for container_name, container_data in data.containers.items():
            if (
                DOMAIN,
                f"{config_entry.entry_id}_{data.endpoint.id}_{container_name}",
            ) in device.identifiers:
                return data.endpoint.id, container_data.container.id

    raise ServiceValidationError(
        translation_domain=DOMAIN,
        translation_key="invalid_target",
    )


async def prune_images(call: ServiceCall) -> None:
    """Prune unused images in Portainer, with more controls."""
    device = _async_get_device(call, call.data[ATTR_DEVICE_ID])
    config_entry = _async_get_entry_from_device(call, device)
    coordinator = config_entry.runtime_data
    endpoint_id = _async_get_endpoint_id(device, config_entry)

    try:
        await coordinator.portainer.images_prune(
            endpoint_id=endpoint_id,
            until=call.data.get(ATTR_DATE_UNTIL),
            dangling=call.data.get(ATTR_DANGLING, False),
        )
    except PortainerAuthenticationError as err:
        raise HomeAssistantError(
            translation_domain=DOMAIN,
            translation_key="invalid_auth",
        ) from err
    except PortainerConnectionError as err:
        raise HomeAssistantError(
            translation_domain=DOMAIN,
            translation_key="cannot_connect",
        ) from err
    except PortainerTimeoutError as err:
        raise HomeAssistantError(
            translation_domain=DOMAIN,
            translation_key="timeout_connect",
        ) from err


async def recreate_container(call: ServiceCall) -> None:
    """Recreate a container in Portainer, with more controls."""
    device = _async_get_device(call, call.data[ATTR_CONTAINER_DEVICE_ID])
    config_entry = _async_get_entry_from_device(call, device)
    coordinator = config_entry.runtime_data
    endpoint_id, container_id = _async_get_container_and_endpoint_ids(
        device, config_entry
    )
    timeout: timedelta | None = call.data.get(ATTR_TIMEOUT)

    try:
        await coordinator.portainer.container_recreate(
            endpoint_id=endpoint_id,
            container_id=container_id,
            **({"timeout": timeout} if timeout is not None else {}),
            pull_image=call.data.get(ATTR_PULL_IMAGE, False),
        )
    except PortainerAuthenticationError as err:
        raise HomeAssistantError(
            translation_domain=DOMAIN,
            translation_key="invalid_auth",
        ) from err
    except PortainerConnectionError as err:
        raise HomeAssistantError(
            translation_domain=DOMAIN,
            translation_key="cannot_connect",
        ) from err
    except PortainerTimeoutError as err:
        raise HomeAssistantError(
            translation_domain=DOMAIN,
            translation_key="timeout_connect",
        ) from err

    await coordinator.async_request_refresh()


async def async_setup_services(hass: HomeAssistant) -> None:
    """Set up services."""

    hass.services.async_register(
        DOMAIN,
        SERVICE_PRUNE_IMAGES,
        prune_images,
        SERVICE_PRUNE_IMAGES_SCHEMA,
    )

    service.async_register_admin_service(
        hass,
        DOMAIN,
        SERVICE_RECREATE_CONTAINER,
        recreate_container,
        SERVICE_RECREATE_CONTAINER_SCHEMA,
    )