From 5755bfb6508caa663bfa7443072a0624f88fe8da Mon Sep 17 00:00:00 2001
From: Azalea <22280294+hykilpikonna@users.noreply.github.com>
Date: Tue, 10 Mar 2026 19:17:49 -0400
Subject: [PATCH] [O] File name whitelist

---
 workflow.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/workflow.py b/workflow.py
index 9886b34..0b338ae 100644
--- a/workflow.py
+++ b/workflow.py
@@ -27,11 +27,15 @@ from utils_ai import (
 )
 
 
+VALID_MEDIA_EXTS = {".mkv", ".mp4", ".avi", ".ts", ".m2ts", ".webm", ".srt", ".ass", ".ssa", ".sub", ".idx"}
+
 def format_file_tree(file_tree: list) -> str:
     """Helper to convert the qB file tree output into simple relative paths for the LLM prompt"""
     lines = []
     for f in file_tree:
-        lines.append(f.get("name", ""))
+        name = f.get("name", "")
+        if Path(name).suffix.lower() in VALID_MEDIA_EXTS:
+            lines.append(name)
     return "\n".join(lines)
 
 def sanitize_filename(name: str) -> str: