Add snapshot service to image entity

homeassistant/components/image/__init__.py

@@ -8,37 +8,59 @@
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 import logging
+import os
 from random import SystemRandom
 from typing import Final, final
 
 from aiohttp import hdrs, web
 import httpx
 from propcache import cached_property
+import voluptuous as vol
 
 from homeassistant.components.http import KEY_AUTHENTICATED, KEY_HASS, HomeAssistantView
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONTENT_TYPE_MULTIPART, EVENT_HOMEASSISTANT_STOP
-from homeassistant.core import Event, EventStateChangedData, HomeAssistant, callback
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    CONTENT_TYPE_MULTIPART,
+    EVENT_HOMEASSISTANT_STOP,
+)
+from homeassistant.core import (
+    Event,
+    EventStateChangedData,
+    HomeAssistant,
+    ServiceCall,
+    callback,
+)
 from homeassistant.exceptions import HomeAssistantError
-from homeassistant.helpers import config_validation as cv
+import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import Entity, EntityDescription
 from homeassistant.helpers.entity_component import EntityComponent
 from homeassistant.helpers.event import (
     async_track_state_change_event,
     async_track_time_interval,
 )
 from homeassistant.helpers.httpx_client import get_async_client
-from homeassistant.helpers.typing import UNDEFINED, ConfigType, UndefinedType
+from homeassistant.helpers.template import Template
+from homeassistant.helpers.typing import (
+    UNDEFINED,
+    ConfigType,
+    UndefinedType,
+    VolDictType,
+)
 
 from .const import DATA_COMPONENT, DOMAIN, IMAGE_TIMEOUT
 
 _LOGGER = logging.getLogger(__name__)
 
+SERVICE_SNAPSHOT: Final = "snapshot"
+
 ENTITY_ID_FORMAT: Final = DOMAIN + ".{}"
 PLATFORM_SCHEMA = cv.PLATFORM_SCHEMA
 PLATFORM_SCHEMA_BASE = cv.PLATFORM_SCHEMA_BASE
 SCAN_INTERVAL: Final = timedelta(seconds=30)
 
+ATTR_FILENAME: Final = "filename"
+
 DEFAULT_CONTENT_TYPE: Final = "image/jpeg"
 ENTITY_IMAGE_URL: Final = "/api/image_proxy/{0}?token={1}"
 
@@ -47,10 +69,13 @@
 
 GET_IMAGE_TIMEOUT: Final = 10
 
+
 FRAME_BOUNDARY = "frame-boundary"
 FRAME_SEPARATOR = bytes(f"\r\n--{FRAME_BOUNDARY}\r\n", "utf-8")
 LAST_FRAME_MARKER = bytes(f"\r\n--{FRAME_BOUNDARY}--\r\n", "utf-8")
 
+IMAGE_SERVICE_SNAPSHOT: VolDictType = {vol.Required(ATTR_FILENAME): cv.template}
+
 
 class ImageEntityDescription(EntityDescription, frozen_or_thawed=True):
     """A class that describes image entities."""
@@ -115,6 +140,10 @@ def unsub_track_time_interval(_event: Event) -> None:
 
     hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, unsub_track_time_interval)
 
+    component.async_register_entity_service(
+        SERVICE_SNAPSHOT, IMAGE_SERVICE_SNAPSHOT, async_handle_snapshot_service
+    )
+
     return True
 
 
@@ -380,3 +409,37 @@ async def handle(
     ) -> web.StreamResponse:
         """Serve image stream."""
         return await async_get_still_stream(request, image_entity)
