ha-nvidia-shield/custom_components/nvidia_shield/image.py

"""Screen-capture image entity — an OPT-IN diagnostic surface (t66).

The user asked for "an image entity of the device screen capture? Can that be
added". This is that entity, and it is the only thing in the integration that
grabs the screen. Three boundaries are deliberate and load-bearing:

* **It never drives state.** An image entity's state IS its
  ``image_last_updated`` timestamp, and nothing here writes to the manager's
  store, to the power ladder or to any other entity. A capture cannot make the
  media player, the sensors or the power switch claim something the device did
  not report. (The user's original constraint: screen capture must never drive
  state.)
* **It is opt-in, default OFF** (``CONF_SCREEN_CAPTURE_ENABLED``). With the
  option absent ``async_setup_entry`` adds no entity at all.
* **It never polls.** ``ImageEntity`` ships ``_attr_should_poll = False`` ("no
  need to poll image entities") and this module keeps it. Nothing runs on a
  timer; a capture happens only when something asks for the bytes (the image
  proxy serving ``entity_picture``) or when the user calls
  ``homeassistant.update_entity``.

BINARY-SAFE CAPTURE — base64 over the shell, and why not a device file
=====================================================================
``AdbDevice.shell()`` returns ``str``: it decodes the shell channel as text, so
raw PNG bytes cannot travel through it — ``\\x89PNG`` is not UTF-8 and the
undecodable bytes would be mangled on the way out. There are two binary-safe
routes and this module takes the SECOND:

1. ``screencap -p /sdcard/<name>.png`` then ``AdbDevice.pull()`` — the library's
   real binary path (the ADB sync protocol). It needs the raw ``AdbDevice``
   object, which the ADB client owns as a private attribute, and it writes to
   the device's shared storage: something must delete that file afterwards, and
   a crash between write and delete leaves a full-resolution screenshot behind
   on the user's device.
2. ``screencap -p | base64`` then ``base64.b64decode`` locally. Base64 output is
   pure ASCII, so the shell's text decoding cannot corrupt it, and there is no
   device-side file at all — nothing to clean up, nothing to leak.

Route 2 costs ~33% more bytes on the socket than route 1, and one claim is
INFERRED rather than measured: no Shield was available here, so how the
persistent ADB socket handles a few hundred kilobytes of base64 in one
``shell()`` call is untested on hardware (the README says so). It is still the
right trade: it uses the PUBLIC seam the watchdog's socket already offers
(``manager.adb_poller.shell``), leaves no artifact on the device, and its
correctness on the wire is provable here with a real PNG fixture.

One capture costs exactly ONE shell, whatever the caller does: concurrent
requests share a lock, and a successful capture is reused for
``CAPTURE_REUSE_SECONDS`` so a dashboard repainting cannot hammer the device.
"""

from __future__ import annotations

import asyncio
import base64
import binascii
import logging
import time

from homeassistant.components.image import ImageEntity
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.util import dt as dt_util

from .const import CONF_SCREEN_CAPTURE_ENABLED, DOMAIN

_LOGGER = logging.getLogger(__name__)

#: ``screencap -p`` writes PNG to stdout; ``base64`` makes it ASCII, which is the
#: whole trick — see the module docstring.
SCREEN_CAPTURE_COMMAND = "screencap -p | base64"

#: The PNG file signature. Nothing but real PNG bytes ever leaves this module.
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"

#: A successful capture is reused for this long. Several image fetches in a row
#: (a dashboard reload, two clients, a stream preview) must not become one screen
#: grab each — "cannot hammer the device" is a property of the code, not a hope.
CAPTURE_REUSE_SECONDS = 5.0

#: The reason the entity is (un)available, published as ``capture_status``.
STATUS_OK = "ok"
STATUS_ADB_DISABLED = "adb_disabled"
STATUS_ADB_DISCONNECTED = "adb_disconnected"
STATUS_CAPTURE_FAILED = "capture_failed"

#: The attribute HA reads beside the state — declared unrecorded (G12).
ATTR_CAPTURE_STATUS = "capture_status"


def decode_screen_capture(output: str | None) -> bytes | None:
    """Decode one ``screencap -p | base64`` answer into PNG bytes, or None.

    Pure on purpose, so the one part that can corrupt an image is testable
    without a device: the answer is base64 with ``base64``'s line wrapping (the
    whitespace is not part of the payload), the decoded bytes MUST carry the PNG
    signature, and every other outcome — empty output, a shell that never
    answered, a device that printed an error line instead — is a failure that
    returns None rather than a corrupt or zero-byte image.
    """
    if not output:
        return None
    try:
        data = base64.b64decode("".join(output.split()), validate=True)
    except (binascii.Error, ValueError):
        return None
    if not data.startswith(PNG_SIGNATURE):
        return None
    # The signature alone is not a PNG: the spec puts IHDR as the FIRST chunk,
    # immediately after it (length 13, type "IHDR"). A shell that returned the
    # signature plus junk — or a device that answered with a header in another
    # format — fails here. What CANNOT be detected locally is a truncation AFTER
    # the header: those bytes are still a well-formed PNG prefix, and the
    # consumer (HA's image proxy, then the browser) is what rejects them.
    if data[12:16] != b"IHDR":
        return None
    return data


