Only take env variables needed for sudo

This helps with how `sudo -E` doesn't work in sudo-rs.

Signed-off-by: Mario Limonciello <mario.limonciello@amd.com>

Author: superm1

src/amd_debug/common.py

@@ -382,7 +382,19 @@ def relaunch_sudo() -> None:
     """Relaunch the script with sudo if not already running as root"""
     if not is_root():
         logging.debug("Relaunching with sudo")
-        os.execvp("sudo", ["sudo", "-E"] + sys.argv)
+        env_vars = []
+        for var in ["DISPLAY", "WAYLAND_DISPLAY", "XAUTHORITY", "XDG_RUNTIME_DIR",
+                    "DBUS_SESSION_BUS_ADDRESS", "XDG_SESSION_TYPE"]:
+            value = os.environ.get(var)
+            if value:
+                env_vars.append(f"{var}={value}")
+
+        if env_vars:
+            sudo_cmd = ["sudo"] + [f"--preserve-env={var.split('=')[0]}" for var in env_vars] + sys.argv
+        else:
+            sudo_cmd = ["sudo"] + sys.argv
+
+        os.execvp("sudo", sudo_cmd)
 
 
 def running_ssh():

src/amd_debug/s2idle.py

@@ -60,14 +60,20 @@ def display_report_file(fname, fmt) -> None:
         return
     user = os.environ.get("SUDO_USER")
     if user:
-        # ensure that xdg tools will know how to display the file
-        # (user may need to call tool with sudo -E)
-        if os.environ.get("XDG_SESSION_TYPE"):
-            subprocess.call(["sudo", "-E", "-u", user, "xdg-open", fname])
+        env_vars = []
+        for var in ["DISPLAY", "WAYLAND_DISPLAY", "XAUTHORITY", "XDG_RUNTIME_DIR",
+                    "DBUS_SESSION_BUS_ADDRESS", "XDG_SESSION_TYPE"]:
+            value = os.environ.get(var)
+            if value:
+                env_vars.append(f"{var}={value}")
+
+        if env_vars:
+            cmd = ["sudo", "-u", user] + [f"env"] + env_vars + ["xdg-open", fname]
+            subprocess.call(cmd)
         else:
             print(
-                "To display report automatically in browser launch tool "
-                f"with '-E' argument (Example: sudo -E {sys.argv[0]})"
+                "Unable to detect graphical session environment. "
+                "Report saved to: " + fname
             )
 
 

src/test_s2idle.py

@@ -614,10 +614,13 @@ def test_display_report_file_html_root_with_user(
         display_report_file("report.html", "html")
         mock_is_root.assert_called_once()
         mock_env_get.assert_any_call("SUDO_USER")
-        mock_env_get.assert_any_call("XDG_SESSION_TYPE")
-        mock_subprocess_call.assert_called_once_with(
-            ["sudo", "-E", "-u", "testuser", "xdg-open", "report.html"]
-        )
+        call_args = mock_subprocess_call.call_args[0][0]
+        assert call_args[0] == "sudo"
+        assert call_args[1] == "-u"
+        assert call_args[2] == "testuser"
+        assert call_args[3] == "env"
+        assert "xdg-open" in call_args
+        assert "report.html" in call_args
 
     @patch("amd_debug.s2idle.is_root", return_value=True)
     @patch("os.environ.get", side_effect=lambda key: None)