deja-vu/boot_diags/src/app.rs
2Shirt 72130109cb
Refactor tasks (again)
Avoids copying/recreating the TaskType enum (previously the Task enum).
2025-02-11 21:36:00 -08:00

763 lines
27 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 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::{self, SafeMode},
cpu::get_cpu_name,
drivers,
},
tasks::{TaskType, 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>,
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<Mode>,
boot_modes: Vec<SafeMode>,
selections: Vec<Option<usize>>,
system32: String,
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());
let mut list = StatefulList::default();
list.set_items(vec![
Mode::BootDiags,
Mode::BootSetup,
Mode::InjectDrivers,
Mode::SetBootMode,
]);
Ok(Self {
// TUI
action_rx,
action_tx,
components: vec![
Box::new(Title::new("Boot Diagnostics")),
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::Home,
list,
boot_modes: vec![SafeMode::Enable, SafeMode::Disable],
system32: String::new(),
selections: vec![None, None],
tasks,
})
}
pub fn inject_driver(&mut self, index: usize) {
if let Some(driver) = self.clone.driver_list.get(index) {
if let Some(disk_index) = self.clone.disk_index_dest {
let disk_list = self.clone.disk_list.lock().unwrap();
if let Some(disk) = disk_list.get(disk_index) {
if let Some(boot_index) = self.clone.part_index_boot {
if let Ok(task) = boot::inject_driver(
driver,
disk.get_part_letter(boot_index).as_str(),
&self.system32,
) {
self.tasks.add(task);
}
}
}
}
}
}
pub fn next_mode(&mut self) -> Mode {
match self.cur_mode {
Mode::Home => Mode::ScanDisks,
Mode::InstallDrivers => Mode::ScanDisks,
Mode::ScanDisks => Mode::SelectDisks,
Mode::SelectDisks => Mode::SelectParts,
Mode::SelectParts => Mode::DiagMenu,
Mode::BootDiags | Mode::BootSetup => Mode::DiagMenu, // TODO: add Mode::ProgressReport
Mode::InjectDrivers | Mode::SetBootMode => Mode::DiagMenu,
Mode::Done => Mode::DiagMenu,
Mode::Failed => Mode::Failed,
// Default to current mode
_ => self.cur_mode,
}
}
pub fn set_boot_mode(&mut self, boot_mode: SafeMode) {
let new_mode = match boot_mode {
SafeMode::Disable => "Normal",
SafeMode::Enable => "Safe Mode (minimal)",
};
info!("Setting boot mode to: {new_mode}");
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) {
if let Some(boot_index) = self.clone.part_index_boot {
if let Ok(task) = boot::set_mode(
disk.get_part_letter(boot_index).as_str(),
boot_mode,
&self.system32,
&disk.part_type,
) {
self.tasks.add(task);
};
}
}
}
}
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::DiagMenu => {
self.selections[0] = None;
self.selections[1] = None;
self.list.select_first_item();
}
Mode::InjectDrivers | Mode::InstallDrivers => self.clone.scan_drivers(),
Mode::ScanDisks => {
if self.tasks.idle() {
self.tasks.add(TaskType::ScanDisks);
}
self.action_tx.send(Action::DisplayPopup(
popup::Type::Info,
String::from("Scanning Disks..."),
))?;
}
_ => {}
}
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();
// Late init
self.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(".")
};
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(..);
// Check background task(s)
match self.cur_mode {
Mode::BootDiags => {
// Check result of background task (if finished)
if let Some(task) = self.tasks.poll()? {
// TODO: Impl logic
}
}
Mode::ScanDisks => {
// Check background task (Once all are complete Action::NextScreen is sent)
self.tasks.poll()?;
}
_ => {}
}
}
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::NextScreen => {
let next_mode = self.next_mode();
self.cur_mode = next_mode;
self.action_tx.send(Action::DismissPopup)?;
self.action_tx.send(Action::SetMode(next_mode))?;
}
Action::Process => match self.cur_mode {
Mode::DiagMenu => {
// Use highlighted entry
if let Some(new_mode) = self.list.get_selected() {
self.action_tx.send(Action::SetMode(new_mode))?;
}
}
Mode::Done => {
self.action_tx.send(Action::NextScreen)?;
}
Mode::BootDiags | Mode::BootSetup => {
let new_mode = self.next_mode();
self.action_tx.send(Action::SetMode(new_mode))?;
}
_ => {}
},
Action::Resize(w, h) => self.handle_resize(tui, w, h)?,
Action::Render => self.render(tui)?,
Action::ScanDisks => self.action_tx.send(Action::SetMode(Mode::ScanDisks))?,
Action::Select(one, two) => match self.cur_mode {
Mode::InjectDrivers => {
if let Some(index) = one {
self.inject_driver(index);
}
}
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_dest = one;
}
Mode::SelectParts => {
self.clone.part_index_boot = one;
self.clone.part_index_os = two;
}
Mode::SetBootMode => {
if let Some(index) = one {
if let Some(boot_mode) = self.boot_modes.get(index) {
self.set_boot_mode(boot_mode.to_owned());
}
}
}
_ => {}
},
Action::SetMode(new_mode) => {
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.action_tx.send(build_right_items(self))?;
self.action_tx.send(Action::Select(None, None))?;
}
_ => {}
}
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::BootDiags | Mode::BootSetup | Mode::Home | Mode::PEMenu | Mode::ScanDisks => {
String::from("(q) to quit")
}
Mode::InstallDrivers | Mode::InjectDrivers | Mode::SetBootMode => {
String::from("(Enter) to select / (q) to quit")
}
Mode::DiagMenu | Mode::SelectParts => {
String::from("(Enter) to select / (s) to start over / (q) to quit")
}
Mode::Done => String::from("(Enter) to continue / (q) to quit"),
Mode::SelectDisks => String::from(
"(Enter) to select / / (i) to install driver / (r) to rescan / (q) to quit",
),
Mode::Failed => String::from("(Enter) or (q) to quit"),
// Invalid states
Mode::Confirm | Mode::Clone | Mode::PreClone | Mode::PostClone | Mode::SelectTableType => {
panic!("This shouldn't happen?")
}
}
}
fn build_left_items(app: &App) -> Action {
let mut items = Vec::new();
let mut labels = vec![String::new(), String::new()];
let select_num: usize;
let title: String;
match app.cur_mode {
Mode::Home => {
select_num = 0;
title = String::from("Home");
}
Mode::DiagMenu => {
select_num = 0;
title = String::from("Troubleshooting");
app.list.items.iter().for_each(|mode| {
let (name, _) = get_mode_strings(*mode);
items.push(name);
});
}
Mode::InstallDrivers => {
select_num = 1;
title = String::from("Install Drivers");
app.clone
.driver_list
.iter()
.for_each(|driver| items.push(driver.to_string()));
}
Mode::InjectDrivers => {
select_num = 1;
title = String::from("Select Drivers");
app.clone
.driver_list
.iter()
.for_each(|driver| items.push(driver.to_string()));
}
Mode::SelectDisks => {
select_num = 1;
title = String::from("Select Disk");
let disk_list = app.clone.disk_list.lock().unwrap();
disk_list
.iter()
.for_each(|disk| items.push(disk.description.to_string()));
}
Mode::ScanDisks => {
select_num = 0;
title = String::from("Processing");
}
Mode::SelectParts => {
select_num = 2;
title = String::from("Select Boot and OS Partitions");
labels[0] = String::from("boot");
labels[1] = 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::BootDiags | Mode::BootSetup => {
select_num = 0;
let (new_title, _) = get_mode_strings(app.cur_mode);
title = new_title;
}
Mode::SetBootMode => {
select_num = 1;
title = String::from("Set Boot Mode");
app.boot_modes.iter().for_each(|entry| {
items.push(format!("{:?} Safe Mode", entry));
});
}
Mode::Done | Mode::Failed => {
select_num = 0;
title = String::from("Done");
}
// Invalid states
Mode::SelectTableType
| Mode::PEMenu
| Mode::Confirm
| Mode::PreClone
| Mode::Clone
| Mode::PostClone => {
panic!("This shouldn't happen?")
}
};
Action::UpdateLeft(title, labels, items, select_num)
}
fn build_right_items(app: &App) -> Action {
let mut items: Vec<Vec<DVLine>> = Vec::new();
let mut labels: Vec<Vec<DVLine>> = Vec::new();
let mut start_index = 0;
// TODO: DELETE THIS SECTION
start_index = 1;
items.push(vec![
DVLine {
line_parts: vec![format!("Mode: {:?}", app.cur_mode)],
line_colors: vec![Color::Reset],
},
DVLine {
line_parts: vec![format!(
"Parts: {:?} // {:?}",
app.clone.part_index_boot, app.clone.part_index_os
)],
line_colors: vec![Color::Reset],
},
DVLine {
line_parts: vec![format!(
"Selected: {:?} // {:?}",
app.list.selected(),
app.list.get_selected(),
)],
line_colors: vec![Color::Reset],
},
DVLine {
line_parts: vec![String::from("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")],
line_colors: vec![Color::Reset],
},
]);
// TODO: DELETE THIS SECTION
match app.cur_mode {
Mode::DiagMenu => {
let mut header_lines: Vec<DVLine> = Vec::new();
if let Some(index) = app.clone.disk_index_dest {
let disk_list = app.clone.disk_list.lock().unwrap();
if let Some(disk) = disk_list.get(index) {
let mut parts: Vec<usize> = Vec::new();
if let Some(index) = app.clone.part_index_boot {
parts.push(index);
}
if let Some(index) = app.clone.part_index_os {
parts.push(index);
}
// Disk Details
header_lines.append(&mut vec![
DVLine {
line_parts: vec![String::from("Disk")],
line_colors: vec![Color::Cyan],
},
DVLine {
line_parts: vec![String::new()],
line_colors: vec![Color::Reset],
},
]);
header_lines.append(&mut get_disk_description_right(&disk, Some(parts)));
// Add header
if !header_lines.is_empty() {
items[0].append(&mut header_lines);
// TODO: Replace line above with lines below
// items.push(header_lines);
// start_index = 1;
}
}
}
app.list.items.iter().for_each(|mode| {
let (name, description) = get_mode_strings(*mode);
items.push(vec![
DVLine {
line_parts: vec![name],
line_colors: vec![Color::Cyan],
},
DVLine {
line_parts: vec![String::new()],
line_colors: vec![Color::Reset],
},
DVLine {
line_parts: vec![description],
line_colors: vec![Color::Reset],
},
]);
});
}
Mode::InjectDrivers | 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 => {
let dest_dv_line = DVLine {
line_parts: vec![String::from("Disk")],
line_colors: vec![Color::Cyan],
};
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, None)));
}
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, None));
// Partition Details
disk.parts
.iter()
.for_each(|part| items.push(get_part_description(&part)));
}
}
}
Mode::SetBootMode => {
app.boot_modes.iter().for_each(|mode| {
match mode {
SafeMode::Disable => {
items.push(vec![DVLine {
line_parts: vec![String::from("Disable Safe Mode")],
line_colors: vec![Color::Reset],
}]);
}
SafeMode::Enable => {
items.push(vec![DVLine {
line_parts: vec![String::from("Enable Safe Mode (minimal)")],
line_colors: vec![Color::Reset],
}]);
}
};
});
}
_ => {}
}
Action::UpdateRight(labels, start_index, items)
}
fn get_mode_strings(mode: Mode) -> (String, String) {
match mode {
Mode::BootDiags => (
String::from("Boot Diagnostics"),
String::from("Check for common Windows boot issues"),
),
Mode::BootSetup => (
String::from("Boot Setup"),
String::from("Create or recreate boot files"),
),
Mode::InjectDrivers => (
String::from("Inject Drivers"),
String::from("Inject drivers into existing Windows environment"),
),
Mode::SetBootMode => (
String::from("Toggle Safe Mode"),
String::from("Enable or disable safe mode"),
),
_ => panic!("This shouldn't happen"),
}
}