Encrypt saved images and add explorer

app.py

@@ -3,6 +3,7 @@ import random
 import secrets
 import shutil
 import sys
+import tempfile
 import threading
 import time
 import urllib.parse
@@ -32,17 +33,21 @@ BaseEventLoop.__del__ = close_loop
 
 import gradio as gr
 import numpy as np
+import py7zr
 import spaces
 import torch
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
 from fastapi import Body, Depends, HTTPException, Query
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import Response
+from fastapi.responses import FileResponse, Response
 from gradio_client import Client
 from gradio.routes import App
 from huggingface_hub import download_bucket_files, list_bucket_tree
 from PIL import Image
 from PIL.PngImagePlugin import PngInfo
 from pydantic import BaseModel, Field
+from starlette.background import BackgroundTask
 from workflow_api import image_metadata
 
 
@@ -53,9 +58,16 @@ DATA_DIR = Path(os.environ.get("DATA_DIR", DEFAULT_DATA_DIR))
 LOCAL_MODEL_DIR = Path(os.environ.get("LOCAL_MODEL_DIR", "/tmp/models"))
 COMFYUI_PATH = Path(os.environ.get("COMFYUI_PATH", Path.cwd() / "ComfyUI"))
 OUTPUT_DIR = DATA_DIR / "output"
+IMAGE_DIR = DATA_DIR / "img"
 TIMEZONE = ZoneInfo("Asia/Singapore")
 MIB = 1024 * 1024
 COPY_SIZE = 16 * MIB
+SALT_SIZE = 16
+NONCE_SIZE = 12
+SCRYPT_N = 2**14
+FILE_MAGIC = b"EPNG1"
+IMAGE_SUFFIXES = (".epng", ".png")
+PREVIEW_QUALITY = 75
 RETURN_SCALE = .5
 PORT = int(os.environ.get("PORT", "7860"))
 LOCAL_URL = os.environ.get("LOCAL_URL", f"http://127.0.0.1:{PORT}")
@@ -65,6 +77,9 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 PASSWORD = os.environ.get("pass")
 if not PASSWORD:
     raise RuntimeError("pass environment variable is required")
+HAIP = os.environ.get("haip")
+if not HAIP:
+    raise RuntimeError("haip environment variable is required")
 
 BUCKET_ID = "q6/C-storage"
 MODEL_KINDS = ("checkpoints", "vae", "loras")
@@ -151,12 +166,11 @@ def cleanup_mount():
         shutil.rmtree(OUTPUT_DIR)
         log("Deleted output directory")
 
-    image_dir = DATA_DIR / "img"
-    if image_dir.is_dir():
-        for path in image_dir.rglob("*"):
+    if IMAGE_DIR.is_dir():
+        for path in IMAGE_DIR.rglob("*"):
             if path.is_file() and path.suffix.casefold() == ".7z":
                 path.unlink()
-                log(f"Deleted img/{path.relative_to(image_dir)}")
+                log(f"Deleted img/{path.relative_to(IMAGE_DIR)}")
 
 
 def download(url, target):
@@ -243,6 +257,10 @@ class DirectRequest(BaseModel):
     denoise: float = Field(DEFAULT_DENOISE, ge=0, le=1)
 
 
+class DownloadRequest(BaseModel):
+    items: list[str]
+
+
 def ensure_comfy():
     if COMFYUI_PATH.is_dir():
         log(f"Found ComfyUI at {COMFYUI_PATH}")
@@ -584,20 +602,37 @@ def image_png_bytes(image):
     return png_bytes(image, pnginfo)
 
 
+def image_key(salt):
+    return Scrypt(
+        salt=salt,
+        length=32,
+        n=SCRYPT_N,
+        r=8,
+        p=1,
+    ).derive(HAIP.encode())
+
+
 def save_image(image, tag):
-    path = DATA_DIR / "img" / datetime.now(TIMEZONE).date().isoformat()
+    path = IMAGE_DIR / datetime.now(TIMEZONE).date().isoformat()
     path.mkdir(parents=True, exist_ok=True)
     number = (
         sum(
-            file.suffix.casefold() in (".epng", ".png")
+            file.suffix.casefold() in IMAGE_SUFFIXES
             for file in path.iterdir()
         )
         + 1
     )
-    name = f"{number}{tag}.png"
+    name = f"{number}{tag}.epng"
     output = path / name
