forensics/runner.py

"""
Unified runner for classic digital forensics tools.
Wraps ELA, Spatial Gradient, Bit Plane, MinMax Deviation, and Wavelet Noise.
"""

import time
import numpy as np
from PIL import Image
from .ela import ELA
from .gradient import gradient_processing
from .bitplane import bit_plane_extractor
from .minmax import minmax_process
from .wavelet import noise_estimation
from xai_engine.viz import to_b64_png


def run_forensic_tool(
    image: Image.Image,
    tool_name: str,
    quality: int = 75,
    scale: int = 50,
    contrast: int = 20,
    gradient_intensity: int = 90,
    blue_mode: str = "Abs",
    channel: str = "Luminance",
    bit: int = 0,
    minmax_radius: int = 2,
) -> dict:
    """
    Execute selected forensic tool on input image.
    
    Returns dict: {
        "tool": str,
        "result_b64": str,
        "info": str,
        "compute_time_ms": float
    }
    """
    t0 = time.perf_counter()
    img_np = np.array(image.convert("RGB"))
    info_text = ""

    if tool_name == "ela":
        name = "Error Level Analysis (ELA)"
        res_img = ELA(img_np, quality=quality, scale=scale, contrast=contrast)
        info_text = f"JPEG Compression Quality: {quality}%, Scale: {scale}, Contrast: {contrast}%"

    elif tool_name == "gradient":
        name = "Spatial Gradient Analysis"
        res_img = gradient_processing(image, intensity=gradient_intensity, blue_mode=blue_mode)
        info_text = f"Gradient Intensity: {gradient_intensity}%, Blue Mode: {blue_mode}"

    elif tool_name == "bitplane":
        name = "Bit Plane Extractor"
        res_img = bit_plane_extractor(image, channel=channel, bit=bit)
        info_text = f"Channel: {channel}, Bit Plane: {bit} (0=LSB, 7=MSB)"

    elif tool_name == "minmax":
        name = "MinMax Local Deviation"
        res_img = minmax_process(img_np, channel=4, radius=minmax_radius)
        info_text = f"Local Block Radius: {minmax_radius} px"

    elif tool_name == "wavelet":
        name = "Wavelet Noise Estimation"
        noise_var = noise_estimation(img_np)
        info_text = f"Estimated High-Frequency Noise Variance: {noise_var:.4f}"
        res_img = image.convert("RGB")  # Return original with overlay text or noise metric

    else:
        raise ValueError(f"Unknown forensic tool: {tool_name}")

    elapsed_ms = (time.perf_counter() - t0) * 1000

    return {
        "tool": name,
        "result_b64": to_b64_png(res_img),
        "info": info_text,
        "compute_time_ms": round(elapsed_ms, 1),
    }