app.py

import gradio as gr
import cv2
import os
import random
import zipfile
from PIL import Image
from datetime import datetime
import io
import hashlib

def procesar_video(video):
    try:
        # Validar y obtener el nombre y ruta del video
        if isinstance(video, dict):
            original_name = video.get("name", "video")
            video_path = video.get("file", video.get("data"))
        else:
            original_name = os.path.basename(video)
            video_path = video

        allowed_extensions = ('.mp4', '.avi', '.mov', '.mkv')
        if not original_name.lower().endswith(allowed_extensions):
            raise gr.Error("Solo se permiten archivos de video (mp4, avi, mov, mkv)")

        # Timestamp para la cadena de custodia
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

        # Crear directorio temporal
        temp_dir = f"temp_{datetime.now().strftime('%Y%m%d%H%M%S')}"
        os.makedirs(temp_dir, exist_ok=True)

        # Extraer TODOS los fotogramas del video
        cap = cv2.VideoCapture(video_path)
        frame_count = 0
        frame_paths = []
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(frame_rgb)
            img_path = os.path.join(temp_dir, f"frame_{frame_count:04d}.jpg")
            img.save(img_path)
            frame_paths.append(img_path)
            frame_count += 1
        cap.release()

        if frame_count == 0:
            raise gr.Error("No se pudieron extraer fotogramas del video")

        # Crear collage de 4 fotogramas aleatorios con separación
        selected_frames = random.sample(frame_paths, min(4, len(frame_paths)))
        images = [Image.open(img) for img in selected_frames]
        margin = 10
        collage_width = images[0].width * 2 + margin
        collage_height = images[0].height * 2 + margin
        collage = Image.new('RGB', (collage_width, collage_height), color=(255,255,255))
        positions = [
            (0, 0),
            (images[0].width + margin, 0),
            (0, images[0].height + margin),
            (images[0].width + margin, images[0].height + margin)
        ]
        for i, img in enumerate(images):
            collage.paste(img, positions[i])
        collage_path = os.path.join(temp_dir, "collage.jpg")
        collage.save(collage_path)

        # ===============================
        # Generación de la cadena de custodia
        # -------------------------------
        # 1. Crear ZIP provisional en memoria con SOLO los fotogramas
        provisional_zip = io.BytesIO()
        with zipfile.ZipFile(provisional_zip, mode="w") as zipf:
            for img in frame_paths:
                zipf.write(img, os.path.basename(img))
        provisional_zip.seek(0)
        provisional_data = provisional_zip.getvalue()
        md5_provisional = hashlib.md5(provisional_data).hexdigest()

        # 2. Crear un ZIP provisional en memoria que incluya un archivo TXT con marcador de posición para el MD5 final
        final_zip_mem = io.BytesIO()
        chain_placeholder = "MD5 final del ZIP: <pending>"
        chain_initial = (
            f"Nombre del archivo: {original_name}\n"
            f"Fecha de carga y extracción: {timestamp}\n"
            f"Número de fotogramas: {frame_count}\n"
            f"MD5 provisional (frames): {md5_provisional}\n"
            f"{chain_placeholder}\n"
        )
        with zipfile.ZipFile(final_zip_mem, mode="w") as zipf:
            for img in frame_paths:
                zipf.write(img, os.path.basename(img))
            zipf.writestr("chain_of_custody.txt", chain_initial)
        final_zip_mem.seek(0)
        final_zip_data = final_zip_mem.getvalue()
        md5_final = hashlib.md5(final_zip_data).hexdigest()

        # 3. Reconstruir el ZIP final en memoria con el TXT actualizado (incluyendo ambos MD5)
        final_chain = (
            f"Nombre del archivo: {original_name}\n"
            f"Fecha de carga y extracción: {timestamp}\n"
            f"Número de fotogramas: {frame_count}\n"
            f"MD5 provisional (frames): {md5_provisional}\n"
            f"MD5 final del ZIP: {md5_final}\n"
        )
        final_zip_mem2 = io.BytesIO()
        with zipfile.ZipFile(final_zip_mem2, mode="w") as zipf:
            for img in frame_paths:
                zipf.write(img, os.path.basename(img))
            zipf.writestr("chain_of_custody.txt", final_chain)
        final_zip_mem2.seek(0)
        final_zip_data2 = final_zip_mem2.getvalue()

        # Escribir el ZIP final a disco con el nombre basado en el video original
        base_name = os.path.splitext(original_name)[0]
        zip_filename = f"{base_name}.zip"
        final_zip_path = os.path.join(temp_dir, zip_filename)
        with open(final_zip_path, "wb") as f:
            f.write(final_zip_data2)
        # ===============================
        
        return collage_path, final_zip_path, temp_dir

    except Exception as e:
        raise gr.Error(f"Error al procesar el video: {str(e)}")

