fix(powershell): update command arguments to include execution policy (#1041)

* fix(powershell): update command arguments to include execution policy

* fix(powershell): format command arguments for better readability

* fix(powershell): refactor command arguments for Windows execution policy and common options

* fix(powershell): improve execution policy handling and refactor argument passing

* refactor(powershell): streamline command execution and improve profile handling

* fix(powershell): improve code formatting for better readability and maintainability

* fix(powershell): enhance argument validation for improved error handling

* refactor(powershell): simplify Powershell struct methods and enhance command preparation

* refactor(powershell): streamline command building and enhance argument handling

* refactor(powershell): change add_common_args to use instance method for better encapsulation

* refactor(powershell): enhance command building by introducing build_command method and improving argument handling

* refactor(powershell): update build_command return type to match executor.execute() output

* refactor(powershell): refactor command building by introducing build_command_internal and adding profile getter

* refactor(powershell): improve profile retrieval and command execution logic

* refactor(powershell): add support for Windows update and Microsoft Store commands

* refactor(powershell): change args method to arg for accepting all Windows updates

* refactor(powershell): change arg method to args for accepting all Windows updates

* refactor(powershell): change args method to arg for accepting all Windows updates

* refactor(powershell): streamline command construction for Windows updates

* refactor(powershell): update args method for has_module function

* refactor(powershell): improve execution policy handling and command construction for Windows updates

* refactor(powershell): update Microsoft Store update command execution and streamline output handling

* refactor(powershell): clean up whitespace and improve code readability in execution policy checks

---------

Co-authored-by: nistee <lo9s4b7qp@mozmail.com>

src/steps/powershell.rs

@@ -1,5 +1,3 @@
-#[cfg(windows)]
-use std::path::Path;
 use std::path::PathBuf;
 use std::process::Command;
 
@@ -18,22 +16,9 @@ pub struct Powershell {
 }
 
 impl Powershell {
-    /// Returns a powershell instance.
-    ///
-    /// If the powershell binary is not found, or the current terminal is dumb
-    /// then the instance of this struct will skip all the powershell steps.
     pub fn new() -> Self {
         let path = which("pwsh").or_else(|| which("powershell")).filter(|_| !is_dumb());
-
+        let profile = path.as_ref().and_then(Self::get_profile);
-        let profile = path.as_ref().and_then(|path| {
-            Command::new(path)
-                .args(["-NoProfile", "-Command", "Split-Path $profile"])
-                .output_checked_utf8()
-                .map(|output| PathBuf::from(output.stdout.trim()))
-                .and_then(super::super::utils::PathExt::require)
-                .ok()
-        });
-
         Powershell { path, profile }
     }
 
@@ -45,117 +30,178 @@ impl Powershell {
         }
     }
 
+    pub fn profile(&self) -> Option<&PathBuf> {
+        self.profile.as_ref()
+    }
+
+    fn get_profile(path: &PathBuf) -> Option<PathBuf> {
+        Self::execute_with_command(path, &["-NoProfile", "-Command", "Split-Path $PROFILE"], |stdout| {
+            Ok(stdout)
+        })
+        .ok() // Convert the Result<String> to Option<String>
+        .and_then(|s| super::super::utils::PathExt::require(PathBuf::from(s)).ok())
+    }
+
+    fn execute_with_command<F>(path: &PathBuf, args: &[&str], f: F) -> Result<String>
+    where
+        F: FnOnce(String) -> Result<String>,
+    {
+        let output = Command::new(path).args(args).output_checked_utf8()?;
+        let stdout = output.stdout.trim().to_string();
+        f(stdout)
+    }
+
+    /// Builds a command with common arguments and optional sudo support.
+    fn build_command_internal<'a>(
+        &self,
+        ctx: &'a ExecutionContext,
+        additional_args: &[&str],
+    ) -> Result<impl CommandExt + 'a> {
+        let powershell = require_option(self.path.as_ref(), t!("Powershell is not installed").to_string())?;
+        let executor = &mut ctx.run_type();
+        let mut command = if let Some(sudo) = ctx.sudo() {
+            let mut cmd = executor.execute(sudo);
+            cmd.arg(powershell);
+            cmd
+        } else {
+            executor.execute(powershell)
+        };
+
+        #[cfg(windows)]
+        {
+            // Check execution policy and return early if it's not set correctly
+            self.execution_policy_args_if_needed()?;
+        }
+
+        command.args(Self::common_args()).args(additional_args);
+        Ok(command)
+    }
+
+    pub fn update_modules(&self, ctx: &ExecutionContext) -> Result<()> {
+        print_separator(t!("Powershell Modules Update"));
+        let mut cmd_args = vec!["Update-Module"];
+
+        if ctx.config().verbose() {
+            cmd_args.push("-Verbose");
+        }
+        if ctx.config().yes(Step::Powershell) {
+            cmd_args.push("-Force");
+        }
+        println!("{}", t!("Updating modules..."));
+        self.build_command_internal(ctx, &cmd_args)?.status_checked()
+    }
+
+    fn common_args() -> &'static [&'static str] {
+        &["-NoProfile"]
+    }
+
     #[cfg(windows)]
