205 lines
6.6 KiB
Python
205 lines
6.6 KiB
Python
"""The Picture integration."""
|
|
import asyncio
|
|
import logging
|
|
import pathlib
|
|
import secrets
|
|
import shutil
|
|
import typing
|
|
|
|
from PIL import Image, ImageOps, UnidentifiedImageError
|
|
from aiohttp import hdrs, web
|
|
from aiohttp.web_request import FileField
|
|
import voluptuous as vol
|
|
|
|
from homeassistant.components.http.view import HomeAssistantView
|
|
from homeassistant.const import CONF_ID
|
|
from homeassistant.core import HomeAssistant, callback
|
|
from homeassistant.helpers import collection
|
|
from homeassistant.helpers.storage import Store
|
|
import homeassistant.util.dt as dt_util
|
|
|
|
from .const import DOMAIN
|
|
|
|
_LOGGER = logging.getLogger(__name__)
|
|
STORAGE_KEY = DOMAIN
|
|
STORAGE_VERSION = 1
|
|
VALID_SIZES = {256, 512}
|
|
MAX_SIZE = 1024 * 1024 * 10
|
|
|
|
CREATE_FIELDS = {
|
|
vol.Required("file"): FileField,
|
|
}
|
|
|
|
UPDATE_FIELDS = {
|
|
vol.Optional("name"): vol.All(str, vol.Length(min=1)),
|
|
}
|
|
|
|
|
|
async def async_setup(hass: HomeAssistant, config: dict):
|
|
"""Set up the Image integration."""
|
|
image_dir = pathlib.Path(hass.config.path(DOMAIN))
|
|
hass.data[DOMAIN] = storage_collection = ImageStorageCollection(hass, image_dir)
|
|
await storage_collection.async_load()
|
|
collection.StorageCollectionWebsocket(
|
|
storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS,
|
|
).async_setup(hass, create_create=False)
|
|
|
|
hass.http.register_view(ImageUploadView)
|
|
hass.http.register_view(ImageServeView(image_dir, storage_collection))
|
|
return True
|
|
|
|
|
|
class ImageStorageCollection(collection.StorageCollection):
|
|
"""Image collection stored in storage."""
|
|
|
|
CREATE_SCHEMA = vol.Schema(CREATE_FIELDS)
|
|
UPDATE_SCHEMA = vol.Schema(UPDATE_FIELDS)
|
|
|
|
def __init__(self, hass: HomeAssistant, image_dir: pathlib.Path) -> None:
|
|
"""Initialize media storage collection."""
|
|
super().__init__(
|
|
Store(hass, STORAGE_VERSION, STORAGE_KEY),
|
|
logging.getLogger(f"{__name__}.storage_collection"),
|
|
)
|
|
self.async_add_listener(self._change_listener)
|
|
self.image_dir = image_dir
|
|
|
|
async def _process_create_data(self, data: typing.Dict) -> typing.Dict:
|
|
"""Validate the config is valid."""
|
|
data = self.CREATE_SCHEMA(dict(data))
|
|
uploaded_file: FileField = data["file"]
|
|
|
|
if not uploaded_file.content_type.startswith("image/"):
|
|
raise vol.Invalid("Only images are allowed")
|
|
|
|
data[CONF_ID] = secrets.token_hex(16)
|
|
data["filesize"] = await self.hass.async_add_executor_job(self._move_data, data)
|
|
|
|
data["content_type"] = uploaded_file.content_type
|
|
data["name"] = uploaded_file.filename
|
|
data["uploaded_at"] = dt_util.utcnow().isoformat()
|
|
|
|
return data
|
|
|
|
def _move_data(self, data):
|
|
"""Move data."""
|
|
uploaded_file: FileField = data.pop("file")
|
|
|
|
# Verify we can read the image
|
|
try:
|
|
image = Image.open(uploaded_file.file)
|
|
except UnidentifiedImageError:
|
|
raise vol.Invalid("Unable to identify image file")
|
|
|
|
# Reset content
|
|
uploaded_file.file.seek(0)
|
|
|
|
media_folder: pathlib.Path = (self.image_dir / data[CONF_ID])
|
|
media_folder.mkdir(parents=True)
|
|
|
|
media_file = media_folder / "original"
|
|
|
|
# Raises if path is no longer relative to the media dir
|
|
media_file.relative_to(media_folder)
|
|
|
|
_LOGGER.debug("Storing file %s", media_file)
|
|
|
|
with media_file.open("wb") as target:
|
|
shutil.copyfileobj(uploaded_file.file, target)
|
|
|
|
image.close()
|
|
|
|
return media_file.stat().st_size
|
|
|
|
@callback
|
|
def _get_suggested_id(self, info: typing.Dict) -> str:
|
|
"""Suggest an ID based on the config."""
|
|
return info[CONF_ID]
|
|
|
|
async def _update_data(self, data: dict, update_data: typing.Dict) -> typing.Dict:
|
|
"""Return a new updated data object."""
|
|
return {**data, **self.UPDATE_SCHEMA(update_data)}
|
|
|
|
async def _change_listener(self, change_type, item_id, data):
|
|
"""Handle change."""
|
|
if change_type != collection.CHANGE_REMOVED:
|
|
return
|
|
|
|
await self.hass.async_add_executor_job(shutil.rmtree, self.image_dir / item_id)
|
|
|
|
|
|
class ImageUploadView(HomeAssistantView):
|
|
"""View to upload images."""
|
|
|
|
url = "/api/image/upload"
|
|
name = "api:image:upload"
|
|
|
|
async def post(self, request):
|
|
"""Handle upload."""
|
|
# Increase max payload
|
|
request._client_max_size = MAX_SIZE # pylint: disable=protected-access
|
|
|
|
data = await request.post()
|
|
item = await request.app["hass"].data[DOMAIN].async_create_item(data)
|
|
return self.json(item)
|
|
|
|
|
|
class ImageServeView(HomeAssistantView):
|
|
"""View to download images."""
|
|
|
|
url = "/api/image/serve/{image_id}/{filename}"
|
|
name = "api:image:serve"
|
|
requires_auth = False
|
|
|
|
def __init__(
|
|
self, image_folder: pathlib.Path, image_collection: ImageStorageCollection
|
|
):
|
|
"""Initialize image serve view."""
|
|
self.transform_lock = asyncio.Lock()
|
|
self.image_folder = image_folder
|
|
self.image_collection = image_collection
|
|
|
|
async def get(self, request: web.Request, image_id: str, filename: str):
|
|
"""Serve image."""
|
|
image_size = filename.split("-", 1)[0]
|
|
try:
|
|
parts = image_size.split("x", 1)
|
|
width = int(parts[0])
|
|
height = int(parts[1])
|
|
except (ValueError, IndexError):
|
|
raise web.HTTPBadRequest
|
|
|
|
if not width or width != height or width not in VALID_SIZES:
|
|
raise web.HTTPBadRequest
|
|
|
|
image_info = self.image_collection.data.get(image_id)
|
|
|
|
if image_info is None:
|
|
raise web.HTTPNotFound()
|
|
|
|
hass = request.app["hass"]
|
|
target_file = self.image_folder / image_id / f"{width}x{height}"
|
|
|
|
if not target_file.is_file():
|
|
async with self.transform_lock:
|
|
# Another check in case another request already finished it while waiting
|
|
if not target_file.is_file():
|
|
await hass.async_add_executor_job(
|
|
_generate_thumbnail,
|
|
self.image_folder / image_id / "original",
|
|
image_info["content_type"],
|
|
target_file,
|
|
(width, height),
|
|
)
|
|
|
|
return web.FileResponse(
|
|
target_file, headers={hdrs.CONTENT_TYPE: image_info["content_type"]}
|
|
)
|
|
|
|
|
|
def _generate_thumbnail(original_path, content_type, target_path, target_size):
|
|
"""Generate a size."""
|
|
image = ImageOps.exif_transpose(Image.open(original_path))
|
|
image.thumbnail(target_size)
|
|
image.save(target_path, format=content_type.split("/", 1)[1])
|