// This file is part of Deja-vu. // // Deja-vu is free software: you can redistribute it and/or modify it // under the terms of the GNU General Public License as published by // the Free Software Foundation, either version 3 of the License, or // (at your option) any later version. // // Deja-vu is distributed in the hope that it will be useful, but // WITHOUT ANY WARRANTY; without even the implied warranty of // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. // See the GNU General Public License for more details. // // You should have received a copy of the GNU General Public License // along with Deja-vu. If not, see . // use core::{ action::Action, components::{ footer::Footer, fps::FpsCounter, left::Left, popup, right::Right, state::StatefulList, title::Title, Component, }, config::Config, line::{get_disk_description_right, get_part_description, DVLine}, state::{CloneSettings, Mode}, system::{ boot, cpu::get_cpu_name, disk::PartitionTableType, diskpart::build_dest_format_script, drivers, }, tasks::{Task, Tasks}, tui::{Event, Tui}, }; use std::{ env, iter::zip, sync::{Arc, Mutex}, }; use color_eyre::Result; use ratatui::{ crossterm::event::KeyEvent, layout::{Constraint, Direction, Layout}, prelude::Rect, style::Color, }; use tokio::sync::mpsc; use tracing::{debug, info}; pub struct App { // TUI action_rx: mpsc::UnboundedReceiver, action_tx: mpsc::UnboundedSender, components: Vec>, config: Config, frame_rate: f64, last_tick_key_events: Vec, should_quit: bool, should_suspend: bool, tick_rate: f64, // App clone: CloneSettings, cur_mode: Mode, list: StatefulList, prev_mode: Mode, selections: Vec>, tasks: Tasks, } impl App { pub fn new(tick_rate: f64, frame_rate: f64) -> Result { let (action_tx, action_rx) = mpsc::unbounded_channel(); let disk_list_arc = Arc::new(Mutex::new(Vec::new())); let mut tasks = Tasks::new(action_tx.clone(), disk_list_arc.clone()); tasks.add(Task::ScanDisks); Ok(Self { // TUI action_rx, action_tx, components: vec![ Box::new(Title::new()), Box::new(FpsCounter::new()), Box::new(Left::new()), Box::new(Right::new()), Box::new(Footer::new()), Box::new(popup::Popup::new()), ], config: Config::new()?, frame_rate, last_tick_key_events: Vec::new(), should_quit: false, should_suspend: false, tick_rate, // App clone: CloneSettings::new(disk_list_arc), cur_mode: Mode::ScanDisks, list: StatefulList::default(), prev_mode: Mode::ScanDisks, selections: vec![None, None], tasks, }) } pub fn prev_mode(&mut self) -> Option { let new_mode = match self.cur_mode { Mode::Failed => Some(Mode::Failed), Mode::Done => Some(Mode::Done), Mode::SelectParts => Some(Mode::SelectParts), Mode::Confirm => Some(Mode::SelectTableType), Mode::SelectTableType => Some(Mode::SelectDisks), Mode::SelectDisks => Some(Mode::ScanDisks), // Mode::InstallDrivers | Mode::ScanDisks | Mode::PreClone | Mode::Clone | Mode::PostClone => None, }; new_mode } pub fn next_mode(&mut self) -> Option { let new_mode = match self.cur_mode { Mode::InstallDrivers => Mode::ScanDisks, Mode::ScanDisks => Mode::SelectDisks, Mode::SelectDisks => Mode::SelectTableType, Mode::SelectTableType => Mode::Confirm, Mode::Confirm => Mode::PreClone, Mode::PreClone => Mode::Clone, Mode::Clone => Mode::SelectParts, Mode::SelectParts => Mode::PostClone, Mode::PostClone | Mode::Done => Mode::Done, Mode::Failed => Mode::Failed, }; if new_mode == self.cur_mode { // No mode change needed None } else { Some(new_mode) } } pub fn set_mode(&mut self, new_mode: Mode) -> Result<()> { info!("Setting mode to {new_mode:?}"); self.cur_mode = new_mode; match new_mode { Mode::InstallDrivers => self.clone.scan_drivers(), Mode::ScanDisks => { self.prev_mode = self.cur_mode; if self.tasks.idle() { self.tasks.add(Task::ScanDisks); } self.action_tx.send(Action::DisplayPopup( popup::Type::Info, String::from("Scanning Disks..."), ))?; } Mode::PreClone => { self.action_tx.send(Action::DisplayPopup( popup::Type::Info, String::from("Formatting destination disk"), ))?; // Build Diskpart script to format destination disk let disk_list = self.clone.disk_list.lock().unwrap(); if let Some(disk_index) = self.clone.disk_index_dest { if let Some(disk) = disk_list.get(disk_index) { let table_type = self.clone.table_type.clone().unwrap(); let diskpart_script = build_dest_format_script(disk.id, &table_type); self.tasks.add(Task::Diskpart(diskpart_script)); } } } Mode::Clone => { self.action_tx.send(Action::DisplayPopup( popup::Type::Info, String::from("Running Clone Tool"), ))?; self.tasks.add(Task::Command( self.config.clone_app_path.clone(), Vec::new(), )); if let Some(dest_index) = self.clone.disk_index_dest { self.tasks.add(Task::UpdateDestDisk(dest_index)); } } Mode::PostClone => { self.action_tx.send(Action::DisplayPopup( popup::Type::Info, String::from("Updating boot configuration"), ))?; // Get System32 path let system32 = if cfg!(windows) { if let Ok(path) = env::var("SYSTEMROOT") { format!("{path}/System32") } else { self.action_tx.send(Action::Error(String::from( "ERROR\n\n\nFailed to find SYSTEMROOT", )))?; return Ok(()); } } else { String::from(".") }; // Add actions let disk_list = self.clone.disk_list.lock().unwrap(); if let Some(disk_index) = self.clone.disk_index_dest { if let Some(disk) = disk_list.get(disk_index) { let table_type = self.clone.table_type.clone().unwrap(); let letter_boot = disk.get_part_letter(self.clone.part_index_boot.unwrap()); let letter_os = disk.get_part_letter(self.clone.part_index_os.unwrap()); // 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, &system32, table_type) { self.tasks.add(task); } // Inject driver(s) (if selected) if let Some(driver) = &self.clone.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::Done => { self.action_tx.send(Action::DisplayPopup( popup::Type::Success, String::from("COMPLETE\n\n\nThank you for using this tool!"), ))?; } _ => {} } Ok(()) } pub async fn run(&mut self) -> Result<()> { let mut tui = Tui::new()? // .mouse(true) // uncomment this line to enable mouse support .tick_rate(self.tick_rate) .frame_rate(self.frame_rate); tui.enter()?; for component in &mut self.components { component.register_action_handler(self.action_tx.clone())?; } for component in &mut self.components { component.register_config_handler(self.config.clone())?; } for component in &mut self.components { component.init(tui.size()?)?; } let action_tx = self.action_tx.clone(); loop { self.handle_events(&mut tui).await?; self.handle_actions(&mut tui)?; if self.should_suspend { tui.suspend()?; action_tx.send(Action::Resume)?; action_tx.send(Action::ClearScreen)?; // tui.mouse(true); tui.enter()?; } else if self.should_quit { tui.stop()?; break; } } tui.exit()?; Ok(()) } async fn handle_events(&mut self, tui: &mut Tui) -> Result<()> { let Some(event) = tui.next_event().await else { return Ok(()); }; let action_tx = self.action_tx.clone(); match event { Event::Quit => action_tx.send(Action::Quit)?, Event::Tick => action_tx.send(Action::Tick)?, Event::Render => action_tx.send(Action::Render)?, Event::Resize(x, y) => action_tx.send(Action::Resize(x, y))?, Event::Key(key) => self.handle_key_event(key)?, _ => {} } for component in &mut self.components { if let Some(action) = component.handle_events(Some(event.clone()))? { action_tx.send(action)?; } } Ok(()) } fn handle_key_event(&mut self, key: KeyEvent) -> Result<()> { let action_tx = self.action_tx.clone(); let Some(keymap) = self.config.keybindings.get(&self.cur_mode) else { return Ok(()); }; if let Some(action) = keymap.get(&vec![key]) { info!("Got action: {action:?}"); action_tx.send(action.clone())?; } else { // If the key was not handled as a single key action, // then consider it for multi-key combinations. self.last_tick_key_events.push(key); // Check for multi-key combinations if let Some(action) = keymap.get(&self.last_tick_key_events) { info!("Got action: {action:?}"); action_tx.send(action.clone())?; } } Ok(()) } fn handle_actions(&mut self, tui: &mut Tui) -> Result<()> { while let Ok(action) = self.action_rx.try_recv() { if action != Action::Tick && action != Action::Render { debug!("{action:?}"); } match action { Action::Tick => { self.last_tick_key_events.drain(..); match self.cur_mode { Mode::ScanDisks | Mode::PreClone | Mode::Clone | Mode::PostClone => { // Check background task self.tasks.poll()?; // Once all are complete Action::NextScreen is sent } _ => {} } } Action::Quit => self.should_quit = true, Action::Suspend => self.should_suspend = true, Action::Resume => self.should_suspend = false, Action::ClearScreen => tui.terminal.clear()?, Action::KeyUp => self.list.previous(), Action::KeyDown => self.list.next(), Action::Error(ref msg) => { self.action_tx .send(Action::DisplayPopup(popup::Type::Error, msg.clone()))?; self.action_tx.send(Action::SetMode(Mode::Failed))?; } Action::InstallDriver => { self.action_tx.send(Action::SetMode(Mode::InstallDrivers))?; } Action::Process => match self.cur_mode { Mode::Confirm => { self.action_tx.send(Action::NextScreen)?; } _ => {} }, Action::Resize(w, h) => self.handle_resize(tui, w, h)?, Action::Render => self.render(tui)?, Action::PrevScreen => { if let Some(new_mode) = self.prev_mode() { self.prev_mode = new_mode; self.cur_mode = new_mode; self.action_tx.send(Action::SetMode(new_mode))?; } } Action::NextScreen => match self.next_mode() { None => {} Some(next) => { self.prev_mode = self.cur_mode; self.cur_mode = next; self.action_tx.send(Action::DismissPopup)?; self.action_tx.send(Action::SetMode(next))?; } }, Action::ScanDisks => self.action_tx.send(Action::SetMode(Mode::ScanDisks))?, Action::Select(one, two) => { match self.cur_mode { Mode::InstallDrivers => { if let Some(index) = one { if let Some(driver) = self.clone.driver_list.get(index).cloned() { drivers::load(&driver.inf_paths); self.clone.driver = Some(driver); } } } Mode::SelectDisks => { self.clone.disk_index_source = one; self.clone.disk_index_dest = two; } Mode::SelectParts => { self.clone.part_index_boot = one; self.clone.part_index_os = two; } Mode::SelectTableType => { self.clone.table_type = { if let Some(index) = one { match index { 0 => Some(PartitionTableType::Guid), 1 => Some(PartitionTableType::Legacy), index => { panic!("Failed to select PartitionTableType: {}", index) } } } else { None } } } _ => {} } self.selections[0] = one; self.selections[1] = two; } Action::SetMode(new_mode) => { self.set_mode(new_mode)?; self.action_tx .send(Action::UpdateFooter(build_footer_string(self.cur_mode)))?; let (title, labels, items, select_one) = build_left_items(self, self.cur_mode); self.action_tx .send(Action::UpdateLeft(title, labels, items, select_one))?; let (labels, start, items) = build_right_items(self, self.cur_mode); self.action_tx .send(Action::UpdateRight(labels, start, items))?; match new_mode { // Mode::InstallDrivers | Mode::SelectDisks => { // self.action_tx.send(Action::Select(None, None))? // } Mode::SelectTableType | Mode::Confirm => { // Select source/dest disks self.action_tx.send(Action::SelectRight( self.clone.disk_index_source, self.clone.disk_index_dest, ))?; } Mode::SelectParts => { // Select first partition as boot partition self.action_tx.send(Action::Select(Some(0), None))?; // Highlight 2nd or 3rd partition as OS partition let index = if let Some(table_type) = &self.clone.table_type { match table_type { PartitionTableType::Guid => 2, PartitionTableType::Legacy => 1, } } else { 1 }; self.action_tx.send(Action::Highlight(index))?; } _ => {} }; } _ => {} } for component in &mut self.components { if let Some(action) = component.update(action.clone())? { self.action_tx.send(action)?; }; } } Ok(()) } fn handle_resize(&mut self, tui: &mut Tui, w: u16, h: u16) -> Result<()> { tui.resize(Rect::new(0, 0, w, h))?; self.render(tui)?; Ok(()) } fn render(&mut self, tui: &mut Tui) -> Result<()> { tui.draw(|frame| { if let [header, _body, footer, left, right, popup] = get_chunks(frame.area())[..] { let component_areas = vec![ header, // Title Bar header, // FPS Counter left, right, footer, popup, ]; for (component, area) in zip(self.components.iter_mut(), component_areas) { if let Err(err) = component.draw(frame, area) { let _ = self .action_tx .send(Action::Error(format!("Failed to draw: {err:?}"))); } } }; })?; Ok(()) } } fn centered_rect(percent_x: u16, percent_y: u16, r: Rect) -> Rect { // Cut the given rectangle into three vertical pieces let popup_layout = Layout::default() .direction(Direction::Vertical) .constraints([ Constraint::Percentage((100 - percent_y) / 2), Constraint::Percentage(percent_y), Constraint::Percentage((100 - percent_y) / 2), ]) .split(r); // Then cut the middle vertical piece into three width-wise pieces Layout::default() .direction(Direction::Horizontal) .constraints([ Constraint::Percentage((100 - percent_x) / 2), Constraint::Percentage(percent_x), Constraint::Percentage((100 - percent_x) / 2), ]) .split(popup_layout[1])[1] // Return the middle chunk } fn get_chunks(r: Rect) -> Vec { let mut chunks: Vec = Vec::with_capacity(6); // Main sections chunks.extend( Layout::default() .direction(Direction::Vertical) .constraints([ Constraint::Length(3), Constraint::Min(1), Constraint::Length(3), ]) .split(r) .to_vec(), ); // Left/Right chunks.extend( Layout::default() .direction(Direction::Horizontal) .constraints([Constraint::Percentage(50), Constraint::Percentage(50)]) .split(centered_rect(90, 90, chunks[1])) .to_vec(), ); // Popup chunks.push(centered_rect(60, 25, r)); // Done chunks } fn build_footer_string(cur_mode: Mode) -> String { match cur_mode { Mode::ScanDisks | Mode::PreClone | Mode::Clone | Mode::PostClone => { String::from("(q) to quit") } Mode::SelectParts => String::from("(Enter) to select / (s) to start over / (q) to quit"), Mode::SelectDisks => String::from( "(Enter) to select / / (i) to install driver / (r) to rescan / (q) to quit", ), Mode::SelectTableType => String::from("(Enter) to select / (b) to go back / (q) to quit"), Mode::Confirm => String::from("(Enter) to confirm / (b) to go back / (q) to quit"), Mode::Done | Mode::Failed | Mode::InstallDrivers => String::from("(Enter) or (q) to quit"), } } fn build_left_items(app: &App, cur_mode: Mode) -> (String, Vec, Vec, bool) { let title: String; let mut items = Vec::new(); let mut labels: Vec = Vec::new(); let mut select_one = false; match cur_mode { Mode::InstallDrivers => { title = String::from("Install Drivers"); select_one = true; app.clone .driver_list .iter() .for_each(|driver| items.push(driver.to_string())); } Mode::SelectDisks => { title = String::from("Select Source and Destination Disks"); labels.push(String::from("source")); labels.push(String::from("dest")); let disk_list = app.clone.disk_list.lock().unwrap(); disk_list .iter() .for_each(|disk| items.push(disk.description.to_string())); } Mode::SelectTableType => { title = String::from("Select Partition Table Type"); select_one = true; items.push(format!("{}", PartitionTableType::Guid)); items.push(format!("{}", PartitionTableType::Legacy)); } Mode::Confirm => { title = String::from("Confirm Selections"); } Mode::PreClone | Mode::Clone | Mode::PostClone | Mode::ScanDisks => { title = String::from("Processing"); } Mode::SelectParts => { title = String::from("Select Boot and OS Partitions"); labels.push(String::from("boot")); labels.push(String::from("os")); let disk_list = app.clone.disk_list.lock().unwrap(); if let Some(index) = app.clone.disk_index_dest { if let Some(disk) = disk_list.get(index) { disk.get_parts().iter().for_each(|part| { items.push(part.to_string()); }); } } } Mode::Done | Mode::Failed => title = String::from("Done"), }; (title, labels, items, select_one) } fn build_right_items(app: &App, cur_mode: Mode) -> (Vec>, usize, Vec>) { let mut items = Vec::new(); let mut labels: Vec> = Vec::new(); let mut start_index = 0; match cur_mode { Mode::InstallDrivers => { items.push(vec![DVLine { line_parts: vec![String::from("CPU")], line_colors: vec![Color::Cyan], }]); items.push(vec![DVLine { line_parts: vec![get_cpu_name()], line_colors: vec![Color::Reset], }]); start_index = 2; } Mode::SelectDisks | Mode::SelectTableType | Mode::Confirm => { // Labels labels.push(vec![DVLine { line_parts: vec![String::from("Source")], line_colors: vec![Color::Cyan], }]); let dest_dv_line = DVLine { line_parts: vec![ String::from("Dest"), String::from(" (WARNING: ALL DATA WILL BE DELETED!)"), ], line_colors: vec![Color::Cyan, Color::Red], }; if let Some(table_type) = &app.clone.table_type { // Show table type let type_str = match table_type { PartitionTableType::Guid => "GPT", PartitionTableType::Legacy => "MBR", }; labels.push(vec![ dest_dv_line, DVLine { line_parts: vec![format!(" (Will be formatted {type_str})")], line_colors: vec![Color::Yellow], }, ]); } else { labels.push(vec![dest_dv_line]); } let disk_list = app.clone.disk_list.lock().unwrap(); disk_list .iter() .for_each(|disk| items.push(get_disk_description_right(&disk))); } Mode::SelectParts => { vec!["Boot", "OS"].iter().for_each(|s| { labels.push(vec![DVLine { line_parts: vec![String::from(*s)], line_colors: vec![Color::Cyan], }]) }); if let Some(index) = app.clone.disk_index_dest { start_index = 1; let disk_list = app.clone.disk_list.lock().unwrap(); if let Some(disk) = disk_list.get(index) { // Disk Details items.push(get_disk_description_right(&disk)); // Partition Details disk.parts .iter() .for_each(|part| items.push(get_part_description(&part))); } } } _ => {} } (labels, start_index, items) }