deja-vu/deja_vu/src/app.rs

736 lines
28 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

// 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 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 rand::random;
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>,
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
clone: CloneSettings,
cur_mode: Mode,
list: StatefulList<usize>,
prev_mode: Mode,
selections: Vec<Option<usize>>,
tasks: Tasks,
}
impl App {
pub fn new(tick_rate: f64, frame_rate: f64) -> Result<Self> {
let (action_tx, action_rx) = mpsc::unbounded_channel();
let disk_list_arc = Arc::new(Mutex::new(Vec::new()));
let tasks = Tasks::new(action_tx.clone(), disk_list_arc.clone());
Ok(Self {
// TUI
action_rx,
action_tx,
components: vec![
Box::new(Title::new("Clone Tool")),
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::default(),
list: StatefulList::default(),
prev_mode: Mode::default(),
selections: vec![None, None],
tasks,
})
}
pub fn prev_mode(&mut self) -> Option<Mode> {
let new_mode = match self.cur_mode {
Mode::Home => Some(Mode::Home),
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),
// Disallowed moves
Mode::InstallDrivers
| Mode::ScanDisks
| Mode::PreClone
| Mode::Clone
| Mode::PostClone => None,
// Invalid states
Mode::PEMenu => panic!("This shouldn't happen?"),
};
new_mode
}
pub fn next_mode(&mut self) -> Option<Mode> {
let new_mode = match self.cur_mode {
Mode::Home => Mode::ScanDisks,
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,
// Invalid states
Mode::PEMenu => panic!("This shouldn't happen?"),
};
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, fortune()))?;
}
_ => {}
}
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();
action_tx.send(Action::SetMode(Mode::ScanDisks))?;
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) => {
// Clear TableType selection
match new_mode {
Mode::SelectDisks | Mode::SelectTableType => {
self.clone.table_type = None;
}
_ => {}
}
self.set_mode(new_mode)?;
self.action_tx
.send(Action::UpdateFooter(build_footer_string(self.cur_mode)))?;
self.action_tx.send(build_left_items(self, self.cur_mode))?;
self.action_tx
.send(build_right_items(self, self.cur_mode))?;
match new_mode {
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<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 build_footer_string(cur_mode: Mode) -> String {
match cur_mode {
Mode::Home | 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"),
// Invalid states
Mode::PEMenu => panic!("This shouldn't happen?"),
}
}
fn build_left_items(app: &App, cur_mode: Mode) -> Action {
let select_num: usize;
let title: String;
let mut items = Vec::new();
let mut labels: Vec<String> = Vec::new();
match cur_mode {
Mode::Home => {
select_num = 0;
title = String::from("Home");
}
Mode::InstallDrivers => {
select_num = 1;
title = String::from("Install Drivers");
app.clone
.driver_list
.iter()
.for_each(|driver| items.push(driver.to_string()));
}
Mode::SelectDisks => {
select_num = 2;
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 => {
select_num = 1;
title = String::from("Select Partition Table Type");
items.push(format!("{}", PartitionTableType::Guid));
items.push(format!("{}", PartitionTableType::Legacy));
}
Mode::Confirm => {
select_num = 0;
title = String::from("Confirm Selections");
}
Mode::ScanDisks | Mode::PreClone | Mode::Clone | Mode::PostClone => {
select_num = 0;
title = String::from("Processing");
}
Mode::SelectParts => {
select_num = 2;
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 => {
select_num = 0;
title = String::from("Done");
}
// Invalid states
Mode::PEMenu => panic!("This shouldn't happen?"),
};
Action::UpdateLeft(title, labels, items, select_num)
}
fn build_right_items(app: &App, cur_mode: Mode) -> Action {
let mut items = Vec::new();
let mut labels: Vec<Vec<DVLine>> = 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)));
}
}
}
_ => {}
}
Action::UpdateRight(labels, start_index, items)
}
fn fortune() -> String {
String::from(match random::<u8>() / 4 {
0 => "FUN FACT\n\n\nComputers barely work.",
1 => "CRASH OVERRIDE\n\n\n\"Hack the planet!\"",
2 => "CATS\n\n\n\"All your base are belong to us!\"",
3 => "HMM\n\n\nThis has all happened before...\n\nThis will all happen again.",
4 => "CYPHER\n\n\n\"I dont even see the code. All I see is blonde, brunette, red-head.\"",
5 => "CONGRATULATIONS\n\n\nYour did it!",
6 => "DID YOU KNOW?\n\n\nmacOS includes a built-in screen reader!",
7 => "TIP OF THE DAY\n\n\nNever go full Snappy!",
8 => "WORDS OF WISDOM\n\n\n\nIts not DNS,\n\nTheres no way its DNS,\n\nIt was DNS.",
9 => "HAL 9000\n\n\n\"I'm sorry Dave, I'm afraid I can't do that.\"",
_ => "COMPLETE\n\n\nThank you for using this tool!",
})
}