2Shirt
dd03962c84
This simplifies the code quite a bit and the Option<T> logic is only used for string parsing. s.is_empty() is just fine in this case IMO.
629 lines
25 KiB
Rust
629 lines
25 KiB
Rust
// 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 <https://www.gnu.org/licenses/>.
|
|
//
|
|
use std::{
|
|
env,
|
|
iter::zip,
|
|
path::PathBuf,
|
|
sync::{Arc, Mutex},
|
|
thread::{self, sleep, JoinHandle},
|
|
time::Duration,
|
|
};
|
|
|
|
use color_eyre::Result;
|
|
use crossterm::event::KeyEvent;
|
|
use ratatui::{
|
|
layout::{Constraint, Direction, Layout},
|
|
prelude::Rect,
|
|
};
|
|
use serde::{Deserialize, Serialize};
|
|
use tokio::sync::mpsc;
|
|
use tracing::{debug, info};
|
|
|
|
use crate::{
|
|
action::Action,
|
|
components::{
|
|
footer::Footer, fps::FpsCounter, left::Left, popup, right::Right, title::Title, Component,
|
|
},
|
|
config::Config,
|
|
system::{
|
|
disk::{get_disks, Disk, PartitionTableType},
|
|
diskpart::{build_dest_format_script, refresh_disk_info},
|
|
drivers::{self, Driver},
|
|
},
|
|
tui::{Event, Tui},
|
|
};
|
|
|
|
pub struct App {
|
|
// TUI
|
|
action_rx: mpsc::UnboundedReceiver<Action>,
|
|
action_tx: mpsc::UnboundedSender<Action>,
|
|
components: Vec<Box<dyn Component>>,
|
|
config: Config,
|
|
frame_rate: f64,
|
|
last_tick_key_events: Vec<KeyEvent>,
|
|
should_quit: bool,
|
|
should_suspend: bool,
|
|
tick_rate: f64,
|
|
// App
|
|
cur_mode: Mode,
|
|
disk_index_dest: Option<usize>,
|
|
disk_index_source: Option<usize>,
|
|
disk_list: Arc<Mutex<Vec<Disk>>>,
|
|
part_index_boot: Option<usize>,
|
|
part_index_os: Option<usize>,
|
|
driver: Option<Driver>,
|
|
prev_mode: Mode,
|
|
selections: Vec<Option<usize>>,
|
|
table_type: Option<PartitionTableType>,
|
|
task_handles: Vec<JoinHandle<()>>,
|
|
}
|
|
|
|
#[derive(Default, Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
|
pub enum Mode {
|
|
#[default]
|
|
ScanDisks,
|
|
InstallDrivers,
|
|
SelectDisks,
|
|
SelectTableType,
|
|
Confirm,
|
|
PreClone,
|
|
Clone,
|
|
SelectParts,
|
|
PostClone,
|
|
Done,
|
|
}
|
|
|
|
impl App {
|
|
pub fn new(tick_rate: f64, frame_rate: f64) -> Result<Self> {
|
|
let (action_tx, action_rx) = mpsc::unbounded_channel();
|
|
Ok(Self {
|
|
// TUI
|
|
action_rx,
|
|
action_tx,
|
|
components: vec![
|
|
Box::new(Title::new()),
|
|
Box::new(FpsCounter::default()),
|
|
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
|
|
cur_mode: Mode::ScanDisks,
|
|
disk_index_dest: None,
|
|
disk_index_source: None,
|
|
disk_list: Arc::new(Mutex::new(Vec::new())),
|
|
driver: None,
|
|
part_index_boot: None,
|
|
part_index_os: None,
|
|
prev_mode: Mode::ScanDisks,
|
|
selections: vec![None, None],
|
|
table_type: None,
|
|
task_handles: Vec::new(),
|
|
})
|
|
}
|
|
|
|
pub fn next_mode(&mut self) -> Option<Mode> {
|
|
let new_mode = match (self.prev_mode, self.cur_mode) {
|
|
(_, Mode::InstallDrivers) => Mode::ScanDisks,
|
|
(_, Mode::ScanDisks) => Mode::SelectDisks,
|
|
(_, Mode::SelectDisks | Mode::SelectTableType | Mode::SelectParts) => {
|
|
if self.selections[1].is_some() {
|
|
Mode::Confirm
|
|
} else {
|
|
self.cur_mode
|
|
}
|
|
}
|
|
(Mode::SelectDisks, Mode::Confirm) => Mode::SelectTableType,
|
|
(Mode::SelectTableType, Mode::Confirm) => Mode::PreClone,
|
|
(_, Mode::PreClone) => Mode::Clone,
|
|
(_, Mode::Clone) => Mode::SelectParts,
|
|
(Mode::SelectParts, Mode::Confirm) => Mode::PostClone,
|
|
(_, Mode::PostClone) => Mode::Done,
|
|
(_, Mode::Done) => Mode::Done,
|
|
// Invalid states
|
|
(_, Mode::Confirm) => panic!("This shouldn't happen."),
|
|
};
|
|
|
|
if new_mode != self.cur_mode {
|
|
match self.cur_mode {
|
|
// Update prev_mode if appropriate
|
|
Mode::Confirm => {}
|
|
Mode::PreClone | Mode::Clone | Mode::PostClone | Mode::Done => {
|
|
// Override since we're past the point of no return
|
|
self.prev_mode = self.cur_mode;
|
|
}
|
|
_ => self.prev_mode = self.cur_mode,
|
|
}
|
|
Some(new_mode)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
pub async fn run(&mut self) -> Result<()> {
|
|
let disk_list_arc = Arc::clone(&self.disk_list);
|
|
self.task_handles.push(lazy_get_disks(disk_list_arc));
|
|
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 self.components.iter_mut() {
|
|
component.register_action_handler(self.action_tx.clone())?;
|
|
}
|
|
for component in self.components.iter_mut() {
|
|
component.register_config_handler(self.config.clone())?;
|
|
}
|
|
for component in self.components.iter_mut() {
|
|
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 self.components.iter_mut() {
|
|
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(());
|
|
};
|
|
match keymap.get(&vec![key]) {
|
|
Some(action) => {
|
|
info!("Got action: {action:?}");
|
|
action_tx.send(action.clone())?;
|
|
}
|
|
_ => {
|
|
// 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
|
|
if let Ok(_) = &self.disk_list.try_lock() {
|
|
if self.task_handles.is_empty() {
|
|
// All tasks complete
|
|
self.action_tx.send(Action::NextScreen)?;
|
|
} else {
|
|
// Tasks remain, check if current is complete
|
|
if let Some(handle) = self.task_handles.first() {
|
|
if handle.is_finished() {
|
|
self.task_handles.remove(0);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
Action::Quit => self.should_quit = true,
|
|
Action::Suspend => self.should_suspend = true,
|
|
Action::Resume => self.should_suspend = false,
|
|
Action::ClearScreen => tui.terminal.clear()?,
|
|
Action::Command(ref cmd_path, ref cmd_args) => self
|
|
.task_handles
|
|
.push(lazy_command(&cmd_path, cmd_args.clone())),
|
|
Action::Diskpart(ref script) => {
|
|
self.task_handles.push(lazy_diskpart(&script.to_owned()))
|
|
}
|
|
Action::InstallDriver => {
|
|
self.action_tx.send(Action::SetMode(Mode::InstallDrivers))?
|
|
}
|
|
Action::SelectDriver(ref driver) => {
|
|
self.driver = Some(driver.clone());
|
|
drivers::load(&driver.inf_paths);
|
|
self.action_tx.send(Action::NextScreen)?;
|
|
}
|
|
Action::SelectTableType(ref table_type) => {
|
|
self.table_type = Some(table_type.clone());
|
|
self.action_tx.send(Action::NextScreen)?;
|
|
}
|
|
Action::Resize(w, h) => self.handle_resize(tui, w, h)?,
|
|
Action::Render => self.render(tui)?,
|
|
Action::PrevScreen => {
|
|
self.action_tx.send(Action::SetMode(self.prev_mode))?;
|
|
}
|
|
Action::NextScreen => {
|
|
if let Some(mode) = self.next_mode() {
|
|
self.action_tx.send(Action::DismissPopup)?;
|
|
self.action_tx.send(Action::SetMode(mode))?;
|
|
}
|
|
}
|
|
Action::ScanDisks => {
|
|
let disk_list_arc = Arc::clone(&self.disk_list);
|
|
self.task_handles.push(lazy_get_disks(disk_list_arc));
|
|
self.action_tx.send(Action::SetMode(Mode::ScanDisks))?;
|
|
}
|
|
Action::Select(one, two) => {
|
|
match self.cur_mode {
|
|
Mode::SelectDisks => {
|
|
self.disk_index_source = one.clone();
|
|
self.disk_index_dest = two.clone();
|
|
}
|
|
Mode::SelectParts => {
|
|
self.part_index_boot = one.clone();
|
|
self.part_index_os = two.clone();
|
|
}
|
|
_ => {}
|
|
}
|
|
self.selections[0] = one;
|
|
self.selections[1] = two;
|
|
}
|
|
Action::SetMode(new_mode) => {
|
|
self.cur_mode = new_mode;
|
|
match new_mode {
|
|
Mode::ScanDisks => {
|
|
self.prev_mode = self.cur_mode;
|
|
}
|
|
Mode::SelectDisks | Mode::SelectParts => {
|
|
let disk_list = self.disk_list.lock().unwrap();
|
|
self.action_tx
|
|
.send(Action::UpdateDiskList(disk_list.clone()))?;
|
|
}
|
|
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.disk_list.lock().unwrap();
|
|
if let Some(disk_index) = self.disk_index_dest {
|
|
if let Some(disk) = disk_list.get(disk_index) {
|
|
let table_type = self.table_type.clone().unwrap();
|
|
let diskpart_script =
|
|
build_dest_format_script(disk.id, &table_type);
|
|
self.action_tx.send(Action::Diskpart(diskpart_script))?;
|
|
}
|
|
}
|
|
}
|
|
Mode::Clone => {
|
|
self.action_tx.send(Action::DisplayPopup(
|
|
popup::Type::Info,
|
|
String::from("Running Clone Tool"),
|
|
))?;
|
|
self.action_tx.send(Action::Command(
|
|
self.config.clone_app_path.clone(),
|
|
Vec::new(),
|
|
))?;
|
|
}
|
|
Mode::PostClone => {
|
|
// TODO: FIXME
|
|
self.action_tx.send(Action::DisplayPopup(
|
|
popup::Type::Info,
|
|
String::from("Updating boot configuration"),
|
|
))?;
|
|
|
|
// Get System32 path
|
|
let system32 = if cfg!(windows) {
|
|
match env::var("SYSTEMROOT") {
|
|
Ok(path) => path,
|
|
Err(_) => panic!("Failed to find SYSTEMROOT"),
|
|
}
|
|
} else {
|
|
String::from(".")
|
|
};
|
|
|
|
// Add actions
|
|
let disk_list = self.disk_list.lock().unwrap();
|
|
if let Some(disk_index) = self.disk_index_dest {
|
|
if let Some(disk) = disk_list.get(disk_index) {
|
|
let table_type = self.table_type.clone().unwrap();
|
|
let letter_boot: String;
|
|
let letter_os: String;
|
|
if let Some(part) =
|
|
disk.parts.get(self.part_index_boot.unwrap())
|
|
{
|
|
letter_boot = if part.letter.is_empty() {
|
|
String::from("??")
|
|
} else {
|
|
part.letter.clone()
|
|
};
|
|
letter_os = if part.letter.is_empty() {
|
|
String::from("??")
|
|
} else {
|
|
part.letter.clone()
|
|
};
|
|
} else {
|
|
self.action_tx.send(Action::DisplayPopup(
|
|
popup::Type::Error,
|
|
String::from(
|
|
"Failed to get drive letters for destination",
|
|
),
|
|
))?;
|
|
self.action_tx.send(Action::SetMode(Mode::Done))?;
|
|
return Ok(());
|
|
}
|
|
|
|
// Create boot files
|
|
self.action_tx.send(Action::Command(
|
|
PathBuf::from(format!("{system32}/bcdboot.exe")),
|
|
vec![
|
|
format!("{letter_os}:\\Windows"),
|
|
String::from("/s"),
|
|
format!("{letter_boot}:"),
|
|
String::from("/f"),
|
|
String::from(match table_type {
|
|
PartitionTableType::Guid => "UEFI",
|
|
PartitionTableType::Legacy => "BIOS",
|
|
}),
|
|
],
|
|
))?;
|
|
|
|
// Update boot sector (for legacy setups)
|
|
if table_type == PartitionTableType::Legacy {
|
|
//
|
|
self.action_tx.send(Action::Command(
|
|
PathBuf::from(format!("{system32}/bootsect.exe")),
|
|
vec![
|
|
String::from("/nt60"),
|
|
format!("{letter_boot}:"),
|
|
String::from("/force"),
|
|
String::from("/mbr"),
|
|
],
|
|
))?;
|
|
}
|
|
|
|
// Lock in safe mode
|
|
let bcd_path = match table_type {
|
|
PartitionTableType::Guid => {
|
|
format!("{letter_boot}\\EFI\\Microsoft\\Boot\\BCD")
|
|
}
|
|
PartitionTableType::Legacy => {
|
|
format!("{letter_boot}\\Boot\\BCD")
|
|
}
|
|
};
|
|
self.action_tx.send(Action::Command(
|
|
PathBuf::from(format!("{system32}/bcdedit.exe")),
|
|
vec![
|
|
String::from("/store"),
|
|
bcd_path,
|
|
String::from("/set"),
|
|
String::from("{default}"),
|
|
String::from("safeboot"),
|
|
String::from("minimal"),
|
|
],
|
|
))?;
|
|
|
|
// Inject driver(s) (if selected)
|
|
if let Some(driver) = &self.driver {
|
|
if let Some(os_path) = driver.path.to_str() {
|
|
let driver_path_str = String::from(os_path);
|
|
self.action_tx.send(Action::Command(
|
|
PathBuf::from(format!("{system32}/dism.exe")),
|
|
vec![
|
|
format!("/image:{letter_os}:\\"),
|
|
String::from("/add-driver"),
|
|
format!("/driver:\"{}\"", driver_path_str,),
|
|
String::from("/recurse"),
|
|
],
|
|
))?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Mode::Done => {
|
|
self.action_tx.send(Action::DisplayPopup(
|
|
popup::Type::Success,
|
|
String::from("COMPLETE\n\n\nThank you for using this tool!"),
|
|
))?;
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
Action::UpdateDestDisk => {
|
|
let disk_list_arc = Arc::clone(&self.disk_list);
|
|
if let Some(dest_index) = self.disk_index_dest {
|
|
self.task_handles
|
|
.push(lazy_update_dest_disk(disk_list_arc, dest_index));
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
for component in self.components.iter_mut() {
|
|
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<Rect> {
|
|
let mut chunks: Vec<Rect> = 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 lazy_command(cmd_path: &PathBuf, cmd_args: Vec<String>) -> JoinHandle<()> {
|
|
info!("Running Command: {cmd_path:?} {cmd_args:?}");
|
|
if cfg!(windows) {
|
|
// TODO: FIXME
|
|
thread::spawn(|| sleep(Duration::from_secs(1)))
|
|
} else {
|
|
thread::spawn(|| sleep(Duration::from_secs(1)))
|
|
}
|
|
}
|
|
|
|
fn lazy_diskpart(script: &str) -> JoinHandle<()> {
|
|
if cfg!(windows) {
|
|
// TODO: FIXME
|
|
thread::spawn(|| sleep(Duration::from_secs(1)))
|
|
} else {
|
|
info!("Running (lazy) Diskpart: {:?}", &script);
|
|
thread::spawn(|| sleep(Duration::from_secs(1)))
|
|
}
|
|
}
|
|
|
|
fn lazy_get_disks(disk_list_arc: Arc<Mutex<Vec<Disk>>>) -> JoinHandle<()> {
|
|
thread::spawn(move || {
|
|
let mut disks = disk_list_arc.lock().unwrap();
|
|
*disks = get_disks();
|
|
})
|
|
}
|
|
|
|
fn lazy_update_dest_disk(
|
|
disk_list_arc: Arc<Mutex<Vec<Disk>>>,
|
|
dest_index: usize,
|
|
) -> JoinHandle<()> {
|
|
if cfg!(windows) {
|
|
thread::spawn(move || {
|
|
let mut disks = disk_list_arc.lock().unwrap();
|
|
refresh_disk_info(&mut disks[dest_index]);
|
|
})
|
|
} else {
|
|
thread::spawn(|| sleep(Duration::from_secs(2)))
|
|
}
|
|
}
|