app.py

from __future__ import annotations
from pathlib import Path
from threading import Thread
import time
import os
import io
import json
import urllib.parse
import tempfile
import sys
import logging

from flask import Flask, render_template, request, jsonify, send_file, make_response
from werkzeug.utils import secure_filename
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address

sys.path.insert(0, str(Path(__file__).parent))
from image_forensics.core import SuiteConfig
from image_forensics.forensics import ForensicsEngine
from image_forensics.processing import ProcessingEngine
from image_forensics.analysis import AnalysisEngine

# ── Logging (production-friendly, no stack traces leaked to users) ──
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("image_forensics_pro")

app = Flask(__name__)

# ── Security / upload hardening ──
ALLOWED_EXTENSIONS = {"jpg", "jpeg", "png", "tiff", "tif", "webp"}
MAX_CONTENT_LENGTH = 10 * 1024 * 1024  # 10 MB max upload size
app.config["MAX_CONTENT_LENGTH"] = MAX_CONTENT_LENGTH

# Rate limiting — protects the server from being flooded with requests
limiter = Limiter(
    get_remote_address,
    app=app,
    default_limits=["60 per minute"],
    storage_uri="memory://",
)

OUTPUT = Path("./web_output").resolve()
OUTPUT.mkdir(exist_ok=True)

config     = SuiteConfig(
    output_dir=OUTPUT,
    log_to_file=False,
    upscale_weights_dir=Path(os.environ.get("WEIGHTS_DIR", "./weights")).resolve(),
)
forensics  = ForensicsEngine(config)
processing = ProcessingEngine(config)
analysis   = AnalysisEngine(config)


def allowed_file(filename: str) -> bool:
    return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS


def _safe_levels(raw_value, minimum: int = 2, maximum: int = 50) -> int:
    """Clamp the level_sweep 'levels' param to a safe range (matches the UI slider 4-50)."""
    try:
        val = int(raw_value)
    except (TypeError, ValueError):
        return 10
    return max(minimum, min(val, maximum))


def save_upload_safely(file) -> str:
    """Validate and save an uploaded file to a temp path. Raises ValueError on invalid input."""
    if not file or not file.filename:
        raise ValueError("No file uploaded")

    filename = secure_filename(file.filename)
    if not filename or not allowed_file(filename):
        raise ValueError("File type not allowed. Use JPG, PNG, TIFF, or WEBP.")

    suffix = "." + filename.rsplit(".", 1)[1].lower()
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
        file.save(temp_file.name)
        return temp_file.name


def start_cache_cleaner(target_dir: Path, max_age_hours: float = 1.0, check_interval_seconds: int = 600):
    def cleanup_loop():
        max_age_seconds = max_age_hours * 3600
        while True:
            try:
                if target_dir.exists():
                    now = time.time()
                    for file_path in target_dir.iterdir():
                        if file_path.is_file():
                            file_age = now - file_path.stat().st_mtime
                            if file_age > max_age_seconds:
                                try:
                                    file_path.unlink()
                                    logger.info(f"[CLEANUP] Automated forensic cache purged: {file_path.name}")
                                except Exception as delete_err:
                                    logger.warning(f"[CLEANUP] Failed to delete {file_path.name}: {delete_err}")
            except Exception as loop_err:
                logger.error(f"[CLEANUP] Error in cleanup lifecycle: {loop_err}")
            time.sleep(check_interval_seconds)

    cleanup_thread = Thread(target=cleanup_loop, daemon=True)
    cleanup_thread.start()


start_cache_cleaner(OUTPUT, max_age_hours=1.0, check_interval_seconds=600)


@app.route("/")
def index():
    return render_template("index.html")


@app.route("/audit", methods=["POST"])
@limiter.limit("10 per minute")
def audit():
    try:
        path = save_upload_safely(request.files.get("file"))
    except ValueError as e:
        return jsonify({"error": str(e)}), 400

    try:
        results = {
            "ela": forensics.error_level_analysis(path),
            "noise": forensics.noise_analysis(path),
            "clone": forensics.clone_detection(path),
            "metadata": forensics.extract_metadata(path),
            "frequency": analysis.frequency_analysis(path),
            "pca": analysis.pca_analysis(path),
            "gradient": analysis.luminance_gradient(path),
            "histogram": analysis.histogram_analysis(path),
        }

        findings, artifacts = [], {}
        for key, res in results.items():
            for f in res.findings:
                findings.append({
                    "module": key, "label": f.label, "severity": f.severity.value,
                    "description": f.description, "confidence": f.confidence
                })
            for ak, av in res.artifacts.items():
                artifacts[f"{key}_{ak}"] = "<In-Memory Binary Data>" if isinstance(av, bytes) else str(av)

        alerts = sum(1 for f in findings if f["severity"] in ("ALERT", "CRITICAL"))
        warns = sum(1 for f in findings if f["severity"] == "WARNING")

        if alerts >= 2:
            verdict, sev = "HIGH RISK — Multiple manipulation indicators detected", "high"
        elif alerts >= 1 or warns >= 3:
            verdict, sev = "MODERATE — Some anomalies detected, review recommended", "medium"
        else:
            verdict, sev = "CLEAN — No significant manipulation indicators found", "low"

        return jsonify({
            "verdict": verdict, "severity": sev, "alert_count": alerts, "warning_count": warns,
            "total_ms": round(sum(r.elapsed_ms for r in results.values()), 1),
            "image_hash": results["metadata"].image_hash, "findings": findings, "artifacts": artifacts
        })
    except Exception:
        logger.exception("Error in /audit")
        return jsonify({"error": "Failed to process image."}), 500
    finally:
        if os.path.exists(path):
            os.remove(path)


