diff --git a/.env.example b/.env.example index 150aee9..f12363f 100644 --- a/.env.example +++ b/.env.example @@ -1,5 +1,5 @@ # Application configuration -MAX_CONTENT_LENGTH=1048576 +MAX_CONTENT_LENGTH=3145728 MAX_PENDING_TIME=600 MAX_STORE_TIME=259200 CLEAR_AT_RESTART=1 @@ -7,7 +7,6 @@ SKIP_IHDR_FILL=0 # Flask configuration FLASK_DEBUG=0 -FLASK_ENV=development # Redis configuration RQ_DASHBOARD_REDIS_URL=redis://redis:6379/0 diff --git a/Dockerfile b/Dockerfile index 4055aa6..522d13a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -49,4 +49,4 @@ COPY aperisolve/ /app/aperisolve/ RUN pip install --no-cache-dir -r /app/aperisolve/requirements.txt # Commande de lancement -CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:5000", "--access-logfile", "-", "--error-logfile", "-", "--log-level", "info", "--capture-output", "aperisolve.wsgi:application"] +CMD ["gunicorn", "-w", "8", "-b", "0.0.0.0:5000", "--access-logfile", "-", "--error-logfile", "-", "--log-level", "info", "--capture-output", "aperisolve.utils.wsgi:application"] diff --git a/aperisolve/__init__.py b/aperisolve/__init__.py index 538f764..6659e63 100644 --- a/aperisolve/__init__.py +++ b/aperisolve/__init__.py @@ -1,3 +1,6 @@ # flake8: noqa: E203,E501,W503 # pylint: disable=C0413,W0718,R0903,R0801 # mypy: disable-error-code=unused-awaitable +from . import app, config, models, workers + +__all__ = ["app", "config", "models", "workers"] diff --git a/aperisolve/analyzers/base_analyzer.py b/aperisolve/analyzers/base_analyzer.py index 42edcb7..6597f99 100644 --- a/aperisolve/analyzers/base_analyzer.py +++ b/aperisolve/analyzers/base_analyzer.py @@ -13,8 +13,9 @@ from pathlib import Path from shutil import rmtree from typing import Any, Optional, overload +from ..config import MAX_PENDING_TIME + _thread_lock = threading.Lock() -MAX_PENDING_TIME = int(os.getenv("MAX_PENDING_TIME", "600")) # 10 minutes by default class SubprocessAnalyzer(ABC): diff --git a/aperisolve/app.py b/aperisolve/app.py index ae869ea..0927ce0 100644 --- a/aperisolve/app.py +++ b/aperisolve/app.py @@ -16,270 +16,255 @@ from flask import Flask, Response, abort, jsonify, render_template, request, sen from redis import Redis from rq import Queue -from .sentry import initialize_sentry +from .utils.sentry import initialize_sentry initialize_sentry() -from .cleanup import cleanup_old_entries from .config import IMAGE_EXTENSIONS, RESULT_FOLDER, WORKER_FILES -from .models import Image, Submission, db, init_db - -app: Flask = Flask(__name__) -app.json.sort_keys = False # type: ignore -app.config["SQLALCHEMY_DATABASE_URI"] = os.environ.get("DB_URI") -app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False -app.config["MAX_CONTENT_LENGTH"] = int(os.environ.get("MAX_CONTENT_LENGTH", 1024 * 1024)) -RESULT_FOLDER.mkdir(parents=True, exist_ok=True) - -db.init_app(app) -with app.app_context(): - init_db(app) - -redis_conn = Redis(host="redis", port=6379) -queue = Queue(connection=redis_conn) +from .models import Image, Submission, cleanup_old_entries, db -@app.errorhandler(413) -def too_large(_: Any) -> tuple[Response, int]: - """Error Handler for Max File Size.""" - sentry_sdk.set_context( - "upload", - { - "content_length": request.content_length, - "max_allowed": app.config["MAX_CONTENT_LENGTH"], - }, - ) - sentry_sdk.capture_message("Upload failed: 413 Payload Too Large", level="warning") - return jsonify({"error": "Image size exceeded"}), 413 +def create_app() -> Flask: + """Create flask application with routes.""" + app = Flask(__name__) + app.json.sort_keys = False # type: ignore + app.config["SQLALCHEMY_DATABASE_URI"] = os.environ.get("DB_URI") + app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False + app.config["MAX_CONTENT_LENGTH"] = int(os.environ.get("MAX_CONTENT_LENGTH", 1024 * 1024)) + app.config["REDIS_QUEUE"] = Queue(connection=Redis(host="redis", port=6379)) + db.init_app(app) + @app.errorhandler(413) + def too_large(_: Any) -> tuple[Response, int]: + """Error Handler for Max File Size.""" + sentry_sdk.set_context( + "upload", + { + "content_length": request.content_length, + "max_allowed": app.config["MAX_CONTENT_LENGTH"], + }, + ) + sentry_sdk.capture_message("Upload failed: 413 Payload Too Large", level="warning") + return jsonify({"error": "Image size exceeded"}), 413 -@app.errorhandler(404) -def not_found(_: Any) -> str: - """Error Handler for 404 not found.""" - return render_template("error.html", message="Resource not found", statuscode=404) + @app.errorhandler(404) + def not_found(_: Any) -> str: + """Error Handler for 404 not found.""" + return render_template("error.html", message="Resource not found", statuscode=404) + @app.route("/") + def index() -> str: + """Render the main index page.""" + return render_template("index.html") -# Sentry automatically captures 500 errors + @app.route("/faq") + def faq() -> str: + """Render the FAQ page.""" + return render_template("faq.html") + @app.route("/show") + def show() -> str: + """Show all active submissions.""" + db_images = Image.query.all() + images = [] + for img in db_images: + last_sub = img.submissions[-1] + images.append({"file": f"/image/{img.hash}", "link": f"/{last_sub.hash}"}) + return render_template("show.html", images=images) -@app.route("/") -def index() -> str: - """Render the main index page.""" - return render_template("index.html") + @app.route("/", methods=["GET"]) + def result_page(hash_val: str) -> str: + """Render the result page for a given submission hash.""" + submission = Submission.query.get_or_404(hash_val) # type: ignore + Image.query.get_or_404(submission.image_hash) # type: ignore + return render_template("result.html", hash_val=hash_val) -@app.route("/faq") -def faq() -> str: - """Render the FAQ page.""" - return render_template("faq.html") + @app.route("/upload", methods=["POST"]) + def upload_image() -> tuple[Response, int]: + """Handle image upload and initiate analysis.""" + cleanup_old_entries() -@app.route("/show") -def show() -> str: - """Show all active submissions.""" - db_images = Image.query.all() - images = [] - for img in db_images: - last_sub = img.submissions[-1] - images.append({"file": f"/image/{img.hash}", "link": f"/{last_sub.hash}"}) - return render_template("show.html", images=images) + if "image" not in request.files: + return jsonify({"error": "No image provided"}), 400 + image = request.files["image"] + password = request.form.get("password") + deep_analysis = request.form.get("deep") == "true" -@app.route("/", methods=["GET"]) -def result_page(hash_val: str) -> str: - """Render the result page for a given submission hash.""" - submission = Submission.query.get_or_404(hash_val) # type: ignore - Image.query.get_or_404(submission.image_hash) # type: ignore + if image.filename is None or image.filename == "": + return jsonify({"error": "No image provided"}), 400 - return render_template("result.html", hash_val=hash_val) + if "." not in image.filename or Path(image.filename).suffix.lower() not in IMAGE_EXTENSIONS: + return jsonify({"error": "Unsupported file type"}), 400 + # Compute hashes and prepare submission + image_data = image.read() + img_hash = hashlib.md5(image_data).hexdigest() + image_name = img_hash + "." + image.filename.rsplit(".", maxsplit=1)[-1].lower() -@app.route("/upload", methods=["POST"]) -def upload_image() -> tuple[Response, int]: - """Handle image upload and initiate analysis.""" + submission_data = ( + image_data + + image.filename.encode() + + (password.encode() if password else b"") + + (b"deep_analysis" if deep_analysis else b"") + ) + submission_hash = hashlib.md5(submission_data).hexdigest() + submission_path = RESULT_FOLDER / img_hash / submission_hash - cleanup_old_entries() + # Check if hash is present on the DB + submission: Submission = Submission.query.filter_by(hash=submission_hash).first() - if "image" not in request.files: - return jsonify({"error": "No image provided"}), 400 + # Only return if BOTH the file exists AND the DB entry exists + if submission_path.exists() and submission is not None: + return jsonify({"submission_hash": submission.hash}), 200 - image = request.files["image"] - password = request.form.get("password") - deep_analysis = request.form.get("deep") == "true" + # Self-Healing Logic + new_img_path = RESULT_FOLDER / img_hash / image_name - if image.filename is None or image.filename == "": - return jsonify({"error": "No image provided"}), 400 + # This fixes the "Ghost File" issue (DB exists, but file is missing) + if not new_img_path.exists(): + new_img_path.parent.mkdir(parents=True, exist_ok=True) + with open(new_img_path, "wb") as f: + f.write(image_data) - if "." not in image.filename or Path(image.filename).suffix.lower() not in IMAGE_EXTENSIONS: - return jsonify({"error": "Unsupported file type"}), 400 + # Try to find image in DB by hash + sub_img = Image.query.filter_by(hash=img_hash).first() - # Compute hashes and prepare submission - image_data = image.read() - img_hash = hashlib.md5(image_data).hexdigest() - image_name = img_hash + "." + image.filename.rsplit(".", maxsplit=1)[-1].lower() + # If DB entry is missing (even if file exists on disk), create it! + if sub_img is None: + sub_img = Image( + file=str(new_img_path), + hash=img_hash, + size=len(image_data), + upload_count=0, + first_submission_date=datetime.now(timezone.utc), + last_submission_date=datetime.now(timezone.utc), + ) + db.session.add(sub_img) # pylint: disable=no-member + db.session.commit() # pylint: disable=no-member - submission_data = ( - image_data - + image.filename.encode() - + (password.encode() if password else b"") - + (b"deep_analysis" if deep_analysis else b"") - ) - submission_hash = hashlib.md5(submission_data).hexdigest() - submission_path = RESULT_FOLDER / img_hash / submission_hash + # Increment upload count for the image + sub_img.upload_count += 1 + db.session.commit() # pylint: disable=no-member - # Check if hash is present on the DB - submission: Submission = Submission.query.filter_by(hash=submission_hash).first() + # Create new Submission entry + submission_path.mkdir(parents=True, exist_ok=True) - # Only return if BOTH the file exists AND the DB entry exists - if submission_path.exists() and submission is not None: + # Only create a new submission if it does not already exist + # New submission for new image + if not submission: + submission = Submission( + filename=image.filename, + password=password, + deep_analysis=deep_analysis, + hash=submission_hash, + status="pending", + date=time.time(), + image_hash=sub_img.hash, + ) + db.session.add(submission) # pylint: disable=no-member + db.session.commit() # pylint: disable=no-member + + # Re-analysis case to prevent early return + submission.status = "pending" # type: ignore + db.session.commit() # pylint: disable=no-member + + # Start the analysis jobs + app.config["REDIS_QUEUE"].enqueue( + "aperisolve.workers.analyze_image", submission.hash, job_timeout=300 + ) + + # Return submission identifier return jsonify({"submission_hash": submission.hash}), 200 - # Self-Healing Logic - new_img_path = RESULT_FOLDER / img_hash / image_name + @app.route("/status/", methods=["GET"]) + def get_status(hash_val: str) -> Response: + """Get the processing status of a submission.""" + submission = Submission.query.get_or_404(hash_val) # type: ignore + return jsonify({"status": submission.status}) # type: ignore - # This fixes the "Ghost File" issue (DB exists, but file is missing) - if not new_img_path.exists(): - new_img_path.parent.mkdir(parents=True, exist_ok=True) - with open(new_img_path, "wb") as f: - f.write(image_data) + @app.route("/infos/", methods=["GET"]) + def get_infos(hash_val: str) -> Response: + """Get the metadata and information of a submission.""" + submission = Submission.query.get_or_404(hash_val) + image = Image.query.get_or_404(submission.image_hash) # type: ignore + names = [name for name in set(sub.filename for sub in image.submissions) if name] + passwords = [pwd for pwd in set(sub.password for sub in image.submissions) if pwd] + return jsonify( + { + "image_path": "image/" + str(Path(image.file).name), + "names": names, + "size": image.size, + "first_submission_date": image.first_submission_date, + "last_submission_date": image.last_submission_date, + "upload_count": image.upload_count, + "passwords": passwords, + } + ) # type: ignore - # Try to find image in DB by hash - sub_img = Image.query.filter_by(hash=img_hash).first() + @app.route("/result/", methods=["GET"]) + def get_result(hash_val: str) -> tuple[Response, int]: + """Get the analysis results of a submission.""" + submission = Submission.query.get_or_404(hash_val) + image = Image.query.get_or_404(submission.image_hash) # type: ignore - # If DB entry is missing (even if file exists on disk), create it! - if sub_img is None: - sub_img = Image( - file=str(new_img_path), - hash=img_hash, - size=len(image_data), - upload_count=0, - first_submission_date=datetime.now(timezone.utc), - last_submission_date=datetime.now(timezone.utc), - ) - db.session.add(sub_img) # pylint: disable=no-member - db.session.commit() # pylint: disable=no-member + result_path = RESULT_FOLDER / str(image.hash) / str(submission.hash) / "results.json" - # Increment upload count for the image - sub_img.upload_count += 1 - db.session.commit() # pylint: disable=no-member + if not result_path.exists(): + return jsonify({"error": "Results not ready yet..."}), 425 - # Create new Submission entry - submission_path.mkdir(parents=True, exist_ok=True) + with open(result_path, "r", encoding="utf-8") as f: + results = json.load(f) - # Only create a new submission if it does not already exist - # New submission for new image - if not submission: - submission = Submission( - filename=image.filename, - password=password, - deep_analysis=deep_analysis, - hash=submission_hash, - status="pending", - date=time.time(), - image_hash=sub_img.hash, - ) - db.session.add(submission) # pylint: disable=no-member - db.session.commit() # pylint: disable=no-member + return jsonify({"results": results}), 200 - # Re-analysis case to prevent early return - submission.status = "pending" # type: ignore - db.session.commit() # pylint: disable=no-member + @app.route("/download//") + def download_output(hash_val: str, tool: str) -> Response: + """Download the output of a specific analyzer for a given submission hash.""" - # Start the analysis jobs - queue.enqueue("aperisolve.workers.analyze_image", submission.hash, job_timeout=300) - - # Return submission identifier - return jsonify({"submission_hash": submission.hash}), 200 - - -@app.route("/status/", methods=["GET"]) -def get_status(hash_val: str) -> Response: - """Get the processing status of a submission.""" - submission = Submission.query.get_or_404(hash_val) # type: ignore - return jsonify({"status": submission.status}) # type: ignore - - -@app.route("/infos/", methods=["GET"]) -def get_infos(hash_val: str) -> Response: - """Get the metadata and information of a submission.""" - submission = Submission.query.get_or_404(hash_val) - image = Image.query.get_or_404(submission.image_hash) # type: ignore - names = [name for name in set(sub.filename for sub in image.submissions) if name] - passwords = [pwd for pwd in set(sub.password for sub in image.submissions) if pwd] - return jsonify( - { - "image_path": "image/" + str(Path(image.file).name), - "names": names, - "size": image.size, - "first_submission_date": image.first_submission_date, - "last_submission_date": image.last_submission_date, - "upload_count": image.upload_count, - "passwords": passwords, - } - ) # type: ignore - - -@app.route("/result/", methods=["GET"]) -def get_result(hash_val: str) -> tuple[Response, int]: - """Get the analysis results of a submission.""" - submission = Submission.query.get_or_404(hash_val) - image = Image.query.get_or_404(submission.image_hash) # type: ignore - - result_path = RESULT_FOLDER / str(image.hash) / str(submission.hash) / "results.json" - - if not result_path.exists(): - return jsonify({"error": "Results not ready yet..."}), 425 - - with open(result_path, "r", encoding="utf-8") as f: - results = json.load(f) - - return jsonify({"results": results}), 200 - - -@app.route("/download//") -def download_output(hash_val: str, tool: str) -> Response: - """Download the output of a specific analyzer for a given submission hash.""" - - submission = Submission.query.filter_by(hash=hash_val).first_or_404() - image = Image.query.get_or_404(submission.image_hash) # type: ignore - - output_dir = RESULT_FOLDER / str(image.hash) / str(submission.hash) - output_file = output_dir / f"{tool}.7z" - - if tool not in WORKER_FILES or not output_file.exists(): - abort(404, description="Tool output not found.") - - return send_file(output_file, as_attachment=True) - - -@app.route("/image/") -@app.route("/image//") -def get_image(hash_val: Optional[str] = None, img_name: Optional[str] = None) -> Response: - """ - Download the image with a specific name (usually ones generated with - decomposer) for a given submission hash. - """ - - if img_name is None: - return abort(404, description="Image not found or unsupported format") - - if hash_val is not None: submission = Submission.query.filter_by(hash=hash_val).first_or_404() image = Image.query.get_or_404(submission.image_hash) # type: ignore output_dir = RESULT_FOLDER / str(image.hash) / str(submission.hash) - output_file = output_dir / Path(img_name).name - else: - hash_val = img_name.split(".")[0] # Extract hash from the image name - image = Image.query.filter_by(hash=hash_val).first_or_404() - output_file = Path(image.file) # type: ignore + output_file = output_dir / f"{tool}.7z" - if (not output_file.exists()) or (output_file.suffix.lower() not in IMAGE_EXTENSIONS): - return abort(404, description="Image not found or unsupported format") + if tool not in WORKER_FILES or not output_file.exists(): + abort(404, description="Tool output not found.") - return send_file(output_file, as_attachment=True) + return send_file(output_file, as_attachment=True) + @app.route("/image/") + @app.route("/image//") + def get_image(hash_val: Optional[str] = None, img_name: Optional[str] = None) -> Response: + """ + Download the image with a specific name (usually ones generated with + decomposer) for a given submission hash. + """ + + if img_name is None: + return abort(404, description="Image not found or unsupported format") + + if hash_val is not None: + submission = Submission.query.filter_by(hash=hash_val).first_or_404() + image = Image.query.get_or_404(submission.image_hash) # type: ignore + + output_dir = RESULT_FOLDER / str(image.hash) / str(submission.hash) + output_file = output_dir / Path(img_name).name + else: + hash_val = img_name.split(".")[0] # Extract hash from the image name + image = Image.query.filter_by(hash=hash_val).first_or_404() + output_file = Path(image.file) # type: ignore + + if (not output_file.exists()) or (output_file.suffix.lower() not in IMAGE_EXTENSIONS): + return abort(404, description="Image not found or unsupported format") + + return send_file(output_file, as_attachment=True) + + return app if __name__ == "__main__": - debug_mode = os.environ.get("FLASK_DEBUG", "false").lower() == "true" - app.run(host="0.0.0.0", port=5000, debug=debug_mode) + my_app = create_app() + debug_mode = os.environ.get("FLASK_DEBUG", "0") == "1" + my_app.run(host="0.0.0.0", port=5000, debug=debug_mode) diff --git a/aperisolve/cleanup.py b/aperisolve/cleanup.py deleted file mode 100644 index 032a7b6..0000000 --- a/aperisolve/cleanup.py +++ /dev/null @@ -1,79 +0,0 @@ -# flake8: noqa: E203,E501,W503 -# pylint: disable=C0413,W0718,R0903,R0801 -# mypy: disable-error-code=unused-awaitable -"""Cleanup old entries in the database and file system.""" - -import shutil -import time -from datetime import datetime, timezone -from os import getenv -from pathlib import Path - -from .models import Image, Submission, db - -RESULT_FOLDER = Path("./results") -MAX_PENDING_TIME = int(getenv("MAX_PENDING_TIME", "600")) # 10 minutes by default -MAX_STORE_TIME = int(getenv("MAX_STORE_TIME", "259200")) # 3 days by default - - -def cleanup_old_entries() -> None: - """ - Clean up old and incomplete entries from the database and file system. - - This function performs two main cleanup operations: - - 1. Submission cleanup: - - Deletes submissions with "pending" or "running" status that have exceeded - the maximum allowed processing time (MAX_PENDING_TIME). - - Deletes completed submissions ("done" status) that have missing or buggy - results (results.json file not found). - - 2. Image cleanup: - - Deletes images older than MAX_STORE_TIME along with all associated - submissions and their result folders. - - Deletes orphaned images (with no submissions) older than MAX_PENDING_TIME - and removes their associated result folders from the file system. - - All database changes are committed at the end of the operation. - - Returns: - None - """ - now = time.time() - for submission in Submission.query.all(): # type: ignore - if ( - submission.status in ("pending", "running") # type: ignore - and now - submission.date > MAX_PENDING_TIME # type: ignore - ): - # Processing took too long, delete - db.session.delete(submission) # pylint: disable=no-member - elif submission.status == "done": # type: ignore - # Search for buggy results, delete - result_path = RESULT_FOLDER / str(submission.image_hash) / str(submission.hash) - result_file = result_path / "results.json" - if not result_file.exists(): - db.session.delete(submission) # pylint: disable=no-member - shutil.rmtree(result_path) - - # Delete "old" images - for img in Image.query.all(): # type: ignore - if img.last_submission_date.tzinfo is None: - img_date = img.last_submission_date.replace(tzinfo=timezone.utc) - else: - img_date = img.last_submission_date - delay = datetime.now(timezone.utc) - img_date - img_fold = RESULT_FOLDER / img.hash - if delay.total_seconds() > MAX_STORE_TIME: - for s in img.submissions: - db.session.delete(s) # pylint: disable=no-member - if img_fold.exists(): - shutil.rmtree(img_fold) # type: ignore - db.session.delete(img) # pylint: disable=no-member - - # Delete Images with missing submission - if len(img.submissions) == 0 and delay.total_seconds() > MAX_PENDING_TIME: - if img_fold.exists(): - shutil.rmtree(img_fold) # type: ignore - db.session.delete(img) # pylint: disable=no-member - - db.session.commit() # pylint: disable=no-member diff --git a/aperisolve/config.py b/aperisolve/config.py index 5bc34b0..40d4c94 100644 --- a/aperisolve/config.py +++ b/aperisolve/config.py @@ -3,9 +3,13 @@ # mypy: disable-error-code=unused-awaitable """Aperi'Solve configuration variables.""" +from os import getenv from pathlib import Path IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".tiff"] WORKER_FILES = ["binwalk", "foremost", "steghide", "zsteg", "image_resize", "openstego", "pcrt"] RESULT_FOLDER = Path(__file__).parent.resolve() / "results" +MAX_PENDING_TIME = int(getenv("MAX_PENDING_TIME", "600")) # 10 minutes by default +MAX_STORE_TIME = int(getenv("MAX_STORE_TIME", "259200")) # 3 days by default +CLEAR_AT_RESTART = int(getenv("CLEAR_AT_RESTART", "0")) diff --git a/aperisolve/models.py b/aperisolve/models.py index 115a74c..66b2330 100644 --- a/aperisolve/models.py +++ b/aperisolve/models.py @@ -1,18 +1,15 @@ # flake8: noqa: E203,E501,W503 -# pylint: disable=C0413,W0718,R0903,R0801 +# pylint: disable=E1101,C0413,W0212,W0718,R0903,R0801 # mypy: disable-error-code=unused-awaitable """This module defines the database models for the Aperi'Solve application.""" import itertools +import shutil import struct -import sys +import time import zlib from datetime import datetime, timezone -from os import getenv -from pathlib import Path -from shutil import rmtree -from flask import Flask from flask_sqlalchemy import SQLAlchemy from sqlalchemy import ( BigInteger, @@ -25,6 +22,7 @@ from sqlalchemy import ( String, ) +from aperisolve.config import MAX_PENDING_TIME, MAX_STORE_TIME, RESULT_FOLDER from aperisolve.utils.utils import get_resolutions, get_valid_depth_color_pairs db: SQLAlchemy = SQLAlchemy() @@ -117,110 +115,109 @@ def fill_ihdr_db() -> None: Computes CRC for each combination and stores parameters directly. """ + try: + # Skip if table already has entries (check before doing any work) + count_query = db.session.query(IHDR.iid).limit(1) + if db.session.execute(count_query).first() is not None: + print("IHDR table already populated, skipping fill.") + return - # Skip if table already has entries - if db.session.query(IHDR.iid).first() is not None: - print("IHDR table already populated, skipping fill.") - return + resolutions = get_resolutions() + bit_color_pairs = list(get_valid_depth_color_pairs()) + interlace_methods = [0, 1] - resolutions = get_resolutions() - bit_color_pairs = list(get_valid_depth_color_pairs()) - interlace_methods = [0, 1] + combinations = itertools.product( + resolutions, + bit_color_pairs, + interlace_methods, + ) + count = 0 - combinations = itertools.product( - resolutions, - bit_color_pairs, - interlace_methods, - ) - batch = [] - batch_size = 10000 # Increased from 5000 - count = 0 - with db.session.no_autoflush: - for (width, height), (bit_depth, color_type), interlace in combinations: - # Compute CRC directly - crc = IHDR.compute_crc(width, height, bit_depth, color_type, interlace) + with db.session.no_autoflush: + for (width, height), (bit_depth, color_type), interlace in combinations: + # Compute CRC directly + crc = IHDR.compute_crc(width, height, bit_depth, color_type, interlace) - # Check if entry already exists - exists = db.session.execute( - db.select(IHDR.iid).where( - db.and_( - IHDR.crc == crc, - IHDR.width == width, - IHDR.height == height, - IHDR.bit_depth == bit_depth, - IHDR.color_type == color_type, - IHDR.interlace == interlace, + # Add new entry + db.session.add( + IHDR( + crc=crc, + width=width, + height=height, + bit_depth=bit_depth, + color_type=color_type, + interlace=interlace, ) ) - ).first() - if exists: - continue + count += 1 + if count % 5000 == 0: + print(f"Inserted {count} IHDR entries...") + db.session.commit() + print(f"Inserted {count} total IHDR entries.") - batch.append( - { - "crc": crc, - "width": width, - "height": height, - "bit_depth": bit_depth, - "color_type": color_type, - "interlace": interlace, - } - ) + except Exception as e: + print(f"Error filling IHDR table: {e}") + db.session.rollback() - count += 1 - if len(batch) >= batch_size: - db.session.bulk_insert_mappings(IHDR, batch) - db.session.commit() - print(f"Inserted {count} IHDR entries...") - batch = [] - - # Insert remaining entries - if batch: - db.session.bulk_insert_mappings(IHDR, batch) - db.session.commit() - - print(f"Inserted {count} total IHDR entries.") - - -def init_db(app: Flask) -> None: +def cleanup_old_entries() -> None: """ - Initialize the database by creating tables and populating lookup data. + Clean up old and incomplete entries from the database and file system. - This function creates the database schema if it doesn't already exist, - and populates the CRC lookup table for PNG IHDR validation. + This function performs two main cleanup operations: - Args: - app (Flask): The Flask application instance used to establish the application context - for database operations. + 1. Submission cleanup: + - Deletes submissions with "pending" or "running" status that have exceeded + the maximum allowed processing time (MAX_PENDING_TIME). + - Deletes completed submissions ("done" status) that have missing or buggy + results (results.json file not found). + + 2. Image cleanup: + - Deletes images older than MAX_STORE_TIME along with all associated + submissions and their result folders. + - Deletes orphaned images (with no submissions) older than MAX_PENDING_TIME + and removes their associated result folders from the file system. + + All database changes are committed at the end of the operation. Returns: None - - Side Effects: - - Creates all database tables if the "image" table doesn't exist - - Populates the IHDR CRC lookup table with initial data - - Prints status messages to console during initialization """ - # Detect if running as RQ worker by checking command line - # Worker: /usr/local/bin/rq worker ... - # Web: /usr/local/bin/flask run ... or gunicorn/wsgi - is_worker = any("rq" in arg and "flask" not in arg for arg in sys.argv[:2]) + now = time.time() + for submission in Submission.query.all(): # type: ignore + if ( + submission.status in ("pending", "running") # type: ignore + and now - submission.date > MAX_PENDING_TIME # type: ignore + ): + # Processing took too long, delete + db.session.delete(submission) # pylint: disable=no-member + elif submission.status == "done": # type: ignore + # Search for buggy results, delete + result_path = RESULT_FOLDER / str(submission.image_hash) / str(submission.hash) + result_file = result_path / "results.json" + if not result_file.exists(): + db.session.delete(submission) # pylint: disable=no-member + shutil.rmtree(result_path) - with app.app_context(): - # Workers should never clear the database - if not is_worker and getenv("CLEAR_AT_RESTART", "0") == "1": - print("Clearing database and file system at restart...") - db.session.remove() # pylint: disable=no-member - db.drop_all() - rmtree(Path("./results"), ignore_errors=True) # Clear results folder - - # Always create tables (idempotent - safe to call multiple times) - db.create_all() - print("Database structure created successfully.") - - # Workers should never fill IHDR table - if is_worker: - print("Running as worker, skipping IHDR table fill.") + # Delete "old" images + for img in Image.query.all(): # type: ignore + if img.last_submission_date.tzinfo is None: + img_date = img.last_submission_date.replace(tzinfo=timezone.utc) else: - fill_ihdr_db() + img_date = img.last_submission_date + delay = datetime.now(timezone.utc) - img_date + img_fold = RESULT_FOLDER / img.hash + if delay.total_seconds() > MAX_STORE_TIME: + for s in img.submissions: + db.session.delete(s) # pylint: disable=no-member + if img_fold.exists(): + shutil.rmtree(img_fold) # type: ignore + db.session.delete(img) # pylint: disable=no-member + + # Delete Images with missing submission + if len(img.submissions) == 0 and delay.total_seconds() > MAX_PENDING_TIME: + if img_fold.exists(): + shutil.rmtree(img_fold) # type: ignore + db.session.delete(img) # pylint: disable=no-member + + db.session.commit() # pylint: disable=no-member diff --git a/aperisolve/utils/init_db.py b/aperisolve/utils/init_db.py new file mode 100644 index 0000000..abe89a3 --- /dev/null +++ b/aperisolve/utils/init_db.py @@ -0,0 +1,62 @@ +""" +Database initialization script. +This script is meant to be run ONCE at deploy time, never from Gunicorn or runtime code. +If CLEAR_AT_RESTART is set, database will be reset at launch time. +""" + +import os +from shutil import rmtree + +from sqlalchemy import inspect + +from ..app import create_app +from ..config import RESULT_FOLDER +from ..models import db, fill_ihdr_db + + +def main() -> None: + """Database initialization main function.""" + reset = os.getenv("CLEAR_AT_RESTART") == "1" + + app = create_app() + + with app.app_context(): + inspector = inspect(db.engine) + tables = inspector.get_table_names() + + if reset: + + # Delete result folder content + if RESULT_FOLDER.exists(): + for item in RESULT_FOLDER.iterdir(): + try: + if item.is_dir(): + rmtree(item) + else: + item.unlink() + except Exception: + continue + + print("⚠️ CLEAR_AT_RESTART=1 detected") + print("💥 Dropping all tables...") + db.drop_all() + print("🗄️ Recreating schema...") + db.create_all() + + else: + if not tables: + print("🗄️ Creating database schema...") + db.create_all() + else: + print("🗄️ Database schema already exists") + + print("📦 Filling IHDR static data...") + fill_ihdr_db() + + print("✅ Database initialization complete") + + RESULT_FOLDER.mkdir(parents=True, exist_ok=True) + + +if __name__ == "__main__": + main() diff --git a/aperisolve/utils/png.py b/aperisolve/utils/png.py index 2ae2c74..ff319c7 100644 --- a/aperisolve/utils/png.py +++ b/aperisolve/utils/png.py @@ -9,7 +9,7 @@ import struct import zlib from typing import Any, Optional -from ..app import app +from ..app import create_app from ..models import IHDR from .utils import int2hex, str2hex @@ -241,6 +241,7 @@ class PNG: crc_int = struct.unpack("!I", crc)[0] # Database lookup + app = create_app() with app.app_context(): matches = IHDR.query.filter_by(crc=crc_int).all() diff --git a/aperisolve/sentry.py b/aperisolve/utils/sentry.py similarity index 100% rename from aperisolve/sentry.py rename to aperisolve/utils/sentry.py diff --git a/aperisolve/wsgi.py b/aperisolve/utils/wsgi.py similarity index 63% rename from aperisolve/wsgi.py rename to aperisolve/utils/wsgi.py index a220be7..3a90300 100644 --- a/aperisolve/wsgi.py +++ b/aperisolve/utils/wsgi.py @@ -3,7 +3,6 @@ # mypy: disable-error-code=unused-awaitable """Aperi'Solve WSGI runner.""" -from .app import app as application +from aperisolve.app import create_app -if __name__ == "__main__": - application.run() +application = create_app() diff --git a/aperisolve/workers.py b/aperisolve/workers.py index 71b06f5..cb6e467 100644 --- a/aperisolve/workers.py +++ b/aperisolve/workers.py @@ -9,7 +9,7 @@ from typing import Any import sentry_sdk -from .sentry import initialize_sentry +from .utils.sentry import initialize_sentry initialize_sentry() @@ -24,9 +24,9 @@ from .analyzers.pngcheck import analyze_pngcheck from .analyzers.steghide import analyze_steghide from .analyzers.strings import analyze_strings from .analyzers.zsteg import analyze_zsteg -from .app import app, db +from .app import create_app from .config import RESULT_FOLDER -from .models import Image, Submission +from .models import Image, Submission, db def analyze_image(submission_hash: str) -> None: @@ -55,6 +55,7 @@ def analyze_image(submission_hash: str) -> None: - Creates result directories and generates analysis output files - Modifies the database session and commits changes """ + app = create_app() with app.app_context(): submission: Submission = Submission.query.get(submission_hash) # type: ignore image = Image.query.get_or_404(submission.image_hash) # type: ignore diff --git a/compose.dev.yml b/compose.dev.yml index 4d251b5..f095370 100644 --- a/compose.dev.yml +++ b/compose.dev.yml @@ -12,16 +12,14 @@ services: depends_on: postgres: condition: service_healthy + initdb: + condition: service_completed_successfully redis: condition: service_started worker: condition: service_started env_file: - .env.example - environment: - - FLASK_DEBUG=true - - FLASK_ENV=development - - CLEAR_AT_RESTART=0 # Don't clear on Flask debug auto-reload worker: #image: ghcr.io/zeecka/aperisolve:latest @@ -32,14 +30,27 @@ services: volumes: - ./aperisolve:/app/aperisolve depends_on: - - redis - - postgres + postgres: + condition: service_healthy + initdb: + condition: service_completed_successfully + redis: + condition: service_started env_file: - .env.example - environment: - - FLASK_DEBUG=true - - FLASK_ENV=development - - CLEAR_AT_RESTART=0 # Always false for worker to prevent double clearing + + initdb: + build: . + container_name: aperisolve-initdb + command: python -m aperisolve.utils.init_db + volumes: + - ./aperisolve:/app/aperisolve + depends_on: + postgres: + condition: service_healthy + env_file: + - .env.example + restart: "no" redis: image: redis:7 diff --git a/compose.yml b/compose.yml index f4d1d1f..b388c51 100644 --- a/compose.yml +++ b/compose.yml @@ -1,9 +1,9 @@ services: web: - image: ghcr.io/zeecka/aperisolve:latest - #build: . + #image: ghcr.io/zeecka/aperisolve:latest + build: . container_name: aperisolve-web - restart: always + restart: no ports: - "5000:5000" volumes: @@ -11,6 +11,8 @@ services: depends_on: postgres: condition: service_healthy + initdb: + condition: service_completed_successfully redis: condition: service_started worker: @@ -25,20 +27,22 @@ services: start_period: 900s worker: - image: ghcr.io/zeecka/aperisolve:latest - #build: . + #image: ghcr.io/zeecka/aperisolve:latest + build: . container_name: aperisolve-worker - restart: always + restart: no command: rq worker default --url redis://redis:6379/0 volumes: - apresults:/app/aperisolve/results depends_on: - - redis - - postgres + postgres: + condition: service_healthy + initdb: + condition: service_completed_successfully + redis: + condition: service_started env_file: - .env - environment: - - CLEAR_AT_RESTART=0 # Always false for worker to prevent double clearing postgres: image: postgres:18 @@ -54,6 +58,19 @@ services: env_file: - .env + initdb: + build: . + container_name: aperisolve-initdb + command: python -m aperisolve.utils.init_db + volumes: + - apresults:/app/aperisolve/results + depends_on: + postgres: + condition: service_healthy + env_file: + - .env + restart: "no" + redis: image: redis:7 container_name: redis