ComfyAssets-ComfyUI_PromptManager/py/api/history_current.py

"""Native reported locations and already-opened current media, without ownership."""

import asyncio
from dataclasses import dataclass
import errno
import hashlib
import hmac
import io
import json
import mimetypes
import os
from pathlib import Path
import stat
from typing import BinaryIO

from aiohttp import web

from . import preview_store, thumbnail_settings
from .images import _canonical, _public_path, SERVABLE_EXTENSIONS, VIDEO_EXTENSIONS
from .images import in_trash_below_roots
from .images import existing_media_thumbnail, media_preview_path

try:
    from ...utils.execution_outputs import (
        OutputReference,
        OutputRoots,
        snapshot_output_roots,
        validate_output_reference,
    )
except ImportError:
    from utils.execution_outputs import (
        OutputReference,
        OutputRoots,
        snapshot_output_roots,
        validate_output_reference,
    )

_SECRET = os.urandom(32)
_READ_SIZE = 64 * 1024


def _opaque(value: object) -> str:
    raw = json.dumps(value, ensure_ascii=True, separators=(",", ":")).encode("ascii")
    return hmac.new(_SECRET, raw, hashlib.sha256).hexdigest()


@dataclass(frozen=True)
class CurrentOutputScope:
    path_style: str
    lexical_output: str | None
    lexical_temp: str | None
    roots: OutputRoots
    canonical_output: Path | None
    canonical_temp: Path | None
    authorized_roots: tuple[Path, ...]
    current_scope: str


@dataclass(frozen=True)
class ResolvedCurrentCandidate:
    public: dict
    canonical_target: Path | None = None
    namespace: Path | None = None
    gallery_root: Path | None = None


class CurrentLocationUnavailable(Exception):
    """A bounded current-file failure without the underlying private path."""

    def __init__(self, reason: str, *, stale: bool = False) -> None:
        self.reason = reason
        self.stale = stale
        super().__init__(reason)


def _host_directory(name: str) -> str | None:
    try:
        import folder_paths

        value = getattr(folder_paths, name)()
    except (ImportError, AttributeError, OSError):
        return None
    return value if type(value) is str else None


def snapshot_current_output_scope(api: object) -> CurrentOutputScope:
    """Bind host lexical namespaces independently of current gallery authority."""
    from ..config import GalleryConfig

    style = "windows" if os.name == "nt" else "posix"
    output = _host_directory("get_output_directory")
    temp = _host_directory("get_temp_directory")
    roots = snapshot_output_roots(style, output, temp)
    canonical_output = _canonical(output) if roots.output_key is not None else None
    canonical_temp = _canonical(temp) if roots.temp_key is not None else None
    authorized = tuple(
        sorted({_canonical(path) for path in api._get_all_output_dirs()})
    )
    policy = []
    for entry in GalleryConfig.MONITORING_DIRECTORIES:
        valid, _ = GalleryConfig.validate_gallery_root(entry)
        policy.append((entry if type(entry) is str else None, valid))
    current_scope = _opaque(
        [
            "prompt-manager-current-output-v1",
            style,
            output,
            temp,
            roots.output_key,
            roots.temp_key,
            str(canonical_output) if canonical_output is not None else None,
            str(canonical_temp) if canonical_temp is not None else None,
            [str(path) for path in authorized],
            sorted(
                str(_canonical(path))
                for path in GalleryConfig.allowed_gallery_parents()
            ),
            sorted(policy, key=lambda item: (item[0] or "", item[1])),
            sorted(str(_canonical(path)) for path in GalleryConfig.path_anchors()),
            "gallery-validation-v1",
        ]
    )
    return CurrentOutputScope(
        style,
        output,
        temp,
        roots,
        canonical_output,
        canonical_temp,
        authorized,
        current_scope,
    )


def _descriptor_reason(reference: dict, scope: CurrentOutputScope) -> str | None:
    try:
        validate_output_reference(OutputReference("current", 0, **reference))
    except (TypeError, ValueError):
        return "namespace_mismatch"
    if reference["path_style"] != scope.path_style:
        return "foreign_path_style"
    folder = reference["folder_type"]
    expected = scope.roots.output_key if folder == "output" else scope.roots.temp_key
    if reference["root_key"] is None or expected is None:
        return "root_unavailable"
    if reference["root_key"] != expected:
        return "namespace_mismatch"
    return None