@app.route("/analyze", methods=["POST"])
@limiter.limit("20 per minute")
def analyze():
    module = request.form.get("module")
    try:
        path = save_upload_safely(request.files.get("file"))
    except ValueError as e:
        return jsonify({"error": str(e)}), 400

    try:
        dispatch = {
            "ela": lambda: forensics.error_level_analysis(path),
            "noise": lambda: forensics.noise_analysis(path),
            "clone": lambda: forensics.clone_detection(path),
            "metadata": lambda: forensics.extract_metadata(path),
            "frequency": lambda: analysis.frequency_analysis(path),
            "pca": lambda: analysis.pca_analysis(path),
            "gradient": lambda: analysis.luminance_gradient(path),
            "histogram": lambda: analysis.histogram_analysis(path),
            "level_sweep": lambda: analysis.level_sweep(path, levels=_safe_levels(request.form.get("levels", 10))),
        }

        if module not in dispatch:
            return jsonify({"error": f"unknown module: {module}"}), 400

        res = dispatch[module]()

        img_bytes = None
        if res.artifacts and len(res.artifacts) > 0:
            artifact_val = list(res.artifacts.values())[0]
            if isinstance(artifact_val, bytes):
                img_bytes = artifact_val
            elif isinstance(artifact_val, (str, Path)) and os.path.exists(artifact_val):
                with open(artifact_val, "rb") as f:
                    img_bytes = f.read()
                os.remove(artifact_val)

        if not img_bytes:
            with open(path, "rb") as f:
                img_bytes = f.read()

        response = make_response(send_file(io.BytesIO(img_bytes), mimetype="image/png"))
        findings_data = [{"label": f.label, "severity": f.severity.value, "description": f.description} for f in res.findings]
        response.headers["X-Success"] = json.dumps(res.success)
        response.headers["X-Findings"] = urllib.parse.quote(json.dumps(findings_data))
        response.headers["X-Metrics"] = urllib.parse.quote(json.dumps(res.metrics))
        return response
    except Exception:
        logger.exception("Error in /analyze")
        return jsonify({"error": "Failed to analyze image."}), 500
    finally:
        if os.path.exists(path):
            os.remove(path)


@app.route("/process", methods=["POST"])
@limiter.limit("10 per minute")
def process_image():
    ptype = request.form.get("type")
    try:
        path = save_upload_safely(request.files.get("file"))
    except ValueError as e:
        return jsonify({"error": str(e)}), 400

    try:
        if ptype == "denoise":
            res = processing.denoise(path, method=request.form.get("method", "nlm"), strength=float(request.form.get("strength", 10.0)))
        elif ptype == "deblur":
            res = processing.deblur(path, method=request.form.get("method", "wiener"))
        elif ptype == "watermark":
            res = processing.remove_watermark(path)
        elif ptype == "color":
            res = processing.enhance_color(path, method=request.form.get("method", "clahe"))
        elif ptype == "upscale":
            target_dim_val = int(request.form.get("target_dim", 2000))
            if target_dim_val <= 0 or target_dim_val > 4000:
                target_dim_val = 2000
            res = processing.ai_upscale(path, target_min_dim=target_dim_val)
        elif ptype == "contour":
            res = processing.extract_contours(path, mode=request.form.get("mode", "canny"))
        else:
            return jsonify({"error": f"unknown type: {ptype}"}), 400

        if hasattr(res, "error") and res.error:
            return jsonify({"error": res.error}), 400

        result_file_path = None
        for k, v in res.artifacts.items():
            if "mask" not in k:
                result_file_path = v
                break

        if result_file_path and os.path.exists(result_file_path):
            with open(result_file_path, "rb") as f:
                img_bytes = f.read()
            os.remove(result_file_path)

            response = make_response(send_file(io.BytesIO(img_bytes), mimetype="image/png"))
            response.headers["X-Success"] = json.dumps(res.success)
            response.headers["X-Metrics"] = urllib.parse.quote(json.dumps(res.metrics))
            return response
        else:
            return jsonify({"error": "Processing completed without producing an output image."}), 400
    except Exception:
        logger.exception("Error in /process")
        return jsonify({"error": "Failed to process image."}), 500
    finally:
        if os.path.exists(path):
            os.remove(path)


# ── Error handlers: never leak internal details to users ──
@app.errorhandler(413)
def file_too_large(e):
    return jsonify({"error": "File too large. Maximum size is 10MB."}), 413


@app.errorhandler(404)
def not_found(e):
    return jsonify({"error": "Not found"}), 404


@app.errorhandler(500)
def server_error(e):
    logger.exception("Unhandled server error")
    return jsonify({"error": "Internal server error"}), 500


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 5000))
    print("== ImageForensics Pro ==")
    print(f"  http://localhost:{port}")
    print("========================")
    app.run(debug=False, host="0.0.0.0", port=port)