CVE-2024-32023: Insufficient Information (NVD-CWE-noinfo)
value: Kohya_ss is a GUI for Kohya's Stable Diffusion trainers. Kohya_ss is vulnerable to a path injection in the `common_gui.py` `find_and_replace` function. This vulnerability is fixed in 23.1.5.
CWE description: There is insufficient information about the issue to classify it; details are unkown or unspecified.
Published: 2024-04-16T15:15Z
Repository: https://github.com/bmaltais/kohya_ss
Language: Python
Files changed: kohya_gui/basic_caption_gui.py, kohya_gui/blip_caption_gui.py, kohya_gui/class_basic_training.py, kohya_gui/common_gui.py, kohya_gui/git_caption_gui.py, kohya_gui/group_images_gui.py, kohya_gui/manual_caption_gui.py, kohya_gui/wd14_caption_gui.py
Commit date: 2024-04-11 08:32:17 -0400
Commit message: Find replace (#2258)

* Improve security to limit finding and replacing contents of any files on the system, that the attacker knows the contents of by specifying the list of acceptable extensions for captions.

* Update caption extension fields to Dropdown with restricted choices for security reasons
Diff: diff --git a/kohya_gui/basic_caption_gui.py b/kohya_gui/basic_caption_gui.py
index 4866a3f5..c47442ab 100644
--- a/kohya_gui/basic_caption_gui.py
+++ b/kohya_gui/basic_caption_gui.py
@@ -188,13 +188,13 @@ def gradio_basic_caption_gui_tab(headless=False, default_images_dir=None):
                 get_folder_path,
                 outputs=images_dir,
                 show_progress=False,
             )
             # Textbox for caption file extension
