444 lines
14 KiB
Python
444 lines
14 KiB
Python
import os
|
||
import zipfile
|
||
import uuid
|
||
from datetime import datetime
|
||
from werkzeug.utils import secure_filename
|
||
from werkzeug.datastructures import FileStorage
|
||
from flask import (
|
||
Blueprint, request, redirect, url_for,
|
||
flash, send_from_directory, current_app,
|
||
jsonify, abort
|
||
)
|
||
from flask_login import login_required, current_user
|
||
from PIL import Image, UnidentifiedImageError
|
||
|
||
from app import db
|
||
from .models import Media, ZipJob, ImageHeart, FeaturedImage
|
||
from .tasks import process_zip
|
||
|
||
bp = Blueprint(
|
||
"media",
|
||
__name__,
|
||
url_prefix="/media",
|
||
template_folder="templates"
|
||
)
|
||
|
||
# ─── Constants ──────────────────────────────────────────────────────────────
|
||
|
||
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp"}
|
||
DOC_EXTS = {".pdf", ".txt", ".csv"}
|
||
ZIP_EXT = ".zip"
|
||
MAX_ZIP_FILES = 1000
|
||
MAX_IMAGE_PIXELS = 8000 * 8000 # ~64M pixels
|
||
|
||
# ─── Context Processor ──────────────────────────────────────────────────────
|
||
|
||
@bp.app_context_processor
|
||
def inject_helpers():
|
||
"""Expose generate_image_url in all media templates."""
|
||
return dict(generate_image_url=generate_image_url)
|
||
|
||
# ─── Helper Functions ───────────────────────────────────────────────────────
|
||
|
||
def allowed_file(filename: str) -> bool:
|
||
"""
|
||
Return True if the file extension is allowed.
|
||
"""
|
||
ext = os.path.splitext(filename)[1].lower()
|
||
allowed = current_app.config.get(
|
||
"ALLOWED_EXTENSIONS",
|
||
IMAGE_EXTS | DOC_EXTS | {ZIP_EXT}
|
||
)
|
||
return ext in allowed
|
||
|
||
def get_upload_path(plugin: str, related_id: int) -> (str, str):
|
||
"""
|
||
Build and return (absolute_dir, relative_subdir) under UPLOAD_FOLDER.
|
||
"""
|
||
base = current_app.config["UPLOAD_FOLDER"]
|
||
subdir = os.path.join(plugin, str(related_id))
|
||
abs_dir = os.path.abspath(os.path.join(base, subdir))
|
||
if not abs_dir.startswith(os.path.abspath(base) + os.sep):
|
||
raise RuntimeError("Upload path escapes base directory")
|
||
os.makedirs(abs_dir, exist_ok=True)
|
||
return abs_dir, subdir
|
||
|
||
def validate_image(path: str) -> bool:
|
||
"""
|
||
Verify image integrity and enforce pixel-size limit.
|
||
"""
|
||
try:
|
||
with Image.open(path) as img:
|
||
img.verify()
|
||
w, h = Image.open(path).size
|
||
return w * h <= MAX_IMAGE_PIXELS
|
||
except (UnidentifiedImageError, IOError):
|
||
return False
|
||
|
||
def validate_pdf(path: str) -> bool:
|
||
"""
|
||
Quick header check for PDF files.
|
||
"""
|
||
try:
|
||
with open(path, "rb") as f:
|
||
return f.read(5) == b"%PDF-"
|
||
except IOError:
|
||
return False
|
||
|
||
def validate_text(path: str) -> bool:
|
||
"""
|
||
Ensure the file is valid UTF-8 text/CSV.
|
||
"""
|
||
try:
|
||
with open(path, "rb") as f:
|
||
f.read(1024).decode("utf-8")
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
def strip_exif(image: Image.Image) -> Image.Image:
|
||
"""
|
||
Rotate per EXIF orientation and strip metadata.
|
||
"""
|
||
try:
|
||
exif = image._getexif()
|
||
if exif:
|
||
orientation_key = next(
|
||
(k for k, v in Image.ExifTags.TAGS.items()
|
||
if v == "Orientation"), None
|
||
)
|
||
o = exif.get(orientation_key)
|
||
if o == 3:
|
||
image = image.rotate(180, expand=True)
|
||
elif o == 6:
|
||
image = image.rotate(270, expand=True)
|
||
elif o == 8:
|
||
image = image.rotate(90, expand=True)
|
||
except Exception:
|
||
pass
|
||
data = list(image.getdata())
|
||
clean = Image.new(image.mode, image.size)
|
||
clean.putdata(data)
|
||
return clean
|
||
|
||
def generate_image_url(media: Media):
|
||
"""
|
||
Given a Media instance, return its public URL or a placeholder.
|
||
"""
|
||
if media and media.file_url:
|
||
return url_for(
|
||
"media.serve_context_media",
|
||
context=media.plugin,
|
||
context_id=media.related_id,
|
||
filename=media.filename
|
||
)
|
||
# fallback placeholder
|
||
w, h = current_app.config.get("STANDARD_IMG_SIZE", (300, 200))
|
||
return f"https://placehold.co/{w}x{h}"
|
||
|
||
# ─── Core Media Routes ──────────────────────────────────────────────────────
|
||
|
||
@bp.route("/upload", methods=["POST"])
|
||
@login_required
|
||
def upload_media():
|
||
"""
|
||
Accept images, PDFs, text/CSV inline; enqueue ZIPs for async processing.
|
||
"""
|
||
uploaded: FileStorage = request.files.get("media")
|
||
if not uploaded or uploaded.filename == "":
|
||
flash("No file selected.", "warning")
|
||
return redirect(request.referrer or url_for("home"))
|
||
|
||
filename = secure_filename(uploaded.filename)
|
||
ext = os.path.splitext(filename)[1].lower()
|
||
if not allowed_file(filename):
|
||
flash("Unsupported file type.", "danger")
|
||
return redirect(request.referrer)
|
||
|
||
# Determine plugin & ID
|
||
plugin = request.form.get("plugin", "user")
|
||
related_id = int(request.form.get("related_id", current_user.id))
|
||
|
||
# Save location
|
||
abs_dir, subdir = get_upload_path(plugin, related_id)
|
||
save_path = os.path.join(abs_dir, filename)
|
||
uploaded.save(save_path)
|
||
|
||
# Validate & post-process
|
||
if ext in IMAGE_EXTS:
|
||
if not validate_image(save_path):
|
||
os.remove(save_path)
|
||
flash("Invalid or oversized image.", "danger")
|
||
return redirect(request.referrer)
|
||
with Image.open(save_path) as img:
|
||
clean = strip_exif(img)
|
||
clean.save(save_path)
|
||
|
||
elif ext == ".pdf":
|
||
if not validate_pdf(save_path):
|
||
os.remove(save_path)
|
||
flash("Invalid PDF.", "danger")
|
||
return redirect(request.referrer)
|
||
|
||
elif ext in {".txt", ".csv"}:
|
||
if not validate_text(save_path):
|
||
os.remove(save_path)
|
||
flash("Invalid text/CSV.", "danger")
|
||
return redirect(request.referrer)
|
||
|
||
elif ext == ZIP_EXT:
|
||
# Create and enqueue a ZipJob
|
||
job = ZipJob(user_id=current_user.id, filename=filename)
|
||
db.session.add(job)
|
||
db.session.commit()
|
||
process_zip.delay(job.id, save_path)
|
||
flash("ZIP received; processing in background.", "info")
|
||
return redirect(url_for("media.upload_status", job_id=job.id))
|
||
|
||
# Record small-file upload in DB
|
||
media = Media(
|
||
plugin = plugin,
|
||
related_id = related_id,
|
||
filename = filename,
|
||
file_url = f"{subdir}/{filename}",
|
||
uploader_id = current_user.id,
|
||
uploaded_at = datetime.utcnow()
|
||
)
|
||
db.session.add(media)
|
||
db.session.commit()
|
||
flash("File uploaded successfully.", "success")
|
||
return redirect(request.referrer or url_for("home"))
|
||
|
||
@bp.route("/upload/<int:job_id>/status", methods=["GET"])
|
||
@login_required
|
||
def upload_status(job_id: int):
|
||
"""
|
||
Return JSON status for a background ZIP processing job.
|
||
"""
|
||
job = ZipJob.query.get_or_404(job_id)
|
||
if job.user_id != current_user.id:
|
||
abort(403)
|
||
return jsonify({
|
||
"job_id": job.id,
|
||
"status": job.status,
|
||
"error": job.error
|
||
})
|
||
|
||
@bp.route("/<context>/<int:context_id>/<filename>")
|
||
def serve_context_media(context: str, context_id: int, filename: str):
|
||
"""
|
||
Serve a file from UPLOAD_FOLDER/<plugin>/<id>/<filename>,
|
||
with path‐traversal guard and DB check.
|
||
"""
|
||
# Normalize plugin name
|
||
valid = {"user", "plant", "growlog", "vendor"}
|
||
if context in valid:
|
||
plugin_name = context
|
||
elif context.endswith("s") and context[:-1] in valid:
|
||
plugin_name = context[:-1]
|
||
else:
|
||
abort(404)
|
||
|
||
# Sanitize filename
|
||
safe_filename = secure_filename(filename)
|
||
if safe_filename != filename:
|
||
abort(404)
|
||
|
||
# Build and verify path
|
||
base_dir = current_app.config["UPLOAD_FOLDER"]
|
||
dir_path = os.path.join(base_dir, plugin_name, str(context_id))
|
||
full_path = os.path.abspath(os.path.join(dir_path, safe_filename))
|
||
if not full_path.startswith(os.path.abspath(base_dir) + os.sep):
|
||
abort(404)
|
||
|
||
# Confirm DB row
|
||
Media.query.filter_by(
|
||
plugin = plugin_name,
|
||
related_id = context_id,
|
||
filename = filename
|
||
).first_or_404()
|
||
|
||
return send_from_directory(dir_path, filename)
|
||
|
||
# ─── Utility Routes ─────────────────────────────────────────────────────────
|
||
|
||
@bp.route("/heart/<int:media_id>", methods=["POST"])
|
||
@login_required
|
||
def toggle_heart(media_id: int):
|
||
"""
|
||
Toggle a “heart” (like) on an image for the current user.
|
||
"""
|
||
existing = ImageHeart.query.filter_by(
|
||
user_id = current_user.id,
|
||
media_id = media_id
|
||
).first()
|
||
if existing:
|
||
db.session.delete(existing)
|
||
db.session.commit()
|
||
return jsonify(status="unhearted")
|
||
heart = ImageHeart(user_id=current_user.id, media_id=media_id)
|
||
db.session.add(heart)
|
||
db.session.commit()
|
||
return jsonify(status="hearted")
|
||
|
||
@bp.route("/featured/<context>/<int:context_id>/<int:media_id>", methods=["POST"])
|
||
@login_required
|
||
def set_featured_image(context: str, context_id: int, media_id: int):
|
||
"""
|
||
Mark a single image as featured for a given context.
|
||
"""
|
||
valid = {"plant", "growlog", "user", "vendor"}
|
||
if context in valid:
|
||
plugin_name = context
|
||
elif context.endswith("s") and context[:-1] in valid:
|
||
plugin_name = context[:-1]
|
||
else:
|
||
abort(404)
|
||
|
||
media = Media.query.filter_by(
|
||
plugin = plugin_name,
|
||
related_id = context_id,
|
||
id = media_id
|
||
).first_or_404()
|
||
|
||
if media.uploader_id != current_user.id and current_user.role != "admin":
|
||
abort(403)
|
||
|
||
FeaturedImage.query.filter_by(
|
||
context = plugin_name,
|
||
context_id = context_id
|
||
).delete()
|
||
fi = FeaturedImage(
|
||
media_id = media.id,
|
||
context = plugin_name,
|
||
context_id = context_id,
|
||
is_featured = True
|
||
)
|
||
db.session.add(fi)
|
||
db.session.commit()
|
||
flash("Featured image updated.", "success")
|
||
return redirect(request.referrer or url_for("home"))
|
||
|
||
@bp.route("/delete/<int:media_id>", methods=["POST"])
|
||
@login_required
|
||
def delete_media(media_id: int):
|
||
"""
|
||
Delete a media file and its DB record (soft‐delete by permission).
|
||
"""
|
||
media = Media.query.get_or_404(media_id)
|
||
if media.uploader_id != current_user.id and current_user.role != "admin":
|
||
flash("Not authorized to delete this media.", "danger")
|
||
return redirect(request.referrer or url_for("home"))
|
||
|
||
# Remove file on disk
|
||
base = current_app.config["UPLOAD_FOLDER"]
|
||
full = os.path.abspath(os.path.join(base, media.file_url))
|
||
try:
|
||
os.remove(full)
|
||
except OSError:
|
||
current_app.logger.error(f"Failed to delete file {full}")
|
||
|
||
# Remove DB record
|
||
db.session.delete(media)
|
||
db.session.commit()
|
||
flash("Media deleted.", "success")
|
||
return redirect(request.referrer or url_for("home"))
|
||
|
||
@bp.route("/rotate/<int:media_id>", methods=["POST"])
|
||
@login_required
|
||
def rotate_media(media_id: int):
|
||
"""
|
||
Rotate an image −90° and strip its EXIF metadata.
|
||
"""
|
||
media = Media.query.get_or_404(media_id)
|
||
if media.uploader_id != current_user.id and current_user.role != "admin":
|
||
abort(403)
|
||
|
||
base = current_app.config["UPLOAD_FOLDER"]
|
||
full = os.path.abspath(os.path.join(base, media.file_url))
|
||
try:
|
||
with Image.open(full) as img:
|
||
rotated = img.rotate(-90, expand=True)
|
||
clean = strip_exif(rotated)
|
||
clean.save(full)
|
||
flash("Image rotated successfully.", "success")
|
||
except Exception as e:
|
||
current_app.logger.error(f"Rotation failed for {full}: {e}")
|
||
flash("Failed to rotate image.", "danger")
|
||
|
||
return redirect(request.referrer or url_for("home"))
|
||
|
||
# ─── Legacy Helpers for Other Plugins ───────────────────────────────────────
|
||
|
||
def _process_upload_file(
|
||
file: FileStorage,
|
||
uploader_id: int,
|
||
plugin: str = '',
|
||
related_id: int = 0,
|
||
plant_id=None,
|
||
growlog_id=None,
|
||
caption=None
|
||
) -> Media:
|
||
"""
|
||
Handles saving an uploaded file under:
|
||
<UPLOAD_FOLDER>/<plugin>/<related_id>/<uuid>.<ext>
|
||
and returns a new Media object (not yet committed).
|
||
"""
|
||
# 1) Unique filename (uuid + original extension)
|
||
uid = str(uuid.uuid4()).replace('-', '')
|
||
_, ext = os.path.splitext(file.filename)
|
||
filename = f"{uid}{ext.lower()}"
|
||
|
||
# 2) Build the folder
|
||
abs_dir, subdir = get_upload_path(plugin, related_id)
|
||
|
||
# 3) Save to disk
|
||
full_path = os.path.join(abs_dir, filename)
|
||
file.save(full_path)
|
||
|
||
# 4) Record the relative URL fragment (for lookup)
|
||
file_url = f"{subdir}/{filename}"
|
||
|
||
# 5) Build the Media row
|
||
now = datetime.utcnow()
|
||
media = Media(
|
||
plugin = plugin,
|
||
related_id = related_id,
|
||
filename = filename,
|
||
uploaded_at = now,
|
||
uploader_id = uploader_id,
|
||
caption = caption,
|
||
plant_id = plant_id,
|
||
growlog_id = growlog_id,
|
||
created_at = now,
|
||
file_url = file_url
|
||
)
|
||
return media
|
||
|
||
def save_media_file(file: FileStorage, user_id: int, **ctx) -> Media:
|
||
"""
|
||
Simple wrapper for other plugins to save an upload via the same logic.
|
||
"""
|
||
return _process_upload_file(file, user_id, **ctx)
|
||
|
||
def delete_media_file(media: Media):
|
||
"""
|
||
Remove a Media record and its file from disk, commit immediately.
|
||
"""
|
||
base = current_app.config["UPLOAD_FOLDER"]
|
||
full = os.path.normpath(os.path.join(base, media.file_url))
|
||
if os.path.exists(full):
|
||
os.remove(full)
|
||
db.session.delete(media)
|
||
db.session.commit()
|
||
|
||
def rotate_media_file(media: Media):
|
||
"""
|
||
Rotate a Media file −90° in place and commit metadata-only change.
|
||
"""
|
||
base = current_app.config["UPLOAD_FOLDER"]
|
||
full = os.path.normpath(os.path.join(base, media.file_url))
|
||
with Image.open(full) as img:
|
||
img.rotate(-90, expand=True).save(full)
|
||
db.session.commit()
|