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
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
+1 -1
View File
@@ -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"]
+3
View File
@@ -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"]
+2 -1
View File
@@ -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):
+203 -218
View File
@@ -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("/<hash_val>", 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("/<hash_val>", 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/<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
# 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/<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
# Try to find image in DB by hash
sub_img = Image.query.filter_by(hash=img_hash).first()
@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
# 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/<hash_val>/<tool>")
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/<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()
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/<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__":
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)
-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
"""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"))
+92 -95
View File
@@ -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
+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
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()
@@ -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()
+4 -3
View File
@@ -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
+21 -10
View File
@@ -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
+27 -10
View File
@@ -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