-            caption_ext = gr.Textbox(
+            caption_ext = gr.Dropdown(
                 label="Caption file extension",
-                placeholder="Extension for caption file (e.g., .caption, .txt)",
+                choices=[".cap", ".caption", ".txt"],
                 value=".txt",
                 interactive=True,
             )
             # Checkbox to overwrite existing captions
             overwrite = gr.Checkbox(
diff --git a/kohya_gui/blip_caption_gui.py b/kohya_gui/blip_caption_gui.py
index cda0b288..fc4ca166 100644
--- a/kohya_gui/blip_caption_gui.py
+++ b/kohya_gui/blip_caption_gui.py
@@ -140,13 +140,13 @@ def gradio_blip_caption_gui_tab(headless=False, default_train_dir=None):
                 get_folder_path,
                 outputs=train_data_dir,
                 show_progress=False,
             )
         with gr.Row():
-            caption_file_ext = gr.Textbox(
+            caption_file_ext = gr.Dropdown(
                 label="Caption file extension",
-                placeholder="Extension for caption file (e.g., .caption, .txt)",
+                choices=[".cap", ".caption", ".txt"],
                 value=".txt",
                 interactive=True,
             )
 
             prefix = gr.Textbox(
diff --git a/kohya_gui/class_basic_training.py b/kohya_gui/class_basic_training.py
index 8b1a040f..57c8e735 100644
--- a/kohya_gui/class_basic_training.py
+++ b/kohya_gui/class_basic_training.py
@@ -96,14 +96,15 @@ class BasicTraining:
             # Initialize the save every N epochs input
             self.save_every_n_epochs = gr.Number(
                 label="Save every N epochs", value=self.config.get("basic.save_every_n_epochs", 1), precision=0
             )
             # Initialize the caption extension input
-            self.caption_extension = gr.Textbox(
-                label="Caption Extension",
-                placeholder="(Optional) default: .caption",
-                value=self.config.get("basic.caption_extension", ""),
+            self.caption_extension = gr.Dropdown(
+                label="Caption file extension",
+                choices=[".cap", ".caption", ".txt"],
+                value=".txt",
+                interactive=True,
             )
 
     def init_precision_and_resources_controls(self) -> None:
         """
         Initializes the precision and resources controls for the model.
diff --git a/kohya_gui/common_gui.py b/kohya_gui/common_gui.py
index fc5e4944..97addc15 100644
--- a/kohya_gui/common_gui.py
+++ b/kohya_gui/common_gui.py
@@ -713,11 +713,10 @@ def has_ext_files(folder_path: str, file_extension: str) -> bool:
             return True
 
     # If no file with the specified extension is found, return False
     return False
 
-
 def find_replace(
     folder_path: str = "",
     caption_file_ext: str = ".caption",
     search_text: str = "",
     replace_text: str = "",
@@ -745,25 +744,45 @@ def find_replace(
         log.warning(
             "No files with extension {caption_file_ext} were found in {folder_path}..."
         )
         # Exit the function early
         return
+    
+    # Check if the caption file extension is one of the supported extensions
+    if caption_file_ext not in [".caption", ".txt", ".txt2", ".cap"]:
+        log.error(
+            f"Unsupported file extension {caption_file_ext} for caption files. Please use .caption, .txt, .txt2, or .cap."
+        )
+        # Exit the function early
+        return
+    
+    # Check if the folder path exists
+    if not os.path.exists(folder_path):
+        log.error(f"The provided path '{folder_path}' is not a valid folder.")
+        return
 
     # List all caption files in the folder
-    caption_files = [f for f in os.listdir(folder_path) if f.endswith(caption_file_ext)]
+    try:
+        caption_files = [f for f in os.listdir(folder_path) if f.endswith(caption_file_ext)]
+    except Exception as e:
+        log.error(f"Error accessing folder {folder_path}: {e}")
+        return
 
     # Iterate over the list of caption files
     for caption_file in caption_files:
         # Construct the full path for each caption file
         file_path = os.path.join(folder_path, caption_file)
         # Read and replace text
-        with open(file_path, "r", errors="ignore") as f:
-            content = f.read().replace(search_text, replace_text)
-
-        # Write the updated content back to the file
-        with open(file_path, "w") as f:
-            f.write(content)
+        try:
+            with open(file_path, "r", errors="ignore") as f:
+                content = f.read().replace(search_text, replace_text)
+
+            # Write the updated content back to the file
+            with open(file_path, "w") as f:
+                f.write(content)
+        except Exception as e:
+            log.error(f"Error processing file {file_path}: {e}")
 
 
 def color_aug_changed(color_aug):
     """
     Handles the change in color augmentation checkbox.
diff --git a/kohya_gui/git_caption_gui.py b/kohya_gui/git_caption_gui.py
index f0b07def..e2538b25 100644
--- a/kohya_gui/git_caption_gui.py
+++ b/kohya_gui/git_caption_gui.py
@@ -112,13 +112,13 @@ def gradio_git_caption_gui_tab(headless=False, default_train_dir=None):
                 get_folder_path,
                 outputs=train_data_dir,
                 show_progress=False,
             )
         with gr.Row():
-            caption_ext = gr.Textbox(
+            caption_ext = gr.Dropdown(
                 label="Caption file extension",
-                placeholder="Extension for caption file (e.g., .caption, .txt)",
+                choices=[".cap", ".caption", ".txt"],
                 value=".txt",
                 interactive=True,
             )
 
             prefix = gr.Textbox(
diff --git a/kohya_gui/group_images_gui.py b/kohya_gui/group_images_gui.py
index 914aad25..5219fb83 100644
--- a/kohya_gui/group_images_gui.py
+++ b/kohya_gui/group_images_gui.py
@@ -169,13 +169,13 @@ def gradio_group_images_gui_tab(headless=False):
                 label="Generate Captions",
                 value=False,
                 info="Generate caption files for the grouped images based on their folder name",
             )
 
-            caption_ext = gr.Textbox(
-                label="Caption Extension",
-                placeholder="Caption file extension (e.g., .txt)",
+            caption_ext = gr.Dropdown(
+                label="Caption file extension",
+                choices=[".cap", ".caption", ".txt"],
                 value=".txt",
                 interactive=True,
             )
 
         group_images_button = gr.Button("Group images")
diff --git a/kohya_gui/manual_caption_gui.py b/kohya_gui/manual_caption_gui.py
index bc90d215..cca885a9 100644
--- a/kohya_gui/manual_caption_gui.py
+++ b/kohya_gui/manual_caption_gui.py
@@ -296,13 +296,13 @@ def gradio_manual_caption_gui_tab(headless=False, default_images_dir=None):
                 get_folder_path,
                 outputs=images_dir,
                 show_progress=False,
             )
             load_images_button = gr.Button("Load", elem_id="open_folder")
-            caption_ext = gr.Textbox(
+            caption_ext = gr.Dropdown(
                 label="Caption file extension",
-                placeholder="Extension for caption file (e.g., .caption, .txt)",
+                choices=[".cap", ".caption", ".txt"],
                 value=".txt",
                 interactive=True,
             )
             auto_save = gr.Checkbox(
                 label="Autosave", info="Options", value=True, interactive=True
diff --git a/kohya_gui/wd14_caption_gui.py b/kohya_gui/wd14_caption_gui.py
index 9768d315..c5b8269d 100644
--- a/kohya_gui/wd14_caption_gui.py
+++ b/kohya_gui/wd14_caption_gui.py
@@ -186,14 +186,14 @@ def gradio_wd14_caption_gui_tab(
                 info="Useful to force model re download when switching to onnx",
             )
 
         with gr.Row():
 
-            caption_extension = gr.Textbox(
+            caption_extension = gr.Dropdown(
                 label="Caption file extension",
-                placeholder="Extension for caption file (e.g., .caption, .txt)",
-                value=config.get("wd14_caption.caption_extension", ".txt"),
+                choices=[".cap", ".caption", ".txt"],
+                value=".txt",
                 interactive=True,
             )
 
             caption_separator = gr.Textbox(
                 label="Caption Separator",

Vulnerable code:             caption_ext = gr.Textbox(
                placeholder="Extension for caption file (e.g., .caption, .txt)",
            caption_file_ext = gr.Textbox(
                placeholder="Extension for caption file (e.g., .caption, .txt)",
            self.caption_extension = gr.Textbox(
                label="Caption Extension",
                placeholder="(Optional) default: .caption",
                value=self.config.get("basic.caption_extension", ""),

    caption_files = [f for f in os.listdir(folder_path) if f.endswith(caption_file_ext)]
        with open(file_path, "r", errors="ignore") as f:
            content = f.read().replace(search_text, replace_text)

        # Write the updated content back to the file
        with open(file_path, "w") as f:
            f.write(content)
            caption_ext = gr.Textbox(
                placeholder="Extension for caption file (e.g., .caption, .txt)",
            caption_ext = gr.Textbox(
                label="Caption Extension",
                placeholder="Caption file extension (e.g., .txt)",
            caption_ext = gr.Textbox(
                placeholder="Extension for caption file (e.g., .caption, .txt)",
            caption_extension = gr.Textbox(
                placeholder="Extension for caption file (e.g., .caption, .txt)",
                value=config.get("wd14_caption.caption_extension", ".txt"),
Fixed code:             caption_ext = gr.Dropdown(
                choices=[".cap", ".caption", ".txt"],
            caption_file_ext = gr.Dropdown(
                choices=[".cap", ".caption", ".txt"],
            self.caption_extension = gr.Dropdown(
                label="Caption file extension",
                choices=[".cap", ".caption", ".txt"],
                value=".txt",
                interactive=True,
    
    # Check if the caption file extension is one of the supported extensions
    if caption_file_ext not in [".caption", ".txt", ".txt2", ".cap"]:
        log.error(
            f"Unsupported file extension {caption_file_ext} for caption files. Please use .caption, .txt, .txt2, or .cap."
        )
        # Exit the function early
        return
    
    # Check if the folder path exists
    if not os.path.exists(folder_path):
        log.error(f"The provided path '{folder_path}' is not a valid folder.")
        return
    try:
        caption_files = [f for f in os.listdir(folder_path) if f.endswith(caption_file_ext)]
    except Exception as e:
        log.error(f"Error accessing folder {folder_path}: {e}")
        return
        try:
            with open(file_path, "r", errors="ignore") as f:
                content = f.read().replace(search_text, replace_text)

            # Write the updated content back to the file
            with open(file_path, "w") as f:
                f.write(content)
        except Exception as e:
            log.error(f"Error processing file {file_path}: {e}")
            caption_ext = gr.Dropdown(
                choices=[".cap", ".caption", ".txt"],
            caption_ext = gr.Dropdown(
                label="Caption file extension",
                choices=[".cap", ".caption", ".txt"],
            caption_ext = gr.Dropdown(
                choices=[".cap", ".caption", ".txt"],
            caption_extension = gr.Dropdown(
                choices=[".cap", ".caption", ".txt"],
                value=".txt",