+
+
+async def async_handle_snapshot_service(
+    image: ImageEntity, service_call: ServiceCall
+) -> None:
+    """Handle snapshot services calls."""
+    hass = image.hass
+    filename: Template = service_call.data[ATTR_FILENAME]
+    filename.hass = hass
+
+    snapshot_file = filename.async_render(variables={ATTR_ENTITY_ID: image.entity_id})
+
+    # check if we allow to access to that file
+    if not hass.config.is_allowed_path(snapshot_file):
+        raise HomeAssistantError(
+            f"Cannot write `{snapshot_file}`, no access to path; `allowlist_external_dirs` may need to be adjusted in `configuration.yaml`"
+        )
+
+    async with asyncio.timeout(IMAGE_TIMEOUT):
+        image_data = await image.async_image()
+
+    if image_data is None:
+        return
+
+    def _write_image(to_file: str, image_data: bytes) -> None:
+        """Executor helper to write image."""
+        os.makedirs(os.path.dirname(to_file), exist_ok=True)
+        with open(to_file, "wb") as img_file:
+            img_file.write(image_data)
+
+    try:
+        await hass.async_add_executor_job(_write_image, snapshot_file, image_data)
+    except OSError as err:
+        raise HomeAssistantError("Can't write image to file") from err

homeassistant/components/image/icons.json

@@ -3,5 +3,10 @@
     "_": {
       "default": "mdi:image"
     }
+  },
+  "services": {
+    "snapshot": {
+      "service": "mdi:camera"
+    }
   }
 }

homeassistant/components/image/services.yaml

@@ -0,0 +1,12 @@
+# Describes the format for available image services
+
+snapshot:
+  target:
+    entity:
+      domain: image
+  fields:
+    filename:
+      required: true
+      example: "/tmp/snapshot_{{ entity_id }}.jpg"
+      selector:
+        text:

homeassistant/components/image/strings.json

@@ -4,5 +4,17 @@
     "_": {
       "name": "[%key:component::image::title%]"
     }
+  },
+  "services": {
+    "snapshot": {
+      "name": "Take snapshot",
+      "description": "Takes a snapshot from an image.",
+      "fields": {
+        "filename": {
+          "name": "Filename",
+          "description": "Template of a filename. Variable available is `entity_id`."
+        }
+      }
+    }
   }
 }

tests/components/image/conftest.py

@@ -88,6 +88,16 @@ async def async_image(self) -> bytes | None:
         return b"Test"
 
 
+class MockImageNoDataEntity(image.ImageEntity):
+    """Mock image entity."""
+
+    _attr_name = "Test"
+
+    async def async_image(self) -> bytes | None:
+        """Return bytes of image."""
+        return None
+
+
 class MockImageSyncEntity(image.ImageEntity):
     """Mock image entity."""
 

tests/components/image/test_init.py

@@ -3,7 +3,7 @@
 from datetime import datetime
 from http import HTTPStatus
 import ssl
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, mock_open, patch
 
 from aiohttp import hdrs
 from freezegun.api import FrozenDateTimeFactory
