diff --git a/kohya_gui.py b/kohya_gui.py index e8d64a8..0954918 100644 --- a/kohya_gui.py +++ b/kohya_gui.py @@ -102,6 +102,7 @@ def UI(**kwargs): "share": False if kwargs.get("do_not_share", False) else kwargs.get("share", False), "root_path": kwargs.get("root_path", None), "debug": kwargs.get("debug", False), + "allowed_paths": config.allowed_paths, } # This line filters out any key-value pairs from `launch_params` where the value is `None`, ensuring only valid parameters are passed to the `launch` function. diff --git a/kohya_gui/class_gui_config.py b/kohya_gui/class_gui_config.py index 33064a9..92c20da 100644 --- a/kohya_gui/class_gui_config.py +++ b/kohya_gui/class_gui_config.py @@ -16,6 +16,7 @@ class KohyaSSGUIConfig: Initialize the KohyaSSGUIConfig class. """ self.config = self.load_config(config_file_path=config_file_path) + self.allowed_paths = self.get_allowed_paths() def load_config(self, config_file_path: str = "./config.toml") -> dict: """ @@ -81,6 +82,12 @@ class KohyaSSGUIConfig: log.debug(f"Returned {data}") return data + def get_allowed_paths(self) -> list: + """ + Retrieves the list of allowed paths from the [server] section of the config file. + """ + return self.get("server.allowed_paths", []) + def is_config_loaded(self) -> bool: """ Checks if the configuration was loaded from a file. diff --git a/test/test_allowed_paths.py b/test/test_allowed_paths.py new file mode 100644 index 0000000..a99b8a7 --- /dev/null +++ b/test/test_allowed_paths.py @@ -0,0 +1,59 @@ +import os +import tempfile +import unittest +import subprocess +from pathlib import Path + +class TestAllowedPaths(unittest.TestCase): + def setUp(self): + print("Setting up test...") + self.temp_dir = tempfile.TemporaryDirectory() + self.temp_dir_path = Path(self.temp_dir.name) + self.dummy_file = self.temp_dir_path / "dummy.txt" + with open(self.dummy_file, "w") as f: + f.write("dummy content") + + self.config_file = self.temp_dir_path / "config.toml" + with open(self.config_file, "w") as f: + f.write(f''' +[server] +allowed_paths = ["{self.temp_dir.name}"] +''') + print("Setup complete.") + + def tearDown(self): + print("Tearing down test...") + self.temp_dir.cleanup() + print("Teardown complete.") + + def test_allowed_paths(self): + print("Running test_allowed_paths...") + # Run the gui with the new config and check if it can access the dummy file + process = subprocess.Popen( + [ + "python", + "kohya_gui.py", + "--config", + str(self.config_file), + "--headless", + ], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + print("Process started.") + # Give the server some time to start + try: + stdout, stderr = process.communicate(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + stdout, stderr = process.communicate() + + print(f"Stdout: {stdout.decode()}") + print(f"Stderr: {stderr.decode()}") + # Check if there are any errors in the stderr + self.assertNotIn("InvalidPathError", stderr.decode()) + print("Test complete.") + +if __name__ == "__main__": + unittest.main()