def limpiar_cache(temp_dir):
    if temp_dir and os.path.exists(temp_dir):
        for file in os.listdir(temp_dir):
            os.remove(os.path.join(temp_dir, file))
        os.rmdir(temp_dir)

# -------------------------------
# Interfaz Gradio
# -------------------------------
with gr.Blocks(title="Extracción de Fotogramas Forenses") as demo:
    gr.Markdown("# Herramienta de Extracción de Fotogramas Forenses")
    gr.Markdown("**Carga un video para extraer TODOS los fotogramas y generar un collage de muestra.**")
    gr.Markdown("Desarrollado por José R. Leonett para el Grupo de Peritos Forenses Digitales de Guatemala - [www.forensedigital.gt](https://www.forensedigital.gt)")
    
    with gr.Row():
        with gr.Column():
            # Desactivar opciones de webcam y edición (tijera) estableciendo source="upload"
            video_input = gr.Video(source="upload", label="Subir Video (solo archivos de video)", interactive=True)
            # Ambos botones inician desactivados
            procesar_btn = gr.Button("Procesar Fotogramas", interactive=False)
        with gr.Column():
            gallery_output = gr.Image(label="Collage de Muestra")
            download_btn = gr.Button("DESCARGAR FOTOGRAMAS", interactive=False)
            download_file = gr.File(label="Archivo ZIP generado", visible=False)
    
    # Estados para guardar el directorio temporal y la ruta del ZIP
    temp_dir_state = gr.State(None)
    zip_path_state = gr.State(None)
    
    # Al cargar un video: si es válido, se activa el botón "Procesar Fotogramas" y se desactiva "DESCARGAR FOTOGRAMAS"
    def on_video_change(video):
        # Verificar que se haya cargado algo y que contenga información
        if not video:
            return gr.update(interactive=False), gr.update(interactive=False)
        if isinstance(video, dict) and video.get("name"):
            original_name = video["name"]
        elif isinstance(video, str) and video:
            original_name = os.path.basename(video)
        else:
            return gr.update(interactive=False), gr.update(interactive=False)
        allowed_extensions = ('.mp4', '.avi', '.mov', '.mkv')
        if not original_name.lower().endswith(allowed_extensions):
            raise gr.Error("Solo se permiten archivos de video (mp4, avi, mov, mkv)")
        # Al tener un video válido, activa "Procesar Fotogramas" y desactiva "DESCARGAR FOTOGRAMAS"
        return gr.update(interactive=True), gr.update(interactive=False)
    
    video_input.change(
        fn=on_video_change,
        inputs=video_input,
        outputs=[procesar_btn, download_btn],
        queue=False
    )
    
    # Al pulsar "Procesar Fotogramas": se procesa el video, se genera el collage y el ZIP, se desactiva el botón de procesar
    # y se activa el botón de descarga.
    def procesar_y_mostrar(video):
        btn_proc_update = gr.update(interactive=False)
        if temp_dir_state.value:
            limpiar_cache(temp_dir_state.value)
        collage_path, zip_path, temp_dir = procesar_video(video)
        btn_download_update = gr.update(interactive=True)
        return collage_path, zip_path, temp_dir, zip_path, btn_download_update, btn_proc_update
    
    procesar_btn.click(
        fn=procesar_y_mostrar,
        inputs=video_input,
        outputs=[gallery_output, download_file, temp_dir_state, zip_path_state, download_btn, procesar_btn],
    )
    
    # Función para activar la descarga: actualiza el componente File con la ruta del ZIP y lo hace visible
    def trigger_download(zip_path):
        return gr.update(value=zip_path, visible=True)
    
    download_btn.click(
        fn=trigger_download,
        inputs=zip_path_state,
        outputs=download_file,
    )

if __name__ == "__main__":
    demo.launch()