Huge refactoring

This commit is contained in:
Zeecka
2026-01-08 05:34:24 +01:00
parent 4b6a944817
commit 5719ee7c3e
15 changed files with 424 additions and 423 deletions
+1 -2
View File
@@ -1,5 +1,5 @@
# Application configuration # Application configuration
MAX_CONTENT_LENGTH=1048576 MAX_CONTENT_LENGTH=3145728
MAX_PENDING_TIME=600 MAX_PENDING_TIME=600
MAX_STORE_TIME=259200 MAX_STORE_TIME=259200
CLEAR_AT_RESTART=1 CLEAR_AT_RESTART=1
@@ -7,7 +7,6 @@ SKIP_IHDR_FILL=0
# Flask configuration # Flask configuration
FLASK_DEBUG=0 FLASK_DEBUG=0
FLASK_ENV=development
# Redis configuration # Redis configuration
RQ_DASHBOARD_REDIS_URL=redis://redis:6379/0 RQ_DASHBOARD_REDIS_URL=redis://redis:6379/0
+1 -1
View File
@@ -49,4 +49,4 @@ COPY aperisolve/ /app/aperisolve/
RUN pip install --no-cache-dir -r /app/aperisolve/requirements.txt RUN pip install --no-cache-dir -r /app/aperisolve/requirements.txt
# Commande de lancement # 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"]
+3
View File
@@ -1,3 +1,6 @@
# flake8: noqa: E203,E501,W503 # flake8: noqa: E203,E501,W503
# pylint: disable=C0413,W0718,R0903,R0801 # pylint: disable=C0413,W0718,R0903,R0801
# mypy: disable-error-code=unused-awaitable # mypy: disable-error-code=unused-awaitable
from . import app, config, models, workers
__all__ = ["app", "config", "models", "workers"]
+2 -1
View File
@@ -13,8 +13,9 @@ from pathlib import Path
from shutil import rmtree from shutil import rmtree
from typing import Any, Optional, overload from typing import Any, Optional, overload
from ..config import MAX_PENDING_TIME
_thread_lock = threading.Lock() _thread_lock = threading.Lock()
MAX_PENDING_TIME = int(os.getenv("MAX_PENDING_TIME", "600")) # 10 minutes by default
class SubprocessAnalyzer(ABC): class SubprocessAnalyzer(ABC):
+203 -218
View File
@@ -16,270 +16,255 @@ from flask import Flask, Response, abort, jsonify, render_template, request, sen
from redis import Redis from redis import Redis
from rq import Queue from rq import Queue
from .sentry import initialize_sentry from .utils.sentry import initialize_sentry
initialize_sentry() initialize_sentry()
from .cleanup import cleanup_old_entries
from .config import IMAGE_EXTENSIONS, RESULT_FOLDER, WORKER_FILES from .config import IMAGE_EXTENSIONS, RESULT_FOLDER, WORKER_FILES
from .models import Image, Submission, db, init_db from .models import Image, Submission, cleanup_old_entries, 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)
@app.errorhandler(413) def create_app() -> Flask:
def too_large(_: Any) -> tuple[Response, int]: """Create flask application with routes."""
"""Error Handler for Max File Size.""" app = Flask(__name__)
sentry_sdk.set_context( app.json.sort_keys = False # type: ignore
"upload", app.config["SQLALCHEMY_DATABASE_URI"] = os.environ.get("DB_URI")
{ app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
"content_length": request.content_length, app.config["MAX_CONTENT_LENGTH"] = int(os.environ.get("MAX_CONTENT_LENGTH", 1024 * 1024))
"max_allowed": app.config["MAX_CONTENT_LENGTH"], app.config["REDIS_QUEUE"] = Queue(connection=Redis(host="redis", port=6379))
}, db.init_app(app)
)
sentry_sdk.capture_message("Upload failed: 413 Payload Too Large", level="warning")
return jsonify({"error": "Image size exceeded"}), 413
@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) @app.errorhandler(404)
def not_found(_: Any) -> str: def not_found(_: Any) -> str:
"""Error Handler for 404 not found.""" """Error Handler for 404 not found."""
return render_template("error.html", message="Resource not found", statuscode=404) 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("/") @app.route("/<hash_val>", methods=["GET"])
def index() -> str: def result_page(hash_val: str) -> str:
"""Render the main index page.""" """Render the result page for a given submission hash."""
return render_template("index.html") 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") @app.route("/upload", methods=["POST"])
def faq() -> str: def upload_image() -> tuple[Response, int]:
"""Render the FAQ page.""" """Handle image upload and initiate analysis."""
return render_template("faq.html")
cleanup_old_entries()
@app.route("/show") if "image" not in request.files:
def show() -> str: return jsonify({"error": "No image provided"}), 400
"""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)
image = request.files["image"]
password = request.form.get("password")
deep_analysis = request.form.get("deep") == "true"
@app.route("/<hash_val>", methods=["GET"]) if image.filename is None or image.filename == "":
def result_page(hash_val: str) -> str: return jsonify({"error": "No image provided"}), 400
"""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) 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"]) submission_data = (
def upload_image() -> tuple[Response, int]: image_data
"""Handle image upload and initiate analysis.""" + 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: # Only return if BOTH the file exists AND the DB entry exists
return jsonify({"error": "No image provided"}), 400 if submission_path.exists() and submission is not None:
return jsonify({"submission_hash": submission.hash}), 200
image = request.files["image"] # Self-Healing Logic
password = request.form.get("password") new_img_path = RESULT_FOLDER / img_hash / image_name
deep_analysis = request.form.get("deep") == "true"
if image.filename is None or image.filename == "": # This fixes the "Ghost File" issue (DB exists, but file is missing)
return jsonify({"error": "No image provided"}), 400 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: # Try to find image in DB by hash
return jsonify({"error": "Unsupported file type"}), 400 sub_img = Image.query.filter_by(hash=img_hash).first()
# Compute hashes and prepare submission # If DB entry is missing (even if file exists on disk), create it!
image_data = image.read() if sub_img is None:
img_hash = hashlib.md5(image_data).hexdigest() sub_img = Image(
image_name = img_hash + "." + image.filename.rsplit(".", maxsplit=1)[-1].lower() 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 = ( # Increment upload count for the image
image_data sub_img.upload_count += 1
+ image.filename.encode() db.session.commit() # pylint: disable=no-member
+ (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
# Check if hash is present on the DB # Create new Submission entry
submission: Submission = Submission.query.filter_by(hash=submission_hash).first() submission_path.mkdir(parents=True, exist_ok=True)
# Only return if BOTH the file exists AND the DB entry exists # Only create a new submission if it does not already exist
if submission_path.exists() and submission is not None: # 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 return jsonify({"submission_hash": submission.hash}), 200
# Self-Healing Logic @app.route("/status/<hash_val>", methods=["GET"])
new_img_path = RESULT_FOLDER / img_hash / image_name 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) @app.route("/infos/<hash_val>", methods=["GET"])
if not new_img_path.exists(): def get_infos(hash_val: str) -> Response:
new_img_path.parent.mkdir(parents=True, exist_ok=True) """Get the metadata and information of a submission."""
with open(new_img_path, "wb") as f: submission = Submission.query.get_or_404(hash_val)
f.write(image_data) 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 @app.route("/result/<hash_val>", methods=["GET"])
sub_img = Image.query.filter_by(hash=img_hash).first() 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! result_path = RESULT_FOLDER / str(image.hash) / str(submission.hash) / "results.json"
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
# Increment upload count for the image if not result_path.exists():
sub_img.upload_count += 1 return jsonify({"error": "Results not ready yet..."}), 425
db.session.commit() # pylint: disable=no-member
# Create new Submission entry with open(result_path, "r", encoding="utf-8") as f:
submission_path.mkdir(parents=True, exist_ok=True) results = json.load(f)
# Only create a new submission if it does not already exist return jsonify({"results": results}), 200
# 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 @app.route("/download/<hash_val>/<tool>")
submission.status = "pending" # type: ignore def download_output(hash_val: str, tool: str) -> Response:
db.session.commit() # pylint: disable=no-member """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/<hash_val>", 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/<hash_val>", 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/<hash_val>", 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/<hash_val>/<tool>")
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/<img_name>")
@app.route("/image/<hash_val>/<img_name>")
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() submission = Submission.query.filter_by(hash=hash_val).first_or_404()
image = Image.query.get_or_404(submission.image_hash) # type: ignore image = Image.query.get_or_404(submission.image_hash) # type: ignore
output_dir = RESULT_FOLDER / str(image.hash) / str(submission.hash) output_dir = RESULT_FOLDER / str(image.hash) / str(submission.hash)
output_file = output_dir / Path(img_name).name output_file = output_dir / f"{tool}.7z"
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): if tool not in WORKER_FILES or not output_file.exists():
return abort(404, description="Image not found or unsupported format") 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/<img_name>")
@app.route("/image/<hash_val>/<img_name>")
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__": if __name__ == "__main__":
debug_mode = os.environ.get("FLASK_DEBUG", "false").lower() == "true" my_app = create_app()
app.run(host="0.0.0.0", port=5000, debug=debug_mode) debug_mode = os.environ.get("FLASK_DEBUG", "0") == "1"
my_app.run(host="0.0.0.0", port=5000, debug=debug_mode)
-79
View File
@@ -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
+4
View File
@@ -3,9 +3,13 @@
# mypy: disable-error-code=unused-awaitable # mypy: disable-error-code=unused-awaitable
"""Aperi'Solve configuration variables.""" """Aperi'Solve configuration variables."""
from os import getenv
from pathlib import Path from pathlib import Path
IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".tiff"] IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".tiff"]
WORKER_FILES = ["binwalk", "foremost", "steghide", "zsteg", "image_resize", "openstego", "pcrt"] WORKER_FILES = ["binwalk", "foremost", "steghide", "zsteg", "image_resize", "openstego", "pcrt"]
RESULT_FOLDER = Path(__file__).parent.resolve() / "results" 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"))
+92 -95
View File
@@ -1,18 +1,15 @@
# flake8: noqa: E203,E501,W503 # 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 # mypy: disable-error-code=unused-awaitable
"""This module defines the database models for the Aperi'Solve application.""" """This module defines the database models for the Aperi'Solve application."""
import itertools import itertools
import shutil
import struct import struct
import sys import time
import zlib import zlib
from datetime import datetime, timezone 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 flask_sqlalchemy import SQLAlchemy
from sqlalchemy import ( from sqlalchemy import (
BigInteger, BigInteger,
@@ -25,6 +22,7 @@ from sqlalchemy import (
String, 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 from aperisolve.utils.utils import get_resolutions, get_valid_depth_color_pairs
db: SQLAlchemy = SQLAlchemy() db: SQLAlchemy = SQLAlchemy()
@@ -117,110 +115,109 @@ def fill_ihdr_db() -> None:
Computes CRC for each combination and stores parameters directly. 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 resolutions = get_resolutions()
if db.session.query(IHDR.iid).first() is not None: bit_color_pairs = list(get_valid_depth_color_pairs())
print("IHDR table already populated, skipping fill.") interlace_methods = [0, 1]
return
resolutions = get_resolutions() combinations = itertools.product(
bit_color_pairs = list(get_valid_depth_color_pairs()) resolutions,
interlace_methods = [0, 1] bit_color_pairs,
interlace_methods,
)
count = 0
combinations = itertools.product( with db.session.no_autoflush:
resolutions, for (width, height), (bit_depth, color_type), interlace in combinations:
bit_color_pairs, # Compute CRC directly
interlace_methods, crc = IHDR.compute_crc(width, height, bit_depth, color_type, interlace)
)
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)
# Check if entry already exists # Add new entry
exists = db.session.execute( db.session.add(
db.select(IHDR.iid).where( IHDR(
db.and_( crc=crc,
IHDR.crc == crc, width=width,
IHDR.width == width, height=height,
IHDR.height == height, bit_depth=bit_depth,
IHDR.bit_depth == bit_depth, color_type=color_type,
IHDR.color_type == color_type, interlace=interlace,
IHDR.interlace == interlace,
) )
) )
).first()
if exists: count += 1
continue if count % 5000 == 0:
print(f"Inserted {count} IHDR entries...")
db.session.commit()
print(f"Inserted {count} total IHDR entries.")
batch.append( except Exception as e:
{ print(f"Error filling IHDR table: {e}")
"crc": crc, db.session.rollback()
"width": width,
"height": height,
"bit_depth": bit_depth,
"color_type": color_type,
"interlace": interlace,
}
)
count += 1 def cleanup_old_entries() -> None:
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:
""" """
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, This function performs two main cleanup operations:
and populates the CRC lookup table for PNG IHDR validation.
Args: 1. Submission cleanup:
app (Flask): The Flask application instance used to establish the application context - Deletes submissions with "pending" or "running" status that have exceeded
for database operations. 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: Returns:
None 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 now = time.time()
# Worker: /usr/local/bin/rq worker ... for submission in Submission.query.all(): # type: ignore
# Web: /usr/local/bin/flask run ... or gunicorn/wsgi if (
is_worker = any("rq" in arg and "flask" not in arg for arg in sys.argv[:2]) 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(): # Delete "old" images
# Workers should never clear the database for img in Image.query.all(): # type: ignore
if not is_worker and getenv("CLEAR_AT_RESTART", "0") == "1": if img.last_submission_date.tzinfo is None:
print("Clearing database and file system at restart...") img_date = img.last_submission_date.replace(tzinfo=timezone.utc)
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.")
else: 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
+62
View File
@@ -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()
+2 -1
View File
@@ -9,7 +9,7 @@ import struct
import zlib import zlib
from typing import Any, Optional from typing import Any, Optional
from ..app import app from ..app import create_app
from ..models import IHDR from ..models import IHDR
from .utils import int2hex, str2hex from .utils import int2hex, str2hex
@@ -241,6 +241,7 @@ class PNG:
crc_int = struct.unpack("!I", crc)[0] crc_int = struct.unpack("!I", crc)[0]
# Database lookup # Database lookup
app = create_app()
with app.app_context(): with app.app_context():
matches = IHDR.query.filter_by(crc=crc_int).all() matches = IHDR.query.filter_by(crc=crc_int).all()
@@ -3,7 +3,6 @@
# mypy: disable-error-code=unused-awaitable # mypy: disable-error-code=unused-awaitable
"""Aperi'Solve WSGI runner.""" """Aperi'Solve WSGI runner."""
from .app import app as application from aperisolve.app import create_app
if __name__ == "__main__": application = create_app()
application.run()
+4 -3
View File
@@ -9,7 +9,7 @@ from typing import Any
import sentry_sdk import sentry_sdk
from .sentry import initialize_sentry from .utils.sentry import initialize_sentry
initialize_sentry() initialize_sentry()
@@ -24,9 +24,9 @@ from .analyzers.pngcheck import analyze_pngcheck
from .analyzers.steghide import analyze_steghide from .analyzers.steghide import analyze_steghide
from .analyzers.strings import analyze_strings from .analyzers.strings import analyze_strings
from .analyzers.zsteg import analyze_zsteg from .analyzers.zsteg import analyze_zsteg
from .app import app, db from .app import create_app
from .config import RESULT_FOLDER from .config import RESULT_FOLDER
from .models import Image, Submission from .models import Image, Submission, db
def analyze_image(submission_hash: str) -> None: 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 - Creates result directories and generates analysis output files
- Modifies the database session and commits changes - Modifies the database session and commits changes
""" """
app = create_app()
with app.app_context(): with app.app_context():
submission: Submission = Submission.query.get(submission_hash) # type: ignore submission: Submission = Submission.query.get(submission_hash) # type: ignore
image = Image.query.get_or_404(submission.image_hash) # type: ignore image = Image.query.get_or_404(submission.image_hash) # type: ignore
+21 -10
View File
@@ -12,16 +12,14 @@ services:
depends_on: depends_on:
postgres: postgres:
condition: service_healthy condition: service_healthy
initdb:
condition: service_completed_successfully
redis: redis:
condition: service_started condition: service_started
worker: worker:
condition: service_started condition: service_started
env_file: env_file:
- .env.example - .env.example
environment:
- FLASK_DEBUG=true
- FLASK_ENV=development
- CLEAR_AT_RESTART=0 # Don't clear on Flask debug auto-reload
worker: worker:
#image: ghcr.io/zeecka/aperisolve:latest #image: ghcr.io/zeecka/aperisolve:latest
@@ -32,14 +30,27 @@ services:
volumes: volumes:
- ./aperisolve:/app/aperisolve - ./aperisolve:/app/aperisolve
depends_on: depends_on:
- redis postgres:
- postgres condition: service_healthy
initdb:
condition: service_completed_successfully
redis:
condition: service_started
env_file: env_file:
- .env.example - .env.example
environment:
- FLASK_DEBUG=true initdb:
- FLASK_ENV=development build: .
- CLEAR_AT_RESTART=0 # Always false for worker to prevent double clearing 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: redis:
image: redis:7 image: redis:7
+27 -10
View File
@@ -1,9 +1,9 @@
services: services:
web: web:
image: ghcr.io/zeecka/aperisolve:latest #image: ghcr.io/zeecka/aperisolve:latest
#build: . build: .
container_name: aperisolve-web container_name: aperisolve-web
restart: always restart: no
ports: ports:
- "5000:5000" - "5000:5000"
volumes: volumes:
@@ -11,6 +11,8 @@ services:
depends_on: depends_on:
postgres: postgres:
condition: service_healthy condition: service_healthy
initdb:
condition: service_completed_successfully
redis: redis:
condition: service_started condition: service_started
worker: worker:
@@ -25,20 +27,22 @@ services:
start_period: 900s start_period: 900s
worker: worker:
image: ghcr.io/zeecka/aperisolve:latest #image: ghcr.io/zeecka/aperisolve:latest
#build: . build: .
container_name: aperisolve-worker container_name: aperisolve-worker
restart: always restart: no
command: rq worker default --url redis://redis:6379/0 command: rq worker default --url redis://redis:6379/0
volumes: volumes:
- apresults:/app/aperisolve/results - apresults:/app/aperisolve/results
depends_on: depends_on:
- redis postgres:
- postgres condition: service_healthy
initdb:
condition: service_completed_successfully
redis:
condition: service_started
env_file: env_file:
- .env - .env
environment:
- CLEAR_AT_RESTART=0 # Always false for worker to prevent double clearing
postgres: postgres:
image: postgres:18 image: postgres:18
@@ -54,6 +58,19 @@ services:
env_file: env_file:
- .env - .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: redis:
image: redis:7 image: redis:7
container_name: redis container_name: redis