-    temp = output.with_suffix(".png.part")
-    temp.write_bytes(image_png_bytes(image))
+    temp = output.with_suffix(".epng.part")
+    salt = os.urandom(SALT_SIZE)
+    nonce = os.urandom(NONCE_SIZE)
+    encrypted = AESGCM(image_key(salt)).encrypt(
+        nonce,
+        image_png_bytes(image),
+        FILE_MAGIC,
+    )
+    temp.write_bytes(FILE_MAGIC + salt + nonce + encrypted)
     temp.replace(output)
     log(f"Saved img/{path.name}/{name}")
 
@@ -624,6 +659,53 @@ def png_bytes(image, pnginfo=None):
     return data.getvalue()
 
 
+def stored_path(value):
+    root = IMAGE_DIR.resolve()
+    path = (root / value).resolve()
+    try:
+        path.relative_to(root)
+    except ValueError as error:
+        raise HTTPException(404, "image not found") from error
+    if not path.is_file() or path.suffix.casefold() not in IMAGE_SUFFIXES:
+        raise HTTPException(404, "image not found")
+    return path
+
+
+def stored_bytes(path):
+    data = path.read_bytes()
+    if path.suffix.casefold() == ".png":
+        return data
+    if not data.startswith(FILE_MAGIC):
+        raise HTTPException(500, "invalid encrypted image")
+    salt_start = len(FILE_MAGIC)
+    nonce_start = salt_start + SALT_SIZE
+    data_start = nonce_start + NONCE_SIZE
+    return AESGCM(image_key(data[salt_start:nonce_start])).decrypt(
+        data[nonce_start:data_start],
+        data[data_start:],
+        FILE_MAGIC,
+    )
+
+
+def selected_paths(items):
+    root = IMAGE_DIR.resolve()
+    selected = {}
+    for value in items:
+        path = (root / value).resolve()
+        try:
+            path.relative_to(root)
+        except ValueError as error:
+            raise HTTPException(404, "image not found") from error
+        paths = path.iterdir() if path.is_dir() else (stored_path(value),)
+        for image in paths:
+            if image.is_file() and image.suffix.casefold() in IMAGE_SUFFIXES:
+                name = image.relative_to(root).with_suffix(".png").as_posix()
+                selected[name] = image
+    if not selected:
+        raise HTTPException(422, "no images selected")
+    return selected
+
+
 def valid_pass(p):
     return secrets.compare_digest(str(p), PASSWORD)
 
@@ -1023,6 +1105,95 @@ def archive(body: bytes = Body()):
     return {"status": True}
 
 
+@api.get(
+    "/explorer",
+    dependencies=[Depends(require_pass)],
+)
+def explorer():
+    return FileResponse(Path.cwd() / "explorer.html")
+
+
+@api.get("/explorer.css")
+def explorer_css():
+    return FileResponse(Path.cwd() / "explorer.css", media_type="text/css")
+
+
+@api.get("/explorer.js")
+def explorer_js():
+    return FileResponse(
+        Path.cwd() / "explorer.js",
+        media_type="text/javascript",
+    )
+
+
+@api.get(
+    "/api/images",
+    dependencies=[Depends(require_pass)],
+)
+def image_list():
+    root = IMAGE_DIR
+    dates = []
+    if root.is_dir():
+        for folder in sorted(root.iterdir(), reverse=True):
+            if not folder.is_dir():
+                continue
+            items = [
+                {
+                    "name": path.with_suffix(".png").name,
+                    "path": path.relative_to(root).as_posix(),
+                    "size": path.stat().st_size,
+                }
+                for path in sorted(folder.iterdir(), reverse=True)
+                if path.is_file()
+                and path.suffix.casefold() in IMAGE_SUFFIXES
+            ]
+            if items:
+                dates.append({"name": folder.name, "images": items})
+    return {"dates": dates}
+
+
+@api.get(
+    "/api/images/preview",
+    dependencies=[Depends(require_pass)],
+)
+def image_preview(path: str = Query(...)):
+    with Image.open(BytesIO(stored_bytes(stored_path(path)))) as image:
+        data = BytesIO()
+        image.save(data, format="WEBP", quality=PREVIEW_QUALITY)
+    return Response(
+        content=data.getvalue(),
+        media_type="image/webp",
+        headers={"Cache-Control": "private, max-age=86400"},
+    )
+
+
+@api.post(
+    "/api/images/download",
+    dependencies=[Depends(require_pass)],
+)
+def image_download(body: DownloadRequest):
+    temp = tempfile.NamedTemporaryFile(suffix=".7z", delete=False)
+    temp.close()
+    try:
+        with py7zr.SevenZipFile(
+            temp.name,
+            "w",
+            password=PASSWORD,
+            header_encryption=True,
+        ) as archive_file:
+            for name, path in selected_paths(body.items).items():
+                archive_file.writestr(stored_bytes(path), name)
+    except Exception:
+        Path(temp.name).unlink(missing_ok=True)
+        raise
+    return FileResponse(
+        temp.name,
+        filename="images.7z",
+        media_type="application/x-7z-compressed",
+        background=BackgroundTask(Path(temp.name).unlink, missing_ok=True),
+    )
+
+
 def login(p):
     if not valid_pass(p):
         raise gr.Error("Invalid password")