@@ -13,13 +13,16 @@
 
 from homeassistant.components import image
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import ATTR_ENTITY_ID
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.setup import async_setup_component
 
 from .conftest import (
     MockImageEntity,
     MockImageEntityCapitalContentType,
     MockImageEntityInvalidContentType,
+    MockImageNoDataEntity,
     MockImageNoStateEntity,
     MockImagePlatform,
     MockImageSyncEntity,
@@ -381,3 +384,143 @@ async def _wrap_async_get_still_stream(*args, **kwargs):
             await hass.async_block_till_done()
 
     await close_future
+
+
+async def test_snapshot_service(hass: HomeAssistant) -> None:
+    """Test snapshot service."""
+    mopen = mock_open()
+    mock_integration(hass, MockModule(domain="test"))
+    mock_platform(hass, "test.image", MockImagePlatform([MockImageSyncEntity(hass)]))
+    assert await async_setup_component(
+        hass, image.DOMAIN, {"image": {"platform": "test"}}
+    )
+    await hass.async_block_till_done()
+
+    with (
+        patch("homeassistant.components.image.open", mopen, create=True),
+        patch("homeassistant.components.image.os.makedirs"),
+        patch.object(hass.config, "is_allowed_path", return_value=True),
+    ):
+        await hass.services.async_call(
+            image.DOMAIN,
+            image.SERVICE_SNAPSHOT,
+            {
+                ATTR_ENTITY_ID: "image.test",
+                image.ATTR_FILENAME: "/test/snapshot.jpg",
+            },
+            blocking=True,
+        )
+
+        mock_write = mopen().write
+
+        assert len(mock_write.mock_calls) == 1
+        assert mock_write.mock_calls[0][1][0] == b"Test"
+
+
+async def test_snapshot_service_with_template(hass: HomeAssistant) -> None:
+    """Test snapshot service."""
+    mopen = mock_open()
+    mock_integration(hass, MockModule(domain="test"))
+    mock_platform(hass, "test.image", MockImagePlatform([MockImageSyncEntity(hass)]))
+    assert await async_setup_component(
+        hass, image.DOMAIN, {"image": {"platform": "test"}}
+    )
+    await hass.async_block_till_done()
+
+    with (
+        patch("homeassistant.components.image.open", mopen, create=True),
+        patch("homeassistant.components.image.os.makedirs"),
+        patch.object(hass.config, "is_allowed_path", return_value=True),
+    ):
+        await hass.services.async_call(
+            image.DOMAIN,
+            image.SERVICE_SNAPSHOT,
+            {
+                ATTR_ENTITY_ID: "image.test",
+                image.ATTR_FILENAME: "/test/snapshot_{{ entity_id }}.jpg",
+            },
+            blocking=True,
+        )
+
+        mock_write = mopen().write
+
+        assert len(mock_write.mock_calls) == 1
+        assert mock_write.mock_calls[0][1][0] == b"Test"
+
+
+async def test_snapshot_service_no_image(hass: HomeAssistant) -> None:
+    """Test snapshot service with no image."""
+    mopen = mock_open()
+    mock_integration(hass, MockModule(domain="test"))
+    mock_platform(hass, "test.image", MockImagePlatform([MockImageNoDataEntity(hass)]))
+    assert await async_setup_component(
+        hass, image.DOMAIN, {"image": {"platform": "test"}}
+    )
+    await hass.async_block_till_done()
+
+    with (
+        patch("homeassistant.components.image.open", mopen, create=True),
+        patch(
+            "homeassistant.components.image.os.makedirs",
+        ),
+        patch.object(hass.config, "is_allowed_path", return_value=True),
+    ):
+        await hass.services.async_call(
+            image.DOMAIN,
+            image.SERVICE_SNAPSHOT,
+            {
+                ATTR_ENTITY_ID: "image.test",
+                image.ATTR_FILENAME: "/test/snapshot.jpg",
+            },
+            blocking=True,
+        )
+
+        mock_write = mopen().write
+
+        assert len(mock_write.mock_calls) == 0
+
+
+async def test_snapshot_service_not_allowed_path(hass: HomeAssistant) -> None:
+    """Test snapshot service with a not allowed path."""
+    mock_integration(hass, MockModule(domain="test"))
+    mock_platform(hass, "test.image", MockImagePlatform([MockURLImageEntity(hass)]))
+    assert await async_setup_component(
+        hass, image.DOMAIN, {"image": {"platform": "test"}}
+    )
+    await hass.async_block_till_done()
+
+    with pytest.raises(HomeAssistantError, match="/test/snapshot.jpg"):
+        await hass.services.async_call(
+            image.DOMAIN,
+            image.SERVICE_SNAPSHOT,
+            {
+                ATTR_ENTITY_ID: "image.test",
+                image.ATTR_FILENAME: "/test/snapshot.jpg",
+            },
+            blocking=True,
+        )
+
+
+async def test_snapshot_service_os_error(hass: HomeAssistant) -> None:
+    """Test snapshot service with os error."""
+    mock_integration(hass, MockModule(domain="test"))
+    mock_platform(hass, "test.image", MockImagePlatform([MockImageSyncEntity(hass)]))
+    assert await async_setup_component(
+        hass, image.DOMAIN, {"image": {"platform": "test"}}
+    )
+    await hass.async_block_till_done()
+
+    with (
+        patch.object(hass.config, "is_allowed_path", return_value=True),
+        patch("os.makedirs", side_effect=OSError),
+        pytest.raises(HomeAssistantError),
+    ):
+        await hass.services.async_call(
+            image.DOMAIN,
+            image.SERVICE_SNAPSHOT,
+            {
+                ATTR_ENTITY_ID: "image.test",
+                image.ATTR_FILENAME: "/test/snapshot.jpg",
+            },
+            blocking=True,
+        )