diff --git a/src/main.rs b/src/main.rs
index acfc649..b8a1133 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -7,13 +7,18 @@ fn main() {
     let args: Vec<String> = env::args().collect();
 
     let iso_mount_dir = Path::new(args.get(1)
-        .expect("Arguments: rsfus (path to XP ISO mount)"))
+        .expect("Arguments: rsfus (path to XP ISO mount) [disk number]"))
         .canonicalize().expect("Can't canonicalize path to XP ISO mount.");
 
     if !iso_mount_dir.exists() || !iso_mount_dir.is_dir() {
         panic!("XP ISO mount is not a directory.");
     }
 
+    let disk_number = if let Some(i) = args.get(2) {
+        i.parse::<u8>().expect("Not a valid number for disk number!")
+    } else { 1 };
+    let setup_source_line = format!("SetupSourceDevice = \"\\device\\harddisk{disk_number}\\partition1\"");
+
     let amd64_basedir = iso_mount_dir.join("amd64");
     let i386_basedir = iso_mount_dir.join("i386");
 
@@ -47,8 +52,8 @@ fn main() {
         for line in txtsetup_contents.lines() {
             txtsetup_lines.push(line);
             if line == "[SetupData]" {
-                println!("TXTSETP: Found [SetupData]");
-                txtsetup_lines.push(r#"SetupSourceDevice = "\device\harddisk1\partition1""#);
+                println!("TXTSETP: Found [SetupData], writing \"{setup_source_line}\"");
+                txtsetup_lines.push(setup_source_line.as_str());
             }
         }
         println!("Writing {txtsetup_copy:?}");
@@ -83,8 +88,8 @@ fn main() {
 
         for index in 0..(bootmgr_contents.len() - win_nt_bt.len()) {
             if &bootmgr_contents[index..index + rdisk0.len()] == rdisk0 {
-                println!("BOOTMGR: {index:#010x}  Patching rdisk(0)  -->  rdisk(1)");
-                bootmgr_contents[index + 6] = 0x31;
+                println!("BOOTMGR: {index:#010x}  Patching rdisk(0)  -->  rdisk({disk_number})");
+                bootmgr_contents[index + 6] = 0x30 + disk_number;
             }
 
             if &bootmgr_contents[index..index + win_nt_bt.len()] == win_nt_bt