def candidate_locator(reference: dict, scope: CurrentOutputScope) -> str | None:
    """Validate the native namespace before constructing a lexical path; no IO."""
    if _descriptor_reason(reference, scope) is not None:
        return None
    root = (
        scope.lexical_output
        if reference["folder_type"] == "output"
        else scope.lexical_temp
    )
    return os.path.normcase(
        os.path.normpath(
            os.path.join(root, reference["subfolder"], reference["filename"])
        )
    )


def _authorized(target: Path, namespace: Path, scope: CurrentOutputScope) -> bool:
    return target.is_relative_to(namespace) and any(
        target.is_relative_to(root) for root in scope.authorized_roots
    )


def _unavailable(reason: str, *, error: bool = False) -> ResolvedCurrentCandidate:
    return ResolvedCurrentCandidate(
        {
            "status": "error" if error else "unavailable",
            "reason": reason,
            "evidence_class": "locator_candidate",
        }
    )


def resolve_current_candidate(
    reference: dict, scope: CurrentOutputScope, gallery: dict
) -> ResolvedCurrentCandidate:
    """Observe only this exact current locator and its snapshot-owned association."""
    reason = _descriptor_reason(reference, scope)
    if reason is not None:
        return _unavailable(reason)
    lexical = candidate_locator(reference, scope)
    namespace = (
        scope.canonical_output
        if reference["folder_type"] == "output"
        else scope.canonical_temp
    )
    try:
        if Path(reference["filename"]).suffix.lower() not in SERVABLE_EXTENSIONS:
            return _unavailable("unsupported_media")
        if namespace is None or not namespace.is_dir():
            return _unavailable("root_unavailable")
        target = _canonical(lexical)
        if in_trash_below_roots(target, [namespace]):
            # Every kind and the outputs listing come through here (#543).
            return _unavailable("missing")
        if not _authorized(target, namespace, scope):
            return _unavailable("not_authorized")
        if target.suffix.lower() not in SERVABLE_EXTENSIONS:
            return _unavailable("unsupported_media")
        observed = target.stat()
        if not stat.S_ISREG(observed.st_mode):
            return _unavailable("not_regular_file")
        if not os.access(target, os.R_OK) or not observed.st_mode & 0o444:
            return _unavailable("unreadable")
        matching = gallery.get(lexical, {})
        root = max(
            (root for root in scope.authorized_roots if target.is_relative_to(root)),
            key=lambda root: len(root.parts),
        )
        return ResolvedCurrentCandidate(
            {
                "status": "available",
                "reason": None,
                "evidence_class": "locator_candidate",
                "gallery_match": matching.get("gallery_match", "none"),
                "gallery_image_id": matching.get("gallery_image_id"),
                "source_label": _public_path(namespace),
                "relative_path": Path(
                    reference["subfolder"], reference["filename"]
                ).as_posix(),
                "current_scope": scope.current_scope,
                "location_key": _opaque(
                    ["current-location-v1", scope.current_scope, str(target)]
                ),
            },
            target,
            namespace,
            root,
        )
    except FileNotFoundError:
        return _unavailable("missing")
    except PermissionError:
        return _unavailable("unreadable")
    except (OSError, ValueError):
        return _unavailable("temporary_unavailable", error=True)


def _target_authorized(target, namespace, scope, store_root=None) -> bool:
    """A preview in the thumbnail folder (#320, folder mode) is authorized by
    the store of its source (``preview_store.owns``: inside the key folder,
    no link); everything else by the output namespace and the roots."""
    if (
        store_root is not None
        and preview_store.current().mode == preview_store.FOLDER_MODE
    ):
        return preview_store.owns(store_root, target)
    return _authorized(target, namespace, scope)