explorer.css

@@ -0,0 +1,140 @@
+@import url('https://fonts.googleapis.com/css2?family=Spline+Sans:wght@400;500;600&display=swap');
+
+:root {
+    --bg: #222;
+    --panel: #333;
+    --text: #ddd;
+    --muted: #999;
+    --accent: #7dc;
+    --border: #555;
+    --gap: 1rem;
+    --width: 90%;
+}
+
+* {
+    box-sizing: border-box;
+}
+
+body {
+    background: var(--bg);
+    color: var(--text);
+    font-family: "Spline Sans", sans-serif;
+    margin: 0;
+}
+
+image-explorer {
+    display: block;
+    margin: 3rem auto;
+    width: var(--width);
+}
+
+h1 {
+    font-size: 1.5rem;
+    font-weight: 500;
+    margin: 0 0 var(--gap);
+}
+
+.tools {
+    display: flex;
+    flex-wrap: wrap;
+    gap: .5rem;
+    margin-bottom: var(--gap);
+}
+
+button,
+input[type="search"] {
+    background: var(--panel);
+    border: 1px solid var(--border);
+    border-radius: 0;
+    color: var(--text);
+    font: inherit;
+    padding: .65rem .8rem;
+}
+
+input[type="search"] {
+    flex: 1 1 16rem;
+}
+
+button {
+    cursor: pointer;
+}
+
+button.primary {
+    background: var(--accent);
+    border-color: var(--accent);
+    color: var(--bg);
+}
+
+button:disabled {
+    cursor: wait;
+    opacity: .6;
+}
+
+.status,
+.count,
+.size {
+    color: var(--muted);
+}
+
+.status {
+    min-height: 1.5rem;
+}
+
+.date {
+    border-top: 1px solid var(--border);
+    margin-top: 1.5rem;
+    padding-top: 1rem;
+}
+
+.date-title {
+    align-items: center;
+    display: flex;
+    gap: .6rem;
+    margin-bottom: .8rem;
+}
+
+.grid {
+    display: grid;
+    gap: var(--gap);
+    grid-template-columns: repeat(auto-fill, minmax(12rem, 1fr));
+}
+
+.image {
+    background: var(--panel);
+    border: 1px solid var(--border);
+    cursor: pointer;
+    display: grid;
+    gap: .5rem;
+    padding: .5rem;
+}
+
+.image img {
+    aspect-ratio: 1;
+    background: var(--bg);
+    display: block;
+    object-fit: contain;
+    width: 100%;
+}
+
+.image-info {
+    display: grid;
+    gap: .2rem;
+    grid-template-columns: auto 1fr;
+}
+
+.name,
+.size {
+    min-width: 0;
+    overflow: hidden;
+    text-overflow: ellipsis;
+    white-space: nowrap;
+}
+
+.size {
+    grid-column: 2;
+    font-size: .8rem;
+}
+
+[hidden] {
+    display: none !important;
+}

explorer.html

@@ -0,0 +1,14 @@
+<!doctype html>
+<html lang="en">
+<head>
+    <meta charset="utf-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1">
+    <title>Image Explorer</title>
+    <link rel="icon" href="data:,">
+    <link rel="stylesheet" href="/explorer.css">
+    <script type="module" src="/explorer.js"></script>
+</head>
+<body>
+    <image-explorer></image-explorer>
+</body>
+</html>

explorer.js

