Add backup and restore of data

app.py

@@ -1,17 +1,20 @@
 import os
+import shutil
 import sqlite3
+import tempfile
 import uuid
+import zipfile
 from datetime import datetime, timezone
 from pathlib import Path
 
-from flask import Flask, jsonify, render_template, request, send_from_directory
+from flask import Flask, jsonify, render_template, request, send_file, send_from_directory
 from werkzeug.utils import secure_filename
 
 app = Flask(__name__)
 DATABASE_PATH = os.environ.get("DATABASE_PATH", str(Path(__file__).with_name("inventory.db")))
 UPLOAD_FOLDER = Path(os.environ.get("UPLOAD_FOLDER", str(Path(DATABASE_PATH).parent / "uploads")))
 ALLOWED_IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "gif", "webp"}
-app.config["MAX_CONTENT_LENGTH"] = 10 * 1024 * 1024
+app.config["MAX_CONTENT_LENGTH"] = 250 * 1024 * 1024
 
 
 def get_db():
@@ -145,6 +148,52 @@ def files_from_request():
     return request.files.getlist("images") or request.files.getlist("image")
 
 
+def add_directory_to_zip(archive, directory, archive_prefix):
+    if not directory.exists():
+        return
+    for path in directory.rglob("*"):
+        if path.is_file():
+            archive.write(path, f"{archive_prefix}/{path.relative_to(directory).as_posix()}")
+
+
+def valid_backup_member(member):
+    path = Path(member.filename)
+    return not path.is_absolute() and ".." not in path.parts and (member.filename == "inventory.db" or member.filename.startswith("uploads/"))
+
+
+def restore_backup(backup):
+    with tempfile.TemporaryDirectory() as temporary_directory:
+        temporary_path = Path(temporary_directory)
+        archive_path = temporary_path / "backup.zip"
+        backup.save(archive_path)
+        with zipfile.ZipFile(archive_path) as archive:
+            members = archive.infolist()
+            if not members or not all(valid_backup_member(member) for member in members):
+                raise ValueError("The backup file has an invalid format.")
+            if "inventory.db" not in archive.namelist():
+                raise ValueError("The backup does not contain inventory.db.")
+            database_path = temporary_path / "inventory.db"
+            uploads_path = temporary_path / "uploads"
+            archive.extract("inventory.db", temporary_path)
+            archive.extractall(temporary_path, [member for member in members if member.filename.startswith("uploads/")])
+
+        with sqlite3.connect(database_path) as connection:
+            if connection.execute("PRAGMA integrity_check").fetchone()[0] != "ok":
+                raise ValueError("The backup database failed its integrity check.")
+        Path(DATABASE_PATH).parent.mkdir(parents=True, exist_ok=True)
+        staged_database_path = Path(DATABASE_PATH).with_name(f".{Path(DATABASE_PATH).name}.restore")
+        shutil.copy2(database_path, staged_database_path)
+        os.replace(staged_database_path, DATABASE_PATH)
+        if UPLOAD_FOLDER.exists():
+            shutil.rmtree(UPLOAD_FOLDER)
+        if uploads_path.exists():
+            staged_uploads_path = UPLOAD_FOLDER.parent / f".{UPLOAD_FOLDER.name}.restore"
+            shutil.copytree(uploads_path, staged_uploads_path, dirs_exist_ok=True)
+            os.replace(staged_uploads_path, UPLOAD_FOLDER)
+        else:
+            UPLOAD_FOLDER.mkdir(parents=True, exist_ok=True)
+
+
 @app.route("/")
 def index():
     return render_template("index.html")
@@ -187,6 +236,38 @@ def summary():
     return jsonify({"total": total, "available": available, "checked_out": checked_out, "maintenance": maintenance})
 
 
+@app.get("/api/backup")
+def backup():
+    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as temporary_database:
+        temporary_database_path = Path(temporary_database.name)
+    try:
+        with get_db() as source, sqlite3.connect(temporary_database_path) as destination:
+            source.backup(destination)
+        archive = tempfile.NamedTemporaryFile(suffix=".zip", delete=False)
+        archive_path = Path(archive.name)
+        archive.close()
+        with zipfile.ZipFile(archive_path, "w", zipfile.ZIP_DEFLATED) as zip_archive:
+            zip_archive.write(temporary_database_path, "inventory.db")
+            add_directory_to_zip(zip_archive, UPLOAD_FOLDER, "uploads")
+        response = send_file(archive_path, as_attachment=True, download_name=f"stockroom-backup-{datetime.now().strftime('%Y%m%d-%H%M%S')}.zip", mimetype="application/zip")
+        response.call_on_close(lambda: archive_path.unlink(missing_ok=True))
+        return response
+    finally:
+        temporary_database_path.unlink(missing_ok=True)
+
+
+@app.post("/api/restore")
+def restore():
+    backup_file = request.files.get("backup")
+    if not backup_file or not backup_file.filename:
+        return jsonify({"error": "Choose a backup ZIP file to restore."}), 400
+    try:
+        restore_backup(backup_file)
+    except (ValueError, zipfile.BadZipFile, sqlite3.DatabaseError) as error:
+        return jsonify({"error": str(error) or "The backup could not be restored."}), 400
+    return jsonify({"restored": True})
+
+
 @app.post("/api/items")
 def create_item():
     data = item_payload(request_data())