def _open_current_target(
    target: Path,
    namespace: Path,
    scope: CurrentOutputScope,
    preview: bool = False,
    store_root: Path | None = None,
) -> BinaryIO:
    """Open only the canonical target and verify authority against the actual handle.

    A *preview* (a file under ``thumbnails/``) may also be JPEG XL (#317); in
    folder mode it is authorized by the store of *store_root* (#320)."""
    if _canonical(target) != target:
        raise CurrentLocationUnavailable("not_authorized", stale=True)
    if in_trash_below_roots(target, [namespace]):
        # Also at the open: the cover chooses its preview after the candidate
        # check, and a preview can be a link into a trash folder (#543).
        raise CurrentLocationUnavailable("missing")
    if not _target_authorized(target, namespace, scope, store_root):
        raise CurrentLocationUnavailable("not_authorized")
    suffix = target.suffix.lower()
    if suffix not in SERVABLE_EXTENSIONS and not (
        preview and suffix in thumbnail_settings.THUMBNAIL_EXTENSIONS
    ):
        raise CurrentLocationUnavailable("unsupported_media")
    flags = (
        os.O_RDONLY
        | getattr(os, "O_BINARY", 0)
        | getattr(os, "O_NOFOLLOW", 0)
        | getattr(os, "O_NONBLOCK", 0)
    )
    try:
        fd = os.open(target, flags)
    except FileNotFoundError:
        raise CurrentLocationUnavailable("missing") from None
    except PermissionError:
        raise CurrentLocationUnavailable("unreadable") from None
    except OSError as error:
        if error.errno in (errno.ELOOP, errno.ENOTDIR):
            raise CurrentLocationUnavailable("not_authorized", stale=True) from None
        raise CurrentLocationUnavailable("temporary_unavailable") from None
    try:
        observed = os.fstat(fd)
        if not stat.S_ISREG(observed.st_mode):
            raise CurrentLocationUnavailable("not_regular_file")
        if not observed.st_mode & 0o444 or not os.access(target, os.R_OK):
            raise CurrentLocationUnavailable("unreadable")
        current = os.stat(target, follow_symlinks=False)
        if (
            _canonical(target) != target
            or (observed.st_dev, observed.st_ino) != (current.st_dev, current.st_ino)
            or not _target_authorized(target, namespace, scope, store_root)
        ):
            raise CurrentLocationUnavailable("not_authorized", stale=True)
        return os.fdopen(fd, "rb")
    except FileNotFoundError:
        os.close(fd)
        raise CurrentLocationUnavailable("missing") from None
    except PermissionError:
        os.close(fd)
        raise CurrentLocationUnavailable("unreadable") from None
    except BaseException:
        os.close(fd)
        raise


def _reauthorize_current_target(
    resolved: ResolvedCurrentCandidate, scope: CurrentOutputScope
) -> None:
    if _canonical(
        resolved.canonical_target
    ) != resolved.canonical_target or not _authorized(
        resolved.canonical_target, resolved.namespace, scope
    ):
        raise CurrentLocationUnavailable("not_authorized", stale=True)


@dataclass
class _OpenedCurrentMedia:
    handle: BinaryIO
    size: int
    content_type: str

    def close(self) -> None:
        self.handle.close()


def _preview_writable(root: Path, preview: Path) -> bool:
    """A preview may be written: the store of *root* and every link segment
    of the preview's folder resolve inside it (the serve route's rule, #249;
    ``preview_store.owns``)."""
    return preview_store.owns(root, preview)


def _is_jxl(path) -> bool:
    return Path(path).suffix.lower() == ".jxl"


def _preview_without_jxl(root: Path, relative: Path) -> Path | None:
    """The first existing preview *relative* owns that is not JPEG XL."""
    return next(
        (
            path
            for path in thumbnail_settings.owned_thumbnails(root, relative)
            if not _is_jxl(path)
        ),
        None,
    )


def accepts_jxl(request) -> bool:
    """Whether the request's ``Accept`` lists ``image/jxl``."""
    accept = request.headers.get("Accept", "")
    return any(
        part.split(";", 1)[0].strip().lower() == "image/jxl"
        for part in accept.split(",")
    )


def _placeholder_media(placeholder: bytes) -> "_OpenedCurrentMedia":
    return _OpenedCurrentMedia(io.BytesIO(placeholder), len(placeholder), "image/jpeg")