-    pub fn has_module(powershell: &Path, command: &str) -> bool {
+    pub fn execution_policy_args_if_needed(&self) -> Result<()> {
+        if !self.is_execution_policy_set("RemoteSigned") {
+            Err(color_eyre::eyre::eyre!(
+                "PowerShell execution policy is too restrictive. \
+                Please run 'Set-ExecutionPolicy RemoteSigned -Scope CurrentUser' in PowerShell \
+                (or use Unrestricted/Bypass if you're sure about the security implications)"
+            ))
+        } else {
+            Ok(())
+        }
+    }
+
+    #[cfg(windows)]
+    fn is_execution_policy_set(&self, policy: &str) -> bool {
+        if let Some(powershell) = &self.path {
+            // These policies are ordered from most restrictive to least restrictive
+            let valid_policies = ["Restricted", "AllSigned", "RemoteSigned", "Unrestricted", "Bypass"];
+
+            // Find the index of our target policy
+            let target_idx = valid_policies.iter().position(|&p| p == policy);
+
+            let output = Command::new(powershell)
+                .args(["-NoProfile", "-Command", "Get-ExecutionPolicy"])
+                .output_checked_utf8();
+
+            if let Ok(output) = output {
+                let current_policy = output.stdout.trim();
+
+                // Find the index of the current policy
+                let current_idx = valid_policies.iter().position(|&p| p == current_policy);
+
+                // Check if current policy exists and is at least as permissive as the target
+                return match (current_idx, target_idx) {
+                    (Some(current), Some(target)) => current >= target,
+                    _ => false,
+                };
+            }
+        }
+        false
+    }
+}
+
+#[cfg(windows)]
+impl Powershell {
+    pub fn supports_windows_update(&self) -> bool {
+        windows::supports_windows_update(self)
+    }
+
+    pub fn windows_update(&self, ctx: &ExecutionContext) -> Result<()> {
+        windows::windows_update(self, ctx)
+    }
+
+    pub fn microsoft_store(&self, ctx: &ExecutionContext) -> Result<()> {
+        windows::microsoft_store(self, ctx)
+    }
+}
+
+#[cfg(windows)]
+mod windows {
+    use super::*;
+
+    pub fn supports_windows_update(powershell: &Powershell) -> bool {
+        powershell
+            .path
+            .as_ref()
+            .map(|p| has_module(p, "PSWindowsUpdate"))
+            .unwrap_or(false)
+    }
+
+    #[cfg(windows)]
+    pub fn windows_update(powershell: &Powershell, ctx: &ExecutionContext) -> Result<()> {
+        debug_assert!(supports_windows_update(powershell));
+
+        // Build the full command string
+        let mut command_str = "Install-WindowsUpdate -Verbose".to_string();
+        if ctx.config().accept_all_windows_updates() {
+            command_str.push_str(" -AcceptAll");
+        }
+
+        // Pass the command string using the -Command flag
+        powershell
+            .build_command_internal(ctx, &["-Command", &command_str])?
+            .status_checked()
+    }
+
+    pub fn microsoft_store(powershell: &Powershell, ctx: &ExecutionContext) -> Result<()> {
+        println!("{}", t!("Scanning for updates..."));
+        let update_command = "Start-Process powershell -Verb RunAs -ArgumentList '-Command', \
+            '(Get-CimInstance -Namespace \"Root\\cimv2\\mdm\\dmmap\" \
+            -ClassName \"MDM_EnterpriseModernAppManagement_AppManagement01\" | \
+            Invoke-CimMethod -MethodName UpdateScanMethod).ReturnValue'";
+
+        powershell
+            .build_command_internal(ctx, &["-Command", update_command])?
+            .status_checked()
+    }
+
+    fn has_module(powershell: &PathBuf, command: &str) -> bool {
         Command::new(powershell)
             .args([
                 "-NoProfile",
                 "-Command",
-                &format!("Get-Module -ListAvailable {command}"),
+                &format!("Get-Module -ListAvailable {}", command),
             ])
             .output_checked_utf8()
             .map(|result| !result.stdout.is_empty())
             .unwrap_or(false)
     }
-
-    pub fn profile(&self) -> Option<&PathBuf> {
-        self.profile.as_ref()
-    }
-
-    pub fn update_modules(&self, ctx: &ExecutionContext) -> Result<()> {
-        let powershell = require_option(self.path.as_ref(), t!("Powershell is not installed").to_string())?;
-
-        print_separator(t!("Powershell Modules Update"));
-
-        let mut cmd = vec!["Update-Module"];
-
-        if ctx.config().verbose() {
-            cmd.push("-Verbose");
-        }
-
-        if ctx.config().yes(Step::Powershell) {
-            cmd.push("-Force");
-        }
-
-        println!("{}", t!("Updating modules..."));
-        ctx.run_type()
-            .execute(powershell)
-            // This probably doesn't need `shell_words::join`.
-            .args(["-NoProfile", "-Command", &cmd.join(" ")])
-            .status_checked()
-    }
-
-    #[cfg(windows)]
-    pub fn supports_windows_update(&self) -> bool {
-        self.path
-            .as_ref()
-            .map(|p| Self::has_module(p, "PSWindowsUpdate"))
-            .unwrap_or(false)
-    }
-
-    #[cfg(windows)]
-    pub fn windows_update(&self, ctx: &ExecutionContext) -> Result<()> {
-        let powershell = require_option(self.path.as_ref(), t!("Powershell is not installed").to_string())?;
-
-        debug_assert!(self.supports_windows_update());
-
-        let accept_all = if ctx.config().accept_all_windows_updates() {
-            "-AcceptAll"
-        } else {
-            ""
-        };
-
-        let install_windowsupdate_verbose = "Install-WindowsUpdate -Verbose".to_string();
-
-        let mut command = if let Some(sudo) = ctx.sudo() {
-            let mut command = ctx.run_type().execute(sudo);
-            command.arg(powershell);
-            command
-        } else {
-            ctx.run_type().execute(powershell)
-        };
-
-        command
-            .args(["-NoProfile", &install_windowsupdate_verbose, accept_all])
-            .status_checked()
-    }
-
-    #[cfg(windows)]
-    pub fn microsoft_store(&self, ctx: &ExecutionContext) -> Result<()> {
-        let powershell = require_option(self.path.as_ref(), t!("Powershell is not installed").to_string())?;
-
-        let mut command = if let Some(sudo) = ctx.sudo() {
-            let mut command = ctx.run_type().execute(sudo);
-            command.arg(powershell);
-            command
-        } else {
-            ctx.run_type().execute(powershell)
-        };
-
-        println!("{}", t!("Scanning for updates..."));
-
-        // Scan for updates using the MDM UpdateScanMethod
-        // This method is also available for non-MDM devices
-        let update_command = "(Get-CimInstance -Namespace \"Root\\cimv2\\mdm\\dmmap\" -ClassName \"MDM_EnterpriseModernAppManagement_AppManagement01\" | Invoke-CimMethod -MethodName UpdateScanMethod).ReturnValue";
-
-        command.args(["-NoProfile", update_command]);
-
-        command
-            .output_checked_with_utf8(|output| {
-                if output.stdout.trim() == "0" {
-                    println!(
-                        "{}",
-                        t!("Success, Microsoft Store apps are being updated in the background")
-                    );
-                    Ok(())
-                } else {
-                    println!(
-                        "{}",
-                        t!("Unable to update Microsoft Store apps, manual intervention is required")
-                    );
-                    Err(())
-                }
-            })
-            .map(|_| ())
-    }
 }