238 lines
9.2 KiB
Python
238 lines
9.2 KiB
Python
import base64
|
|
import ipaddress
|
|
import json
|
|
import socket
|
|
from io import BytesIO
|
|
from urllib.parse import urlsplit
|
|
from uuid import uuid4
|
|
|
|
from PIL import Image, ImageFilter, ImageOps, ImageStat
|
|
|
|
from pigway_frame_flow.errors import ErrorCode as E
|
|
from pigway_frame_flow.errors import FlowError
|
|
from pigway_frame_flow.storage.database import now_iso
|
|
|
|
MAX_IMAGE_PIXELS = 25_000_000
|
|
Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
|
|
|
|
|
|
def convert(data, fit="contain", profile="color", *, size=(1200, 1600)):
|
|
try:
|
|
with Image.open(BytesIO(data)) as source:
|
|
if source.format not in {"JPEG", "PNG", "WEBP"}:
|
|
raise ValueError()
|
|
if source.width * source.height > MAX_IMAGE_PIXELS:
|
|
raise ValueError()
|
|
img = ImageOps.exif_transpose(source).convert("RGBA")
|
|
canvas = Image.new("RGBA", img.size, "white")
|
|
canvas.alpha_composite(img)
|
|
img = canvas.convert("RGB")
|
|
if fit == "contain":
|
|
result = ImageOps.pad(
|
|
img, size, method=Image.Resampling.LANCZOS, color="white"
|
|
)
|
|
else:
|
|
center = (0.5, 0.5)
|
|
if fit == "smart":
|
|
# Lightweight edge saliency: slide an aspect-correct window
|
|
# over a thumbnail; retain the region with most detail.
|
|
thumb = img.copy()
|
|
thumb.thumbnail((240, 240))
|
|
edges = thumb.convert("L").filter(ImageFilter.FIND_EDGES)
|
|
w, h = edges.size
|
|
cw, ch = (
|
|
(int(h * size[0] / size[1]), h)
|
|
if w / h > size[0] / size[1]
|
|
else (w, int(w * size[1] / size[0]))
|
|
)
|
|
choices = []
|
|
for step in range(11):
|
|
x, y = int((w - cw) * step / 10), int((h - ch) * step / 10)
|
|
score = ImageStat.Stat(edges.crop((x, y, x + cw, y + ch))).sum[
|
|
0
|
|
]
|
|
# Favor centered crops in ties, suppress border artifacts.
|
|
choices.append((score - abs(step - 5) * 0.01, x, y))
|
|
_, x, y = max(choices)
|
|
center = (
|
|
(x / (w - cw)) if w != cw else 0.5,
|
|
(y / (h - ch)) if h != ch else 0.5,
|
|
)
|
|
result = ImageOps.fit(
|
|
img, size, method=Image.Resampling.LANCZOS, centering=center
|
|
)
|
|
if profile == "grayscale":
|
|
result = ImageOps.grayscale(result).convert("RGB")
|
|
out = BytesIO()
|
|
result.save(out, "JPEG", quality=93, optimize=True)
|
|
return out.getvalue()
|
|
except (OSError, ValueError, Image.DecompressionBombError):
|
|
raise FlowError(E.INVALID_REQUEST, "Unsupported or invalid image") from None
|
|
|
|
|
|
class Assets:
|
|
def __init__(self, db, settings):
|
|
self.db, self.settings = db, settings
|
|
|
|
def put(self, data, name="image.jpg", kind="images"):
|
|
from pathlib import Path
|
|
|
|
name = Path(name).stem[:110] + ".jpg"
|
|
asset_id = uuid4().hex
|
|
(self.settings.data_dir / kind / (asset_id + ".jpg")).write_bytes(data)
|
|
self.db.execute(
|
|
"INSERT INTO assets VALUES(?,?,?,?)", (asset_id, name, kind, now_iso())
|
|
)
|
|
return self.info(asset_id)
|
|
|
|
def record_origin(self, asset_id, operation):
|
|
job = self.db.one(
|
|
"SELECT jobs.id,jobs.payload FROM job_runs JOIN jobs ON jobs.id=job_runs.job_id WHERE operation_id=?",
|
|
(operation["id"],),
|
|
)
|
|
import json
|
|
|
|
payload = json.loads(operation["payload"])
|
|
content = payload.get("content", payload)
|
|
origin = (
|
|
"task"
|
|
if job
|
|
else "template"
|
|
if content.get("type") in ("template", "auto_collage")
|
|
else "manual"
|
|
if operation["client_id"].startswith("user:")
|
|
else "external"
|
|
)
|
|
self.db.execute(
|
|
"INSERT OR IGNORE INTO asset_metadata(asset_id,origin,job_id,job_name) VALUES(?,?,?,?)",
|
|
(
|
|
asset_id,
|
|
origin,
|
|
job["id"] if job else None,
|
|
json.loads(job["payload"])["name"] if job else None,
|
|
),
|
|
)
|
|
|
|
def promote(self, asset_id):
|
|
row = self.db.one("SELECT * FROM assets WHERE id=?", (asset_id,))
|
|
if row and row["kind"] == "cache":
|
|
source = self.db.one(
|
|
"SELECT payload FROM asset_sources WHERE asset_id=?", (asset_id,)
|
|
)
|
|
kind = (
|
|
"generated"
|
|
if source and json.loads(source["payload"]).get("type") == "template"
|
|
else "images"
|
|
)
|
|
target = self.settings.data_dir / kind / (asset_id + ".jpg")
|
|
self.path(asset_id).replace(target)
|
|
self.db.execute("UPDATE assets SET kind=? WHERE id=?", (kind, asset_id))
|
|
return self.info(asset_id)
|
|
|
|
def info(self, asset_id):
|
|
row = self.db.one("SELECT * FROM assets WHERE id=?", (asset_id,))
|
|
if not row:
|
|
raise FlowError(E.NOT_FOUND, "Asset not found", 404)
|
|
meta = self.db.one("SELECT * FROM asset_metadata WHERE asset_id=?", (asset_id,))
|
|
provenance = self.db.one(
|
|
"SELECT payload FROM asset_sources WHERE asset_id=?", (asset_id,)
|
|
)
|
|
return {
|
|
"source_path": json.loads(provenance["payload"]).get("source_path")
|
|
if provenance
|
|
else None,
|
|
"starred": bool(meta and meta["starred"]),
|
|
"origin": meta["origin"] if meta else "unknown",
|
|
"job_id": meta["job_id"] if meta else None,
|
|
"job_name": meta["job_name"] if meta else None,
|
|
"asset_id": row["id"],
|
|
"source": "generated" if row["kind"] == "generated" else "uploaded",
|
|
"name": row["name"],
|
|
"url": f"/api/v1/assets/{row['id']}",
|
|
"created_at": row["created_at"],
|
|
}
|
|
|
|
def path(self, asset_id):
|
|
row = self.db.one("SELECT * FROM assets WHERE id=?", (asset_id,))
|
|
if not row:
|
|
raise FlowError(E.NOT_FOUND, "Asset not found", 404)
|
|
return self.settings.data_dir / row["kind"] / (row["id"] + ".jpg")
|
|
|
|
def decode(self, value):
|
|
try:
|
|
data = base64.b64decode(value, validate=True)
|
|
except ValueError:
|
|
raise FlowError(E.INVALID_REQUEST, "Invalid base64 image") from None
|
|
if len(data) > self.settings.max_bytes:
|
|
raise FlowError(E.INVALID_REQUEST, "Image exceeds 20 MiB")
|
|
return data
|
|
|
|
def fetch_url(
|
|
self,
|
|
url,
|
|
*,
|
|
method="GET",
|
|
headers=None,
|
|
body=None,
|
|
limit=None,
|
|
accepted_status=(200,),
|
|
timeout=10,
|
|
):
|
|
limit = limit or self.settings.max_bytes
|
|
parsed = urlsplit(url)
|
|
if (
|
|
parsed.scheme != "https"
|
|
or not parsed.hostname
|
|
or parsed.username
|
|
or parsed.password
|
|
or parsed.port not in (None, 443)
|
|
):
|
|
raise FlowError(E.INVALID_REQUEST, "Image URL requires HTTPS on port 443")
|
|
if parsed.hostname.lower() not in self.settings.image_url_hosts:
|
|
raise FlowError(E.FORBIDDEN, "Image host is not allowed", 403)
|
|
try:
|
|
addresses = socket.getaddrinfo(
|
|
parsed.hostname, 443, type=socket.SOCK_STREAM
|
|
)
|
|
if not addresses or any(
|
|
not ipaddress.ip_address(x[4][0]).is_global for x in addresses
|
|
):
|
|
raise FlowError(
|
|
E.FORBIDDEN, "Private image addresses are forbidden", 403
|
|
)
|
|
# Pin the validated resolution for this request, including TLS SNI,
|
|
# instead of performing a second DNS lookup (rebinding).
|
|
import http.client
|
|
import ssl
|
|
|
|
ip = str(addresses[0][4][0])
|
|
|
|
class PinnedHTTPS(http.client.HTTPSConnection):
|
|
def connect(self):
|
|
sock = socket.create_connection((ip, 443), timeout=timeout)
|
|
self.sock = ssl.create_default_context().wrap_socket(
|
|
sock, server_hostname=parsed.hostname
|
|
)
|
|
|
|
connection = PinnedHTTPS(parsed.hostname, timeout=timeout)
|
|
try:
|
|
connection.request(
|
|
method,
|
|
(parsed.path or "/") + ("?" + parsed.query if parsed.query else ""),
|
|
headers=headers or {},
|
|
body=body,
|
|
)
|
|
response = connection.getresponse()
|
|
if response.status not in accepted_status:
|
|
raise FlowError(
|
|
E.INVALID_REQUEST, "Image URL returned an invalid response"
|
|
)
|
|
data = response.read(limit + 1)
|
|
if len(data) > limit:
|
|
raise FlowError(E.INVALID_REQUEST, "Image exceeds 20 MiB")
|
|
return data
|
|
finally:
|
|
connection.close()
|
|
except (OSError, ValueError):
|
|
raise FlowError(E.INVALID_REQUEST, "Unable to fetch image URL") from None
|