@@ -0,0 +1,149 @@
+const $ = (selector, root = document) => root.querySelector(selector)
+const $$ = (selector, root = document) => [...root.querySelectorAll(selector)]
+const p = new URLSearchParams(location.search).get("p")
+const api = (path, query = {}) => `${path}?${new URLSearchParams({ p, ...query })}`
+const preview = path => api("/api/images/preview", { path })
+const size = bytes => `${(bytes / 1024).toFixed(bytes < 10240 ? 1 : 0)} KiB`
+
+class ImageExplorer extends HTMLElement {
+    connectedCallback() {
+        this.innerHTML = `
+            <main>
+                <h1>Images</h1>
+                <div class="tools">
+                    <input type="search" placeholder="Filter images" aria-label="Filter images">
+                    <button class="all" type="button">Select all</button>
+                    <button class="clear" type="button">Clear</button>
+                    <button class="download primary" type="button">Download 7z</button>
+                </div>
+                <div class="status">Loading</div>
+                <div class="dates"></div>
+            </main>`
+        this.load()
+        $(".all", this).addEventListener("click", () => this.select(true))
+        $(".clear", this).addEventListener("click", () => this.select(false))
+        $(".download", this).addEventListener("click", () => this.download())
+        $("input[type=search]", this).addEventListener("input", event => {
+            this.filter(event.target.value)
+        })
+    }
+
+    async load() {
+        const response = await fetch(api("/api/images"))
+        if (!response.ok) {
+            $(".status", this).textContent = "Unable to load images"
+            return
+        }
+        const data = await response.json()
+        for (const date of data.dates) this.addDate(date)
+        this.update()
+    }
+
+    addDate(date) {
+        const section = document.createElement("section")
+        section.className = "date"
+        const title = document.createElement("label")
+        title.className = "date-title"
+        const folder = document.createElement("input")
+        folder.type = "checkbox"
+        folder.className = "folder"
+        const name = document.createElement("strong")
+        name.textContent = date.name
+        const count = document.createElement("span")
+        count.className = "count"
+        count.textContent = `${date.images.length} images`
+        const grid = document.createElement("div")
+        grid.className = "grid"
+        title.append(folder, name, count)
+        section.append(title, grid)
+        for (const image of date.images) grid.append(this.card(image))
+        folder.addEventListener("change", () => {
+            for (const input of $$(".file", section)) input.checked = folder.checked
+            this.update()
+        })
+        $(".dates", this).append(section)
+    }
+
+    card(image) {
+        const card = document.createElement("label")
+        card.className = "image"
+        card.dataset.name = image.name.toLowerCase()
+        const picture = document.createElement("img")
+        picture.loading = "lazy"
+        picture.src = preview(image.path)
+        picture.alt = image.name
+        const info = document.createElement("span")
+        info.className = "image-info"
+        const input = document.createElement("input")
+        input.type = "checkbox"
+        input.className = "file"
+        input.dataset.path = image.path
+        const name = document.createElement("span")
+        name.className = "name"
+        name.textContent = image.name
+        const bytes = document.createElement("span")
+        bytes.className = "size"
+        bytes.textContent = size(image.size)
+        input.addEventListener("change", () => this.update())
+        info.append(input, name, bytes)
+        card.append(picture, info)
+        return card
+    }
+
+    select(checked) {
+        for (const input of $$(".file", this)) input.checked = checked
+        this.update()
+    }
+
+    update() {
+        for (const section of $$(".date", this)) {
+            const files = $$(".file", section)
+            const checked = files.filter(input => input.checked).length
+            const folder = $(".folder", section)
+            folder.checked = checked === files.length
+            folder.indeterminate = checked > 0 && checked < files.length
+        }
+        const count = $$(".file:checked", this).length
+        $(".status", this).textContent = `${count} selected`
+    }
+
+    filter(value) {
+        const query = value.trim().toLowerCase()
+        for (const section of $$(".date", this)) {
+            for (const card of $$(".image", section)) {
+                card.hidden = Boolean(query && !card.dataset.name.includes(query))
+            }
+            section.hidden = $$(".image:not([hidden])", section).length === 0
+        }
+    }
+
+    async download() {
+        const items = $$(".file:checked", this).map(input => input.dataset.path)
+        if (!items.length) return
+        const button = $(".download", this)
+        button.disabled = true
+        $(".status", this).textContent = "Building archive"
+        try {
+            const response = await fetch(api("/api/images/download"), {
+                method: "POST",
+                headers: { "Content-Type": "application/json" },
+                body: JSON.stringify({ items }),
+            })
+            if (!response.ok) throw new Error("Archive failed")
+            const url = URL.createObjectURL(await response.blob())
+            const link = document.createElement("a")
+            link.href = url
+            link.download = "images.7z"
+            link.click()
+            URL.revokeObjectURL(url)
+        } catch (error) {
+            $(".status", this).textContent = error.message
+            return
+        } finally {
+            button.disabled = false
+        }
+        this.update()
+    }
+}
+
+customElements.define("image-explorer", ImageExplorer)

requirements.txt

@@ -1,7 +1,9 @@
 fastapi
+cryptography
 gradio==6.20.0
 huggingface_hub>=1.6
 numpy
 Pillow
+py7zr==1.1.2
 uvicorn
 -r https://raw.githubusercontent.com/Comfy-Org/ComfyUI/master/requirements.txt