Add win-installer partitioning/formatting logic

This commit is contained in:
2Shirt 2025-12-13 08:09:12 -08:00
parent cb7ba1a285
commit a75911cb32
Signed by: 2Shirt
GPG key ID: 152FAC923B0E132C
3 changed files with 242 additions and 24 deletions

View file

@ -32,6 +32,12 @@ use crate::system::disk::{
static DEFAULT_MAX_DISKS: usize = 8;
#[derive(Debug, PartialEq)]
pub enum FormatUseCase {
ApplyWimImage,
Clone,
}
pub struct RegexList {
detail_all_disks: OnceLock<Regex>,
detail_disk: OnceLock<Regex>,
@ -205,7 +211,11 @@ pub fn get_partitions(
}
#[must_use]
pub fn build_dest_format_script(disk_id: usize, part_type: &PartitionTableType) -> String {
pub fn build_dest_format_script(
disk_id: usize,
part_type: &PartitionTableType,
format_use_case: FormatUseCase,
) -> String {
let disk_id = format!("{disk_id}");
let mut script = vec!["automount enable noerr", "select disk {disk_id}", "clean"];
match part_type {
@ -221,6 +231,10 @@ pub fn build_dest_format_script(disk_id: usize, part_type: &PartitionTableType)
script.push("format fs=ntfs quick label=System");
}
}
if format_use_case == FormatUseCase::ApplyWimImage {
script.push("create partition primary");
script.push("format fs=ntfs quick label=Windows");
}
script.join("\r\n").replace("{disk_id}", &disk_id)
}

View file

@ -28,7 +28,10 @@ use core::{
line::{DVLine, get_disk_description_right, get_part_description},
state::Mode,
system::{
boot, cpu::get_cpu_name, disk::PartitionTableType, diskpart::build_dest_format_script,
boot,
cpu::get_cpu_name,
disk::PartitionTableType,
diskpart::{FormatUseCase, build_dest_format_script},
drivers,
},
tasks::{Task, TaskResult, TaskType, Tasks},
@ -209,7 +212,8 @@ impl App {
&& let Some(disk) = disk_list.get(disk_index)
{
let table_type = self.state.table_type.clone().unwrap();
let diskpart_script = build_dest_format_script(disk.id, &table_type);
let diskpart_script =
build_dest_format_script(disk.id, &table_type, FormatUseCase::Clone);
self.tasks.add(TaskType::Diskpart(diskpart_script));
}
}

View file

@ -26,12 +26,20 @@ use core::{
config::Config,
line::{DVLine, get_disk_description_right},
state::Mode,
system::{cpu::get_cpu_name, disk::PartitionTableType, drivers},
tasks::{TaskType, Tasks},
system::{
boot,
cpu::get_cpu_name,
disk::PartitionTableType,
diskpart::{FormatUseCase, build_dest_format_script},
drivers,
},
tasks::{Task, TaskResult, TaskType, Tasks},
tui::{Event, Tui},
};
use std::{
env,
iter::zip,
path::PathBuf,
sync::{Arc, Mutex},
};
@ -111,20 +119,20 @@ impl App {
Mode::SelectWinSource => Mode::SelectWinImage,
Mode::SelectWinImage => Mode::SetUserName,
Mode::SetUserName => Mode::Confirm,
Mode::Confirm => Mode::Process, // i.e. format, apply, etc
Mode::Process | Mode::Done => Mode::Done,
Mode::Confirm => Mode::PreClone,
Mode::PreClone => Mode::Clone,
Mode::Clone => Mode::PostClone,
Mode::PostClone | Mode::Done => Mode::Done,
Mode::Failed => Mode::Failed,
// Invalid States
Mode::BootDiags
| Mode::BootScan
| Mode::BootSetup
| Mode::Clone
| Mode::DiagMenu
| Mode::InjectDrivers
| Mode::LogView
| Mode::PEMenu
| Mode::PreClone
| Mode::PostClone
| Mode::Process
| Mode::SelectParts
| Mode::SetBootMode => panic!("This shouldn't happen?"),
}
@ -139,6 +147,126 @@ impl App {
// self.action_tx
// .send(Action::DisplayPopup(popup::Type::Info, String::from("...")))?;
// }
Mode::PreClone => {
self.action_tx.send(Action::DisplayPopup(
popup::Type::Info,
String::from("Formatting destination disk"),
))?;
// Get System32 path
let system32 = get_system32_path(&self.action_tx);
// (Re)Enable volume mounting
self.tasks.add(TaskType::CommandWait(
PathBuf::from(format!("{system32}/mountvol.exe")),
vec![String::from("/e")],
));
// Build Diskpart script to format destination disk
let disk_list = self.state.disk_list.lock().unwrap();
if let Some(disk_index) = self.state.disk_index_dest
&& let Some(disk) = disk_list.get(disk_index)
{
let table_type = self.state.table_type.clone().unwrap();
let diskpart_script = build_dest_format_script(
disk.id,
&table_type,
FormatUseCase::ApplyWimImage,
);
self.tasks.add(TaskType::Diskpart(diskpart_script));
}
// Update drive letters
self.tasks.add(TaskType::Sleep);
if let Some(dest_index) = self.state.disk_index_dest {
self.tasks.add(TaskType::UpdateDestDisk(dest_index));
}
}
Mode::Clone => {
self.action_tx.send(Action::DisplayPopup(
popup::Type::Info,
String::from("Applying Image"),
))?;
// Get wimlib-imagex path
let program_files = PathBuf::from(get_program_files_path(&self.action_tx));
let wimlib_imagex = program_files.join("wimlib\\wimlib-imagex.exe");
// Get image info
let wim_sources = self.state.wim_sources.lock().unwrap();
let wim_file = wim_sources.get_file(self.state.wim_file_index.unwrap());
let wim_index = format!("{}", self.state.wim_image_index.unwrap());
// Add actions
let disk_list = self.state.disk_list.lock().unwrap();
if let Some(disk_index) = self.state.disk_index_dest
&& let Some(disk) = disk_list.get(disk_index)
{
let num_parts = disk.parts.len();
let dest_path = format!("{}:\\", disk.get_part_letter(num_parts - 1));
self.tasks.add(TaskType::CommandWait(
wimlib_imagex,
vec![String::from("apply"), wim_file.path, wim_index, dest_path],
));
}
}
Mode::PostClone => {
self.action_tx.send(Action::DisplayPopup(
popup::Type::Info,
String::from("Updating boot configuration"),
))?;
// Get System32 path
let system32 = get_system32_path(&self.action_tx);
// Add actions
let disk_list = self.state.disk_list.lock().unwrap();
if let Some(disk_index) = self.state.disk_index_dest
&& let Some(disk) = disk_list.get(disk_index)
{
let table_type = self.state.table_type.clone().unwrap();
let letter_boot = disk.get_part_letter(0);
let letter_os =
disk.get_part_letter(match self.state.table_type.clone().unwrap() {
PartitionTableType::Guid => 2,
PartitionTableType::Legacy => 1,
});
info!("PostClone // Disk: {disk:?}");
info!("\t\tBoot letter: {letter_boot}");
info!("\t\tOS letter: {letter_os}");
// Safety check
if letter_boot.is_empty() || letter_os.is_empty() {
self.action_tx.send(Action::Error(String::from(
"ERROR\n\n\nFailed to get drive letters for the destination",
)))?;
return Ok(());
}
// Create boot files
for task in boot::configure_disk(
&letter_boot,
&letter_os,
boot::SafeMode::Enable,
&system32,
&table_type,
) {
self.tasks.add(task);
}
// Inject driver(s) (if selected)
if let Some(driver) = &self.state.driver {
if let Ok(task) = boot::inject_driver(driver, &letter_os, &system32) {
self.tasks.add(task);
} else {
self.action_tx.send(Action::Error(format!(
"Failed to inject driver:\n{}",
driver.name
)))?;
}
}
}
}
Mode::ScanDisks => {
self.state.reset_all();
if self.tasks.idle() {
@ -261,7 +389,9 @@ impl App {
Action::Tick => {
self.last_tick_key_events.drain(..);
// Check background task(s)
self.tasks.poll()?;
if let Some(task) = self.tasks.poll()? {
self.handle_task(&task)?;
}
if let Ok(mut wim_sources) = self.state.wim_sources.lock() {
wim_sources.poll();
}
@ -317,6 +447,11 @@ impl App {
self.action_tx.send(Action::DismissPopup)?;
self.action_tx.send(Action::SetMode(next_mode))?;
}
Action::DisplayPopup(ref popup_type, ref _popup_text) => {
if *popup_type == popup::Type::Error {
self.action_tx.send(Action::SetMode(Mode::Failed))?;
}
}
Action::PrevScreen => match self.cur_mode {
Mode::SelectTableType => {
self.action_tx.send(Action::SetMode(Mode::SelectDisks))?;
@ -405,6 +540,39 @@ impl App {
Ok(())
}
fn handle_task(&mut self, task: &Task) -> Result<()> {
match task.task_type {
TaskType::CommandWait(_, _) | TaskType::Diskpart(_) => {
// Check result
if let Some(result) = &task.result {
match result {
TaskResult::Error(msg) => {
self.action_tx
.send(Action::Error(format!("{} Failed: {msg}", task.task_type)))?;
}
TaskResult::Output(stdout, stderr, success) => {
if !success {
let msg = if !stdout.is_empty() {
stdout.clone()
} else if !stderr.is_empty() {
stderr.clone()
} else {
String::from("Unknown Error")
};
self.action_tx.send(Action::Error(format!(
"{} Failed: {msg}",
task.task_type
)))?;
}
}
}
}
}
_ => {}
}
Ok(())
}
fn render(&mut self, tui: &mut Tui) -> Result<()> {
tui.draw(|frame| {
if let [header, _body, footer, center, left, right, username, popup] =
@ -426,7 +594,9 @@ impl App {
fn build_footer_string(cur_mode: Mode) -> String {
match cur_mode {
Mode::Home | Mode::ScanDisks => String::from("(q) to quit"),
Mode::Home | Mode::ScanDisks | Mode::PreClone | Mode::Clone | Mode::PostClone => {
String::from("(q) to quit")
}
Mode::InstallDrivers => String::from("(Enter) to select / (q) to quit"),
Mode::SelectDisks => String::from(
"(Enter) to select / / (i) to install driver / (r) to rescan / (q) to quit",
@ -440,18 +610,16 @@ fn build_footer_string(cur_mode: Mode) -> String {
),
Mode::SetUserName => String::from("(Enter) to continue / (Esc) to go back"),
Mode::Confirm => String::from("(Enter) to confirm / (b) to go back / (q) to quit"),
Mode::Done | Mode::Failed | Mode::Process => String::from("(Enter) or (q) to quit"),
Mode::Done | Mode::Failed => String::from("(Enter) or (q) to quit"),
// Invalid States
Mode::BootDiags
| Mode::BootScan
| Mode::BootSetup
| Mode::Clone
| Mode::DiagMenu
| Mode::InjectDrivers
| Mode::LogView
| Mode::PEMenu
| Mode::PreClone
| Mode::PostClone
| Mode::Process
| Mode::SelectParts
| Mode::SetBootMode => panic!("This shouldn't happen?"),
}
@ -475,15 +643,9 @@ fn build_left_items(app: &App) -> Action {
.iter()
.for_each(|driver| items.push(driver.to_string()));
}
Mode::Process => {
select_type = SelectionType::Loop;
title = String::from("Processing");
// TODO: FIXME
}
Mode::ScanWinSources => {
select_type = SelectionType::Loop;
title = String::from("Scanning");
// TODO: FIXME
}
Mode::SelectWinSource => {
select_type = SelectionType::One;
@ -539,6 +701,7 @@ fn build_left_items(app: &App) -> Action {
| Mode::BootScan
| Mode::BootSetup
| Mode::DiagMenu
| Mode::Process
| Mode::InjectDrivers
| Mode::LogView
| Mode::PEMenu
@ -653,7 +816,12 @@ fn build_right_items(app: &App) -> Action {
});
}
}
Mode::SelectWinImage | Mode::SetUserName | Mode::Confirm => {
Mode::SelectWinImage
| Mode::SetUserName
| Mode::Confirm
| Mode::PreClone
| Mode::Clone
| Mode::PostClone => {
info!("Building right items for: {:?}", &app.cur_mode);
let wim_file;
if let Ok(wim_sources) = app.state.wim_sources.lock()
@ -717,7 +885,7 @@ fn build_right_items(app: &App) -> Action {
}])
});
}
Mode::Confirm => {
Mode::Confirm | Mode::PreClone | Mode::Clone | Mode::PostClone => {
if let Some(index) = app.state.wim_image_index
&& let Some(image) = wim_file.images.get(index)
{
@ -842,3 +1010,35 @@ fn get_chunks(r: Rect) -> Vec<Rect> {
// Done
chunks
}
pub fn get_program_files_path(action_tx: &mpsc::UnboundedSender<Action>) -> String {
let mut program_files_path = String::from(".");
if cfg!(windows) {
if let Ok(path) = env::var("PROGRAMFILES") {
program_files_path = path;
} else {
action_tx
.send(Action::Error(String::from(
"ERROR\n\n\nFailed to find PROGRAMFILES",
)))
.expect("Failed to find PROGRAMFILES and then failed to send action");
}
}
program_files_path
}
pub fn get_system32_path(action_tx: &mpsc::UnboundedSender<Action>) -> String {
let mut system32_path = String::from(".");
if cfg!(windows) {
if let Ok(path) = env::var("SYSTEMROOT") {
system32_path = format!("{path}/System32");
} else {
action_tx
.send(Action::Error(String::from(
"ERROR\n\n\nFailed to find SYSTEMROOT",
)))
.expect("Failed to find SYSTEMROOT and then failed to send action");
}
}
system32_path
}