diff --git a/petri/src/vm/hyperv/powershell.rs b/petri/src/vm/hyperv/powershell.rs index 0f7bcd699..b12a2acc2 100644 --- a/petri/src/vm/hyperv/powershell.rs +++ b/petri/src/vm/hyperv/powershell.rs @@ -5,7 +5,8 @@ use anyhow::Context; use core::str; -use std::fmt::Display; +use std::ffi::OsStr; +use std::ffi::OsString; use std::path::Path; use std::process::Command; @@ -18,16 +19,12 @@ pub enum HyperVGeneration { Two, } -impl Display for HyperVGeneration { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - match self { - HyperVGeneration::One => { - write!(f, "1") - } - HyperVGeneration::Two => { - write!(f, "2") - } - } +impl AsRef for HyperVGeneration { + fn as_ref(&self) -> &OsStr { + OsStr::new(match self { + HyperVGeneration::One => "1", + HyperVGeneration::Two => "2", + }) } } @@ -48,28 +45,16 @@ pub enum HyperVGuestStateIsolationType { Disabled, } -impl Display for HyperVGuestStateIsolationType { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - match self { - HyperVGuestStateIsolationType::TrustedLaunch => { - write!(f, "TrustedLaunch") - } - HyperVGuestStateIsolationType::Vbs => { - write!(f, "VBS") - } - HyperVGuestStateIsolationType::Snp => { - write!(f, "SNP") - } - HyperVGuestStateIsolationType::Tdx => { - write!(f, "TDX") - } - HyperVGuestStateIsolationType::OpenHCL => { - write!(f, "OpenHCL") - } - HyperVGuestStateIsolationType::Disabled => { - write!(f, "Disabled") - } - } +impl AsRef for HyperVGuestStateIsolationType { + fn as_ref(&self) -> &OsStr { + OsStr::new(match self { + HyperVGuestStateIsolationType::TrustedLaunch => "TrustedLaunch", + HyperVGuestStateIsolationType::Vbs => "VBS", + HyperVGuestStateIsolationType::Snp => "SNP", + HyperVGuestStateIsolationType::Tdx => "TDX", + HyperVGuestStateIsolationType::OpenHCL => "OpenHCL", + HyperVGuestStateIsolationType::Disabled => "Disabled", + }) } } @@ -86,22 +71,16 @@ pub enum HyperVSecureBootTemplate { OpenSourceShieldedVM, } -impl Display for HyperVSecureBootTemplate { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - match self { - HyperVSecureBootTemplate::SecureBootDisabled => { - write!(f, "SecureBootDisabled") - } - HyperVSecureBootTemplate::MicrosoftWindows => { - write!(f, "MicrosoftWindows") - } +impl AsRef for HyperVSecureBootTemplate { + fn as_ref(&self) -> &OsStr { + OsStr::new(match self { + HyperVSecureBootTemplate::SecureBootDisabled => "SecureBootDisabled", + HyperVSecureBootTemplate::MicrosoftWindows => "MicrosoftWindows", HyperVSecureBootTemplate::MicrosoftUEFICertificateAuthority => { - write!(f, "MicrosoftUEFICertificateAuthority") - } - HyperVSecureBootTemplate::OpenSourceShieldedVM => { - write!(f, "OpenSourceShieldedVM") + "MicrosoftUEFICertificateAuthority" } - } + HyperVSecureBootTemplate::OpenSourceShieldedVM => "OpenSourceShieldedVM", + }) } } @@ -123,31 +102,29 @@ pub struct HyperVNewVMArgs<'a> { /// Runs New-VM with the given arguments. pub fn run_new_vm(args: HyperVNewVMArgs<'_>) -> anyhow::Result<()> { - run_powershell_cmdlet("New-VM", |cmd| { - if let Some(generation) = args.generation { - cmd.arg("-Generation").arg(generation.to_string()); - } - if let Some(guest_state_isolation_type) = args.guest_state_isolation_type { - cmd.arg("-GuestStateIsolationType") - .arg(guest_state_isolation_type.to_string()); - } - if let Some(memory_startup_bytes) = args.memory_startup_bytes { - cmd.arg("-MemoryStartupBytes") - .arg(memory_startup_bytes.to_string()); - } - if let Some(path) = args.path { - cmd.arg("-Path").arg(path); - } - if let Some(vhd_path) = args.vhd_path { - cmd.arg("-VHDPath").arg(vhd_path); - } - cmd.arg("-Name").arg(args.name).arg("-Force") - }) + PowerShellBuilder::new() + .cmdlet("New-VM") + .arg("Name", args.name) + .arg_opt("Generation", args.generation) + .arg_opt("GuestStateIsolationType", args.guest_state_isolation_type) + .arg_opt_string("MemoryStartupBytes", args.memory_startup_bytes) + .arg_opt("Path", args.path) + .arg_opt("VHDPath", args.vhd_path) + .flag("Force") + .finish() + .run() + .context("new_vm") } /// Runs New-VM with the given arguments. pub fn run_remove_vm(name: &str) -> anyhow::Result<()> { - run_powershell_cmdlet("Remove-VM", |cmd| cmd.arg("-Name").arg(name).arg("-Force")) + PowerShellBuilder::new() + .cmdlet("Remove-VM") + .arg("Name", name) + .flag("Force") + .finish() + .run() + .context("remove_vm") } /// Arguments for the Add-VMHardDiskDrive powershell cmdlet @@ -171,20 +148,15 @@ pub struct HyperVAddVMHardDiskDriveArgs<'a> { /// Runs Add-VMHardDiskDrive with the given arguments. pub fn run_add_vm_hard_disk_drive(args: HyperVAddVMHardDiskDriveArgs<'_>) -> anyhow::Result<()> { - run_powershell_cmdlet("Add-VMHardDiskDrive", |cmd| { - if let Some(controller_location) = args.controller_location { - cmd.arg("-ControllerLocation") - .arg(controller_location.to_string()); - } - if let Some(controller_number) = args.controller_number { - cmd.arg("-ControllerNumber") - .arg(controller_number.to_string()); - } - if let Some(path) = args.path { - cmd.arg("-Path").arg(path); - } - cmd.arg("-VMName").arg(args.name) - }) + PowerShellBuilder::new() + .cmdlet("Add-VMHardDiskDrive") + .arg("VMName", args.name) + .arg_opt_string("ControllerLocation", args.controller_location) + .arg_opt_string("ControllerNumber", args.controller_number) + .arg_opt("Path", args.path) + .finish() + .run() + .context("add_vm_hard_disk_drive") } /// Arguments for the Add-VMDvdDrive powershell cmdlet @@ -206,91 +178,47 @@ pub struct HyperVAddVMDvdDriveArgs<'a> { /// Runs Add-VMDvdDrive with the given arguments. pub fn run_add_vm_dvd_drive(args: HyperVAddVMDvdDriveArgs<'_>) -> anyhow::Result<()> { - run_powershell_cmdlet("Add-VMDvdDrive", |cmd| { - if let Some(controller_location) = args.controller_location { - cmd.arg("-ControllerLocation") - .arg(controller_location.to_string()); - } - if let Some(controller_number) = args.controller_number { - cmd.arg("-ControllerNumber") - .arg(controller_number.to_string()); - } - if let Some(path) = args.path { - cmd.arg("-Path").arg(path); - } - cmd.arg("-VMName").arg(args.name) - }) + PowerShellBuilder::new() + .cmdlet("Add-VMDvdDrive") + .arg("VMName", args.name) + .arg_opt_string("ControllerLocation", args.controller_location) + .arg_opt_string("ControllerNumber", args.controller_number) + .arg_opt("Path", args.path) + .finish() + .run() + .context("add_vm_dvd_drive") } /// Runs Add-VMScsiController with the given arguments. pub fn run_add_vm_scsi_controller(name: &str) -> anyhow::Result<()> { - run_powershell_cmdlet("Add-VMScsiController", |cmd| cmd.arg("-VMName").arg(name)) -} - -/// Arguments for creating a new VHD -pub struct CreateVhdArgs<'a> { - /// VHD path - pub path: &'a Path, - /// Filesystem label - pub label: &'a str, -} - -/// Create a new VHD, mount, initialize, and format. Returns drive letter. -pub fn create_vhd(args: CreateVhdArgs<'_>) -> anyhow::Result { - let drive_letter = run_powershell_cmdlet_output("New-VHD", |cmd| { - cmd.arg("-Path") - .arg(args.path) - .arg("-Fixed") - .arg("-SizeBytes") - .arg("64MB"); - - cmd.arg("|").arg("Mount-VHD").arg("-Passthru"); - - cmd.arg("|").arg("Initialize-Disk").arg("-Passthru"); - - cmd.arg("|") - .arg("New-Partition") - .arg("-AssignDriveLetter") - .arg("-UseMaximumSize"); - - cmd.arg("|") - .arg("Format-Volume") - .arg("-FileSystem") - .arg("FAT32") - .arg("-Force") - .arg("-NewFileSystemLabel") - .arg(args.label); - - cmd.arg("|") - .arg("Select-Object") - .arg("-ExpandProperty") - .arg("DriveLetter") - })?; - - if drive_letter.trim().len() != 1 { - anyhow::bail!("invalid drive letter: {drive_letter}"); - } - - drive_letter - .chars() - .next() - .context("could not get drive letter") + PowerShellBuilder::new() + .cmdlet("Add-VMScsiController") + .arg("VMName", name) + .finish() + .run() + .context("add_vm_scsi_controller") } /// Create a new differencing VHD with the provided parent. pub fn create_child_vhd(path: &Path, parent_path: &Path) -> anyhow::Result<()> { - run_powershell_cmdlet("New-VHD", |cmd| { - cmd.arg("-Path") - .arg(path) - .arg("-ParentPath") - .arg(parent_path) - .arg("-Differencing") - }) + PowerShellBuilder::new() + .cmdlet("New-VHD") + .arg("Path", path) + .arg("ParentPath", parent_path) + .flag("Differencing") + .finish() + .run() + .context("create_child_vhd") } /// Runs Dismount-VHD with the given arguments. pub fn run_dismount_vhd(path: &Path) -> anyhow::Result<()> { - run_powershell_cmdlet("Dismount-VHD", |cmd| cmd.arg("-Path").arg(path)) + PowerShellBuilder::new() + .cmdlet("Dismount-VHD") + .arg("Path", path) + .finish() + .run() + .context("dismount_vhd") } /// Arguments for the Set-VMFirmware powershell cmdlet @@ -306,13 +234,13 @@ pub struct HyperVSetVMFirmwareArgs<'a> { /// Runs Set-VMFirmware with the given arguments. pub fn run_set_vm_firmware(args: HyperVSetVMFirmwareArgs<'_>) -> anyhow::Result<()> { - run_powershell_cmdlet("Set-VMFirmware", |cmd| { - if let Some(secure_boot_template) = args.secure_boot_template { - cmd.arg("-SecureBootTemplate") - .arg(secure_boot_template.to_string()); - } - cmd.arg("-VMName").arg(args.name) - }) + PowerShellBuilder::new() + .cmdlet("Set-VMFirmware") + .arg_opt("SecureBootTemplate", args.secure_boot_template) + .arg("VMName", args.name) + .finish() + .run() + .context("set_vm_firmware") } /// Runs Set-VMFirmware with the given arguments. @@ -322,14 +250,17 @@ pub fn run_set_openhcl_firmware( igvm_file: &Path, increase_vtl2_memory: bool, ) -> anyhow::Result<()> { - run_powershell(Some("Set-OpenHCLFirmware"), |cmd| { - cmd.arg("Import-Module").arg(ps_mod).arg(";"); - cmd.arg("Set-OpenHCLFirmware"); - if increase_vtl2_memory { - cmd.arg("-IncreaseVtl2Memory"); - } - cmd.arg("-VMName").arg(name).arg("-IgvmFile").arg(igvm_file) - }) + PowerShellBuilder::new() + .cmdlet("Import-Module") + .positional(ps_mod) + .next() + .cmdlet("Set-OpenHCLFirmware") + .arg("VMName", name) + .arg("IgvmFile", igvm_file) + .flag_opt(increase_vtl2_memory.then_some("IncreaseVtl2Memory")) + .finish() + .run() + .context("set_openhcl_firmware") } /// Sets the initial machine configuration for a VM @@ -338,68 +269,150 @@ pub fn run_set_initial_machine_configuration( ps_mod: &Path, imc_hive: &Path, ) -> anyhow::Result<()> { - run_powershell(Some("Set-InitialMachineConfiguration"), |cmd| { - cmd.arg("Import-Module").arg(ps_mod).arg(";"); - cmd.arg("Set-InitialMachineConfiguration") - .arg("-VMName") - .arg(name) - .arg("-ImcHive") - .arg(imc_hive) - }) + PowerShellBuilder::new() + .cmdlet("Import-Module") + .positional(ps_mod) + .next() + .cmdlet("Set-InitialMachineConfiguration") + .arg("VMName", name) + .arg("ImcHive", imc_hive) + .finish() + .run() + .context("set_initial_machine_configuration") } -/// Runs a powershell cmdlet with the given arguments. -fn run_powershell_cmdlet( - cmdlet: &str, - f: impl FnOnce(&mut Command) -> &mut Command, -) -> anyhow::Result<()> { - run_powershell(None, |cmd| { - cmd.arg(cmdlet); - f(cmd) - }) -} +/// A PowerShell script builder +pub struct PowerShellBuilder(Command); -/// Runs a powershell cmdlet with the given arguments. -fn run_powershell( - name: Option<&str>, - f: impl FnOnce(&mut Command) -> &mut Command, -) -> anyhow::Result<()> { - let mut cmd = Command::new("powershell.exe"); - cmd.arg("-NoProfile"); - f(&mut cmd); - let cmdlet = cmd - .get_args() - .next() - .context("no cmdlet in args")? - .to_string_lossy() - .into_owned(); - let name = name.unwrap_or(&cmdlet); - let status = cmd - .status() - .context(format!("failed to launch powershell cmdlet {name}"))?; - if !status.success() { - anyhow::bail!("powershell cmdlet {name} failed with exit code: {status}"); +impl PowerShellBuilder { + /// Create a new PowerShell command + pub fn new() -> Self { + let mut cmd = Command::new("powershell.exe"); + cmd.arg("-NoProfile"); + Self(cmd) + } + + /// Start a new Cmdlet + pub fn cmdlet>(mut self, cmdlet: S) -> PowerShellCmdletBuilder { + self.0.arg(cmdlet); + PowerShellCmdletBuilder(self.0) + } + + /// Run the PowerShell script + pub fn run(mut self) -> anyhow::Result<()> { + let status = self.0.status().context("failed to launch powershell")?; + if !status.success() { + anyhow::bail!("powershell script failed with exit code: {}", status); + } + Ok(()) + } + + /// Run the PowerShell script and return the output + pub fn output(mut self) -> anyhow::Result { + let output = self.0.output().context("failed to launch powershell")?; + if !output.status.success() { + anyhow::bail!("powershell script failed with exit code: {}", output.status); + } + String::from_utf8(output.stdout).context("powershell output is not utf-8") + } + + /// Use Select-Object to return a property of the returned object + pub fn select_object_property>( + mut self, + property: S, + ) -> PowerShellCmdletBuilder { + self.0 + .arg("Select-Object") + .arg("-ExpandProperty") + .arg(property); + PowerShellCmdletBuilder(self.0) } - Ok(()) } -/// Runs a powershell cmdlet with the given arguments and returns the output -fn run_powershell_cmdlet_output( - cmdlet: &str, - f: impl FnOnce(&mut Command) -> &mut Command, -) -> anyhow::Result { - let mut cmd = Command::new("powershell.exe"); - cmd.arg("-NoProfile"); - cmd.arg(cmdlet); - f(&mut cmd); - let output = cmd - .output() - .context(format!("failed to launch powershell cmdlet {cmdlet}"))?; - if !output.status.success() { - anyhow::bail!( - "powershell cmdlet {cmdlet} failed with exit code: {}", - output.status - ); +/// A PowerShell Cmdlet builder +pub struct PowerShellCmdletBuilder(Command); + +impl PowerShellCmdletBuilder { + /// Add a flag to the cmdlet + pub fn flag>(mut self, flag: S) -> Self { + let mut arg = OsString::from("-"); + arg.push(flag); + self.0.arg(arg); + self + } + + /// Optionally add a flag to the cmdlet + pub fn flag_opt>(self, flag: Option) -> Self { + if let Some(flag) = flag { + self.flag(flag) + } else { + self + } + } + + /// Add a positional argument to the cmdlet + pub fn positional>(mut self, positional: S) -> Self { + self.0.arg(positional); + self + } + + /// Add a positional argument to the cmdlet + pub fn positional_string(self, positional: S) -> Self { + self.positional(positional.to_string()) + } + + /// Optionally add a positional argument to the cmdlet + pub fn positional_opt>(self, positional: Option) -> Self { + if let Some(positional) = positional { + self.positional(positional) + } else { + self + } + } + + /// Optionally add a positional argument to the cmdlet + pub fn positional_opt_string(self, positional: Option) -> Self { + self.positional_opt(positional.map(|x| x.to_string())) + } + + /// Add an argument to the cmdlet + pub fn arg, T: AsRef>(self, name: S, value: T) -> Self { + self.flag(name).positional(value) + } + + /// Add an argument to the cmdlet + pub fn arg_string, T: ToString>(self, name: S, value: T) -> Self { + self.arg(name, value.to_string()) + } + + /// Optionally add an argument to the cmdlet + pub fn arg_opt, T: AsRef>(self, name: S, value: Option) -> Self { + if let Some(value) = value { + self.arg(name, value) + } else { + self + } + } + + /// Optionally add an argument to the cmdlet + pub fn arg_opt_string, T: ToString>(self, name: S, value: Option) -> Self { + self.arg_opt(name, value.map(|x| x.to_string())) + } + + /// Finish the cmdlet + pub fn finish(self) -> PowerShellBuilder { + PowerShellBuilder(self.0) + } + + /// Finish the cmdlet with a pipeline operator + pub fn pipeline(mut self) -> PowerShellBuilder { + self.0.arg("|"); + self.finish() + } + + /// Finish the cmdlet with a semicolon + pub fn next(mut self) -> PowerShellBuilder { + self.0.arg(";"); + self.finish() } - String::from_utf8(output.stdout).context("output is not utf-8") }