diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f7188a9..065d258 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -34,7 +34,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - pip install psutil pytest + pip install psutil platformdirs pytest - name: Install package (no-deps to avoid GUI reqs) run: pip install -e . --no-deps diff --git a/src/lufus/gui/gui.py b/src/lufus/gui/gui.py index 4f384e9..9b96439 100644 --- a/src/lufus/gui/gui.py +++ b/src/lufus/gui/gui.py @@ -1099,10 +1099,16 @@ def dropEvent(self, event): def browse_file(self): # open file dialog to select image :D + from lufus.user_paths import get_best_starting_dir + # ^ Uses the XDG_DOWNALOD_DIR that was detected and shoved into a variable + + starting_dir = get_best_starting_dir() + self.log_message(f"Opening file browser at: {starting_dir}") + file_name, _ = QFileDialog.getOpenFileName( self, self._T.get("dlg_select_image_title", "Select Image"), - "", + starting_dir, self._T.get( "dlg_select_image_filter", "Disk Images (*.iso *.dmg *.img *.bin *.raw);;All Files (*)", diff --git a/src/lufus/gui/start_gui.py b/src/lufus/gui/start_gui.py index 0d2bdeb..2a9542c 100644 --- a/src/lufus/gui/start_gui.py +++ b/src/lufus/gui/start_gui.py @@ -51,6 +51,16 @@ def _show_root_warning() -> None: def launch_gui_with_usb_data() -> None: elevation_attempted = False if os.geteuid() != 0: + # Capture user context (XDG_DOWNLOAD_DIR path) before pkexec elevation + from lufus.user_paths import get_best_starting_dir, ENV_DOWNLOAD_DIR + + try: + detected_path = get_best_starting_dir() + os.environ[ENV_DOWNLOAD_DIR] = detected_path + log.info("Captured user starting directory: %s", detected_path) + except Exception as e: + log.warning("Could not capture user starting directory: %s", e) + _load_initial_theme() elevation_attempted = True elevate_privileges() diff --git a/src/lufus/user_paths.py b/src/lufus/user_paths.py new file mode 100644 index 0000000..785cd83 --- /dev/null +++ b/src/lufus/user_paths.py @@ -0,0 +1,51 @@ +""" +Utility for detecting user's XDG_DOWNLOAD_DIR across root elevation +Falls back to ~/ if no XDG_DOWNLOAD_DIR is found +Made it so it runs as fast as possible before exiting to free like 2-3 KB RAM if you are using an OptiPlex GX270 or smth +""" + +import os +from pathlib import Path +from platformdirs import user_downloads_dir, user_documents_dir +from lufus.lufus_logging import get_logger + +log = get_logger(__name__) + +# Environment variable used to tunnel the path through the root/pkexec barrier +ENV_DOWNLOAD_DIR = "LUFUS_DOWNLOAD_DIR" + + +def get_best_starting_dir() -> str: + """ + Identify the best default dir for the file browser. + + This function works in two steps: + 1: Pre-elevation: Detects the real user's XDG_DOWNLOAD_DIR. If nothing is fonud, it falls back to ~/. + 2: Post-elevation: Gets the path from the tunneled environment variable. + + Returns: + str: An absolute path to a valid directory. + """ + # 1. High Priority: Check if we've already tunneled the path via environment + tunneled = os.environ.get(ENV_DOWNLOAD_DIR) + if tunneled: + if os.path.isdir(tunneled): + log.debug("Using tunneled user path: %s", tunneled) + return tunneled + log.warning("Tunneled path %s is no longer a valid directory.", tunneled) + + # 2. Medium Priority: Use platformdirs to find XDG standard paths + # This works best when called before elevation (geteuid != 0) + try: + # Standard XDG Downloads + dl_path = user_downloads_dir() + if dl_path and os.path.isdir(dl_path): + log.debug("Detected XDG Downloads directory: %s", dl_path) + return str(dl_path) + except Exception as e: + log.error("Failed to resolve XDG Downloads directory: %s", e) + + # 3. Low Priority: Home directory fallback + home = str(Path.home()) + log.debug("Falling back to home directory: %s", home) + return home diff --git a/src/lufus/utils.py b/src/lufus/utils.py index 4065ff4..2226948 100644 --- a/src/lufus/utils.py +++ b/src/lufus/utils.py @@ -27,6 +27,7 @@ def elevate_privileges() -> None: ) # Preserve DISPLAY and XAUTHORITY for GUI apps under pkexec/sudo + # Now also takes the detected XDG_DOWNLOAD_DIR of /src/lufus/user_paths.py to put it into LUFUS_DOWNLOAD_DIR env_vars = [ "DISPLAY", "XAUTHORITY", @@ -34,6 +35,7 @@ def elevate_privileges() -> None: "WAYLAND_DISPLAY", "PYTHONPATH", "LUFUS_THEME", + "LUFUS_DOWNLOAD_DIR", ] cmd = ["pkexec", "env"] @@ -47,7 +49,7 @@ def elevate_privileges() -> None: subprocess.run(cmd, check=True) sys.exit(0) except subprocess.CalledProcessError: - # User likely cancelled or pkexec failed + # User likely cancelled or pkexec failed/isn't installed pass except Exception as e: print(f"Elevation failed: {e}") diff --git a/tests/test_user_paths.py b/tests/test_user_paths.py new file mode 100644 index 0000000..99abf53 --- /dev/null +++ b/tests/test_user_paths.py @@ -0,0 +1,59 @@ +import os +import unittest +from unittest.mock import patch, MagicMock +from lufus.user_paths import get_best_starting_dir, ENV_DOWNLOAD_DIR + + +class TestUserPaths(unittest.TestCase): + def setUp(self): + # Clear the var before each test + if ENV_DOWNLOAD_DIR in os.environ: + del os.environ[ENV_DOWNLOAD_DIR] + + def test_get_best_starting_dir_with_env(self): + """Test that it prioritizes the environment variable if it points to a valid dir.""" + test_path = "/tmp/lufus_test_dir" + os.makedirs(test_path, exist_ok=True) + try: + os.environ[ENV_DOWNLOAD_DIR] = test_path + self.assertEqual(get_best_starting_dir(), test_path) + finally: + if os.path.exists(test_path): + os.rmdir(test_path) + + @patch("lufus.user_paths.user_downloads_dir") + @patch("os.path.isdir") + def test_get_best_starting_dir_with_xdg(self, mock_isdir, mock_downloads): + """Test that it uses XDG_DOWNLOAD_DIR if env is not set""" + mock_downloads.return_value = "/home/user/Downloads" + mock_isdir.side_effect = lambda p: p == "/home/user/Downloads" + + self.assertEqual(get_best_starting_dir(), "/home/user/Downloads") + + @patch("lufus.user_paths.user_downloads_dir") + @patch("os.path.isdir") + def test_get_best_starting_dir_with_custom_name(self, mock_isdir, mock_downloads): + """Test that it handles edge cases/translated names (like ~/Potato or ~/Téléchargements) (~/Potato is a great Downloads folder name btw)""" + custom_path = "/home/user/Téléchargements" + mock_downloads.return_value = custom_path + mock_isdir.side_effect = lambda p: p == custom_path + + self.assertEqual(get_best_starting_dir(), custom_path) + + @patch("lufus.user_paths.user_downloads_dir") + @patch("pathlib.Path.home") + @patch("os.path.isdir") + def test_get_best_starting_dir_fallback_to_home(self, mock_isdir, mock_home, mock_downloads): + """Test fallback to Home if Downloads is not found.""" + mock_downloads.return_value = "/home/user/Downloads" + mock_home.return_value = MagicMock() + mock_home.return_value.__str__.return_value = "/home/user" + + # Downloads does not exist + mock_isdir.return_value = False + + self.assertEqual(get_best_starting_dir(), "/home/user") + + +if __name__ == "__main__": + unittest.main()