async def async_setup_entry(
    hass: HomeAssistant,
    entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Add the screen-capture entity — only when the user opted in (t66)."""
    if not entry.options.get(CONF_SCREEN_CAPTURE_ENABLED, False):
        return  # default OFF: no entity, no capture, no cost

    manager = getattr(entry, "runtime_data", None)
    if manager is None:
        _LOGGER.debug("Screen capture skipped: no runtime manager on the entry")
        return

    async_add_entities([NvidiaShieldScreenImage(hass, entry)])


class NvidiaShieldScreenImage(ImageEntity):
    """The Shield's current screen as a PNG — manual, diagnostic, never a source.

    A capture is fetched on demand and nothing else: the entity publishes no
    command, writes no state source and is not polled. Its state is the
    timestamp of the last successful capture, and its ``capture_status``
    attribute names why it cannot serve an image when it cannot.
    """

    _attr_has_entity_name = True
    _attr_translation_key = "screen_capture"
    _attr_content_type = "image/png"
    # HA's own default for image entities, kept explicitly: this entity must
    # never turn into a timer that grabs the screen every scan interval.
    _attr_should_poll = False
    # G12: a poll-adjacent value must not bloat the recorder database, and the
    # image bytes are never part of the state at all.
    _unrecorded_attributes = frozenset({ATTR_CAPTURE_STATUS})

    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
        """Initialise the entity over the entry's runtime manager.

        The two-argument shape every platform in this integration uses: the
        manager travels on ``entry.runtime_data`` (t40's wiring), so the real-HA
        entity matrix can build this class exactly like the others.
        """
        super().__init__(hass)
        # ImageEntity.__init__ builds HA's shared httpx client but does NOT set
        # ``hass`` (the platform does that in production). The rest of this
        # integration's entities assign it themselves, and without it HA's own
        # state write raises "Attribute hass is None" — measured in the real-HA
        # matrix, which builds entities the same way the tests do.
        self.hass = hass
        self._entry = entry
        self._png: bytes | None = None
        self._captured_at: float = 0.0
        # One shell at a time, however many callers ask for the image.
        self._lock = asyncio.Lock()
        poller = self._adb_poller()
        self._status = (
            STATUS_ADB_DISABLED
            if poller is None
            else STATUS_OK
            if poller.connected
            else STATUS_ADB_DISCONNECTED
        )
        self._attr_unique_id = f"{DOMAIN}_{entry.entry_id}_screen_capture"
        self._attr_device_info = DeviceInfo(
            identifiers={(DOMAIN, entry.entry_id)},
            name=entry.title,
            manufacturer="NVIDIA",
            model="Shield TV",
        )

    # -- availability ------------------------------------------------------

    def _adb_poller(self):
        """The ADB client this entry owns, or None when ADB is off/not built."""
        manager = getattr(self._entry, "runtime_data", None)
        if manager is None:
            return None
        return getattr(manager, "adb_poller", None)

    @property
    def available(self) -> bool:
        """True only while a capture can actually be served.

        Unavailable is the honest answer for "ADB is off", "ADB is down" and "the
        last capture failed"; the alternative — being available and handing HA an
        empty image — is exactly the blank screenshot this must never produce.
        """
        return self._status == STATUS_OK
    @property
    def capture_status(self) -> str:
        """Why the entity is (un)available — the same value as the attribute."""
        return self._status

    @property
    def extra_state_attributes(self) -> dict:
        """The reason behind the availability, as a PROPERTY (HA reads attrs)."""
        return {ATTR_CAPTURE_STATUS: self._status}

    # -- the capture path --------------------------------------------------

    async def async_image(self) -> bytes | None:
        """Return the PNG bytes, capturing on demand (the image proxy seam)."""
        return await self._async_capture()

    async def async_update(self) -> None:
        """Explicit refresh (``homeassistant.update_entity``) — skip the reuse.

        This is also how a failed capture is retried: a failure makes the entity
        unavailable on purpose, and an explicit refresh is the user's way back
        without waiting for the next dashboard fetch.
        """
        await self._async_capture(force=True)

    async def _async_capture(self, *, force: bool = False) -> bytes | None:
        """One screen grab, shared by everyone who asks within the reuse window."""
        async with self._lock:
            if (
                not force
                and self._png is not None
                and time.monotonic() - self._captured_at < CAPTURE_REUSE_SECONDS
            ):
                return self._png

            poller = self._adb_poller()
            if poller is None:
                return self._fail(STATUS_ADB_DISABLED)
            if not poller.connected:
                return self._fail(STATUS_ADB_DISCONNECTED)
            try:
                output = await poller.shell(SCREEN_CAPTURE_COMMAND)
            except Exception as err:  # the transport is remote; it may raise anything
                # ``AdbClientBase.shell`` already turns its own failures into
                # None, so this is defence in depth — a raised error must still
                # end as "unavailable", never as a blank image in HA.
                _LOGGER.warning("Screen capture failed on %s: %s", self._entry.title, err)
                return self._fail(STATUS_CAPTURE_FAILED)
            data = decode_screen_capture(output)
            if data is None:
                # Never a blank or zero-byte image: no answer, a truncated one,
                # an error line or a non-PNG all land here.
                _LOGGER.warning(
                    "Screen capture from %s did not return a PNG (%d bytes of "
                    "shell output)",
                    self._entry.title,
                    len(output or ""),
                )
                return self._fail(STATUS_CAPTURE_FAILED)

            self._png = data
            self._captured_at = time.monotonic()
            self._status = STATUS_OK
            self._attr_image_last_updated = dt_util.utcnow()
            self.async_write_ha_state()
            return data

    def _fail(self, status: str) -> None:
        """Record why no image can be served and stay unavailable."""
        self._status = status
        return None