Huge refactoring
This commit is contained in:
+1
-2
@@ -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
@@ -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"]
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user