def _cover_target(resolved, ensure_preview, placeholder, accepts_jxl):
    """What the cover opens: a preview path, the placeholder media, or None
    for the original (see ``_prepare_current_media``)."""
    target = resolved.canonical_target
    relative = target.relative_to(resolved.gallery_root)
    is_video = target.suffix.lower() in VIDEO_EXTENSIONS
    # #317: any preview the file owns (the current format first); a missing
    # one is made under the current format's name.
    preview = existing_media_thumbnail(resolved.gallery_root, relative)
    if preview is not None and not accepts_jxl and _is_jxl(preview):
        preview = _preview_without_jxl(resolved.gallery_root, relative)
        if preview is None:
            if is_video and placeholder is not None:
                return _placeholder_media(placeholder)
            return None
        return _canonical(preview)
    if preview is not None:
        return _canonical(preview)
    preview = media_preview_path(resolved.gallery_root, relative, create=True)
    if (
        ensure_preview is not None
        and preview is not None
        and _preview_writable(resolved.gallery_root, preview)
    ):
        ensure_preview(target, preview, is_video)
    if is_video and placeholder is not None:
        return _placeholder_media(placeholder)
    return None


def _prepare_current_media(
    resolved: ResolvedCurrentCandidate,
    scope: CurrentOutputScope,
    *,
    thumbnail: bool,
    ensure_preview=None,
    placeholder: bytes | None = None,
    accepts_jxl: bool = True,
) -> _OpenedCurrentMedia:
    """Open the current file, or its preview for ``thumbnail``.

    A missing preview that may be written is handed to
    ``ensure_preview(src, dst, is_video)`` (the route makes it and reads
    again, #249). Without one an image's original is opened and a video is
    the ``placeholder`` JPEG, never the video inside an ``<img>``.

    The cover's URL names no format, so a JPEG XL preview goes only to a
    browser whose ``Accept`` lists ``image/jxl`` (*accepts_jxl*, review M3).
    Another one gets the file's preview in another format when there is one,
    else the original image or a video's placeholder: never a broken image.
    """
    target = resolved.canonical_target
    is_preview = False
    if thumbnail:
        chosen = _cover_target(resolved, ensure_preview, placeholder, accepts_jxl)
        if isinstance(chosen, _OpenedCurrentMedia):
            return chosen
        if chosen is not None:
            target, is_preview = chosen, True
    handle = _open_current_target(
        target,
        resolved.namespace,
        scope,
        is_preview,
        resolved.gallery_root if is_preview else None,
    )
    try:
        return _OpenedCurrentMedia(
            handle,
            os.fstat(handle.fileno()).st_size,
            thumbnail_settings.content_type(target)
            or mimetypes.guess_type(target.name)[0]
            or "application/octet-stream",
        )
    except BaseException:
        handle.close()
        raise


async def _current_media_response(
    request: web.Request, opened: _OpenedCurrentMedia
) -> web.StreamResponse:
    """Stream a bounded single range from a previously authorized handle."""
    try:
        try:
            selected = request.http_range
            start, stop = selected.start, selected.stop
            partial = request.headers.get("Range") is not None
            if (
                partial
                and request.headers["Range"].startswith("bytes=-")
                and start == 0
            ):
                raise ValueError("Empty suffix range")
            if start is not None and start < 0:
                start = max(0, opened.size + start)
            start = start or 0
            stop = min(opened.size, stop) if stop is not None else opened.size
            if partial and (start >= stop or start >= opened.size):
                raise ValueError("Unsatisfiable current media range")
        except ValueError:
            return web.Response(
                status=416,
                headers={
                    "Cache-Control": "no-store",
                    "Content-Range": f"bytes */{opened.size}",
                },
            )
        headers = {
            "Cache-Control": "no-store",
            "Accept-Ranges": "bytes",
            "Content-Type": opened.content_type,
            "Content-Length": str(stop - start),
        }
        if partial:
            headers["Content-Range"] = f"bytes {start}-{stop - 1}/{opened.size}"
        response = web.StreamResponse(status=206 if partial else 200, headers=headers)
        await response.prepare(request)
        if request.method != "HEAD":
            opened.handle.seek(start)
            remaining = stop - start
            while remaining:
                chunk = await asyncio.get_running_loop().run_in_executor(
                    None, opened.handle.read, min(_READ_SIZE, remaining)
                )
                if not chunk:
                    break
                remaining -= len(chunk)
                await response.write(chunk)
        await response.write_eof()
        return response
    finally:
        opened.close()