add logic for non-nvme drives

src/main.rs

@@ -50,8 +50,15 @@ fn format_drive(drive_name: &str) -> rsfdisk::Result<()> {
 }
 
 fn format_partitions(drive_name: &str) {
-    let efi_path = format!("{}p1", drive_name);
+    // Check if drive name ends with a digit (e.g., "nvme0n1")
-    let root_path = format!("{}p2", drive_name);
+    let suffix = if drive_name.chars().last().map(|c| c.is_ascii_digit()).unwrap_or(false) {
+        "p"
+    } else {
+        ""
+    };
+
+    let efi_path = format!("{}{}1", drive_name, suffix);
+    let root_path = format!("{}{}2", drive_name, suffix);
 
     let _efi_partition = Command::new("mkfs.fat")
         .arg("-F32")
@@ -68,8 +75,15 @@ fn format_partitions(drive_name: &str) {
 }
 
 fn mount_partitions(drive_name: &str) {
-    let efi_path = format!("{}p1", drive_name);
+    // Check if drive name ends with a digit (e.g., "nvme0n1")
-    let root_path = format!("{}p2", drive_name);
+    let suffix = if drive_name.chars().last().map(|c| c.is_ascii_digit()).unwrap_or(false) {
+        "p"
+    } else {
+        ""
+    };
+
+    let efi_path = format!("{}{}1", drive_name, suffix);
+    let root_path = format!("{}{}2", drive_name, suffix);
 
     let root_source = Some(Path::new(&root_path));
     let root_target = Path::new("/mnt");
@@ -246,7 +260,8 @@ fn main() {
         println!("---------------");
         println!("");
         println!("1. Lemur Pro 13 (Garrus)");
-        println!("2. Generic");
+        println!("2. Virtual Machine");
+        println!("3. Generic");
         println!("");
 
         println!("Enter your host for installation:");
@@ -289,8 +304,43 @@ fn main() {
                 println!("Exited with status {:?}", install_status);
 
                 break
-            }
+            },
             "2" => {
+                let mut vm_config = Easy::new();
+                vm_config.url("https://gitlab.com/ahoneybun/nix-configs/-/raw/main/hosts/x86_64/vm/configuration.nix").unwrap();
+
+                let mut config_file = File::create("vm.nix").unwrap();
+
+                {
+                    let mut transfer = vm_config.transfer();
+                    transfer.write_function(|data| {
+                        config_file.write_all(data).unwrap();
+                        Ok(data.len())
+                    }).unwrap();
+                    transfer.perform().unwrap();
+                }
+
+                // Copies the system nix files to /mnt/etc/nixos/
+                let _garrus_nix_copy = Command::new("mv")
+                    .args(["vm.nix", "/mnt/etc/nixos"])
+                    .output()
+                    .expect("Failed to copy nix files over");
+
+                let mut nixos_install = Command::new("nixos-install")
+                    .arg("--flake")
+                    .arg("/mnt/etc/nixos#vm")
+                    .arg("--no-root-passwd")
+                    .stdout(Stdio::inherit())
+                    .stderr(Stdio::inherit())
+                    .spawn()
+                    .unwrap();
+
+                let install_status = nixos_install.wait();
+                println!("Exited with status {:?}", install_status);
+
+                break;
+            }
+            "3" => {
                 let mut nixos_install = Command::new("nixos-install")
                     .arg("--flake")
                     .arg("/mnt/etc/nixos#nixos")