Update app.py

app.py

@@ -5,10 +5,12 @@ import random
 import zipfile
 from PIL import Image
 from datetime import datetime
+import io
+import hashlib
 
 def procesar_video(video):
     try:
-        # Determinar la ruta del video y el nombre original
+        # 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"))
@@ -16,60 +18,103 @@ def procesar_video(video):
             original_name = os.path.basename(video)
             video_path = video
 
-        # Configurar directorio temporal
+        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)
-        
-        # Leer el video y extraer TODOS los frames
+
+        # 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
-            
-            # Guardar cada frame
             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 ValueError("No se pudieron extraer frames del video")
-        
-        # Seleccionar 4 frames aleatorios para el collage
+            raise gr.Error("No se pudieron extraer fotogramas del video")
+
+        # Crear collage de 4 fotogramas aleatorios
         selected_frames = random.sample(frame_paths, min(4, len(frame_paths)))
-        
-        # Crear collage
         images = [Image.open(img) for img in selected_frames]
         collage = Image.new('RGB', (images[0].width * 2, images[0].height * 2))
-        
         for i, img in enumerate(images):
             row = i // 2
             col = i % 2
-            x = col * images[0].width
-            y = row * images[0].height
-            collage.paste(img, (x, y))
-        
+            collage.paste(img, (col * images[0].width, row * images[0].height))
         collage_path = os.path.join(temp_dir, "collage.jpg")
         collage.save(collage_path)
-        
-        # Crear archivo ZIP con TODOS los frames usando el nombre original del video
-        base_name = os.path.splitext(original_name)[0]
-        zip_filename = f"{base_name}.zip"
-        zip_path = os.path.join(temp_dir, zip_filename)
-        with zipfile.ZipFile(zip_path, 'w') as zipf:
+
+        # ===============================
+        # 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, zip_path, temp_dir
-    
+        return collage_path, final_zip_path, temp_dir
+
     except Exception as e:
         raise gr.Error(f"Error al procesar el video: {str(e)}")
 
@@ -79,6 +124,9 @@ def limpiar_cache(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.**")
@@ -86,43 +134,57 @@ with gr.Blocks(title="Extracción de Fotogramas Forenses") as demo:
     
     with gr.Row():
         with gr.Column():
-            video_input = gr.Video(label="Subir Video", interactive=True)
+            video_input = gr.Video(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()
-    zip_path_state = gr.State()
-    
-    # Habilita el botón de procesar solo si se ha subido un video
-    def habilitar_procesar(video):
-        return gr.Button.update(interactive=True) if video else gr.Button.update(interactive=False)
+    temp_dir_state = gr.State(None)
+    zip_path_state = gr.State(None)
     
-    # Procesa el video, genera el collage y el ZIP y actualiza los estados
-    def procesar_y_mostrar(video):
-        if temp_dir_state.value:
-            limpiar_cache(temp_dir_state.value)
-        collage_path, zip_path, temp_dir = procesar_video(video)
-        return collage_path, zip_path, temp_dir, zip_path, gr.Button.update(interactive=True)
+    # 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):
+        if video is None:
+            return gr.Button.update(interactive=False), gr.Button.update(interactive=False)
+        if isinstance(video, dict):
+            original_name = video.get("name", "")
+        else:
+            original_name = os.path.basename(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)")
+        # Habilitar "Procesar Fotogramas" y deshabilitar "DESCARGAR FOTOGRAMAS"
+        return gr.Button.update(interactive=True), gr.Button.update(interactive=False)
     
     video_input.change(
-        fn=habilitar_procesar,
+        fn=on_video_change,
         inputs=video_input,
-        outputs=procesar_btn,
+        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):
+        # Desactivar el botón de procesamiento para evitar reprocesos
+        btn_proc_update = gr.Button.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.Button.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],
+        outputs=[gallery_output, download_file, temp_dir_state, zip_path_state, download_btn, procesar_btn],
     )
     
-    # Al pulsar "DESCARGAR FOTOGRAMAS", se envía la ruta del ZIP almacenada en zip_path_state
+    # Al pulsar "DESCARGAR FOTOGRAMAS": se envía la ruta del ZIP generado.
     download_btn.click(
         fn=lambda zip_path: zip_path,
         inputs=zip_path_state,