From 9a43f4c37848fefb88f5e90c603fb0cb1838f75a Mon Sep 17 00:00:00 2001 From: pyro Date: Thu, 21 May 2026 13:15:44 -0500 Subject: [PATCH] added a remove project function, and layed the framework for prompting for user interaction. --- src/funcs.rs | 269 ++++++++++++++++++++++++++++++++++----------------- src/lib.rs | 231 ++++++++++++++++++++++++++++++++++++------- 2 files changed, 379 insertions(+), 121 deletions(-) diff --git a/src/funcs.rs b/src/funcs.rs index ce86d6e..c138f41 100644 --- a/src/funcs.rs +++ b/src/funcs.rs @@ -31,7 +31,7 @@ pub fn run_tui( cursor::Hide, EnableMouseCapture )?; - let backend = CrosstermBackend::new(stdout); + let backend = CrosstermBackend::new(&stdout); let mut terminal = Terminal::new(backend)?; let (event_tx, event_rx) = channel::(); let input_tx = event_tx.clone(); @@ -142,114 +142,207 @@ pub fn run_tui( state.output.push(txt); state.output_scroll = u16::MAX; } + ToolMessage::UpdateProject(index, project) => { + if index < state.projects.len() { + state.projects[index] = project; + } + } + ToolMessage::RebuildDB => { + let project = state.projects[state.selected_project].clone(); + if let Some(db) = project.db { + disable_raw_mode()?; + execute!( + &stdout, + LeaveAlternateScreen, + cursor::Show, + DisableMouseCapture + )?; + std::io::stdout().flush()?; + let mut status = Command::new("distrobox"); + let template_box = state.config.get("template_box").unwrap(); + status + .arg("create") + .arg("--root") + .arg("--clone") + .arg(&template_box) + .arg("--name") + .arg(format!( + "{}-{}-{}", + template_box, project.org_name, project.name + )); + for volume in db.volumes { + let mut dir_name = volume.split("/").last().unwrap(); + if volume.contains("files") { + dir_name = "/pentest"; + } + if volume.contains("/notes") { + dir_name = "notes"; + } + status + .arg("--volume") + .arg(format!("{}:{}:rw", volume, dir_name)); + } + status + .stdin(Stdio::inherit()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()); + match status.status() { + Ok(status) if !status.success() => { + println!( + "\nCommand failed with exit code {status}, press enter to return to tetanus." + ); + let mut discard = String::new(); + let _ = std::io::stdin().read_line(&mut discard); + } + Err(e) => { + println!( + "\nFailed to execure distrobox: {e}, press enter to return to tetanus." + ); + let mut discard = String::new(); + let _ = std::io::stdin().read_line(&mut discard); + } + _ => {} + } + enable_raw_mode()?; + execute!( + &stdout, + EnterAlternateScreen, + DisableMouseCapture, + cursor::Hide + )?; + terminal.clear()?; + } + } + _ => {} }, - AppEvent::Key(key) => match key.code { - KeyCode::Esc => break, - KeyCode::Enter => { - let trimmed = state.curent_intput.trim().to_string(); - if !trimmed.is_empty() { - state.history.push(trimmed.clone()); - state.output.push("\n".to_string()); - state.output.push(format!("[user input] > {}", trimmed)); - state.output.push("\n".to_string()); - let (command, args) = trimmed.split_once(' ').unwrap_or((&trimmed, "")); - match command { - "exit" | "quit" => break, - "reload-modules" => { - state.initialize_modules(); - state.output.push("Reloading module paths...".into()); - } - "help" => { - state.output.push("Available Modules:".into()); - for name in state.module_loader.commands.keys() { - state.output.push(format!(" - {}", name)); + AppEvent::Key(key) => { + match key.code { + KeyCode::Esc => break, + KeyCode::Enter => { + let trimmed = state.curent_intput.trim().to_string(); + if !trimmed.is_empty() { + state.history.push(trimmed.clone()); + state.output.push("\n".to_string()); + state.output.push(format!("[user input] > {}", trimmed)); + state.output.push("\n".to_string()); + if let Some(action) = state.prompt.action.clone() { + match action { + ToolMessage::RemoveProject => { + if trimmed.to_lowercase().contains("y") { + state.execute_command( + "remove_project_confirm", + None, + )?; + } + } + _ => {} } } - "new_project" | "np" => { - if args.split_once(' ').is_some() { - let _ = - state.execute_command(command, Some(args.to_string())); - } else { - state.output.push("Error: USAGE -> np ".into()); + let (command, args) = + trimmed.split_once(' ').unwrap_or((&trimmed, "")); + match command { + "exit" | "quit" => break, + "reload-modules" => { + state.initialize_modules(); + state.output.push("Reloading module paths...".into()); } - } - "current_project" | "cp" => { - let _ = state.execute_command(command, None); - } - command_name => { - if state.module_loader.commands.contains_key(command_name) { - state.output.push(format!( - "[Worker] Executing script '{}'...", - command_name - )); - if let Err(e) = state.execute_command(command_name, None) { + "help" => { + state.output.push("Available Modules:".into()); + for name in state.module_loader.commands.keys() { + state.output.push(format!(" - {}", name)); + } + } + "new_project" | "np" => { + if args.split_once(' ').is_some() { + let _ = state + .execute_command(command, Some(args.to_string())); + } else { state .output - .push(format!("[Error] Pipeline fail: {}", e)); + .push("Error: USAGE -> np ".into()); + } + } + command_name => { + if state.module_loader.commands.contains_key(command_name) { + state.output.push(format!( + "[Worker] Executing script '{}'...", + command_name + )); + if let Err(e) = + state.execute_command(command_name, None) + { + state + .output + .push(format!("[Error] Pipeline fail: {}", e)); + } + } else { + if args == "" { + let _ = state.execute_command(command, None); + } else { + let _ = state.execute_command( + command, + Some(args.to_string()), + ); + } } - } else { - state.output.push(format!( - "[Error] Command '{}' unknown.", - command_name - )); } } + state.curent_intput.clear(); } - state.curent_intput.clear(); + history_index = state.history.len(); } - history_index = state.history.len(); - } - KeyCode::Char(c) => { - state.curent_intput.push(c); - } - KeyCode::Backspace => { - state.curent_intput.pop(); - } - KeyCode::Up if key.modifiers.is_empty() => { - if !state.history.is_empty() && history_index > 0 { - history_index -= 1; - state.curent_intput = state.history[history_index].clone(); + KeyCode::Char(c) => { + state.curent_intput.push(c); } - } - KeyCode::Down if key.modifiers.is_empty() => { - if history_index < state.history.len() { - history_index += 1; - if history_index == state.history.len() { - state.curent_intput.clear(); // Clear back to a fresh prompt - } else { + KeyCode::Backspace => { + state.curent_intput.pop(); + } + KeyCode::Up if key.modifiers.is_empty() => { + if !state.history.is_empty() && history_index > 0 { + history_index -= 1; state.curent_intput = state.history[history_index].clone(); } } - } - KeyCode::Up - if key - .modifiers - .contains(crossterm::event::KeyModifiers::CONTROL) => - { - if let Some(selected) = project_list_state.selected() { - if selected > 0 { - let new_index = selected - 1; - project_list_state.select(Some(new_index)); - state.selected_project = new_index; + KeyCode::Down if key.modifiers.is_empty() => { + if history_index < state.history.len() { + history_index += 1; + if history_index == state.history.len() { + state.curent_intput.clear(); // Clear back to a fresh prompt + } else { + state.curent_intput = state.history[history_index].clone(); + } } } - } - KeyCode::Down - if key - .modifiers - .contains(crossterm::event::KeyModifiers::CONTROL) => - { - if let Some(selected) = project_list_state.selected() { - if selected + 1 < state.projects.len() { - let new_index = selected + 1; - project_list_state.select(Some(new_index)); - state.selected_project = new_index; + KeyCode::Up + if key + .modifiers + .contains(crossterm::event::KeyModifiers::CONTROL) => + { + if let Some(selected) = project_list_state.selected() { + if selected > 0 { + let new_index = selected - 1; + project_list_state.select(Some(new_index)); + state.selected_project = new_index; + } } } + KeyCode::Down + if key + .modifiers + .contains(crossterm::event::KeyModifiers::CONTROL) => + { + if let Some(selected) = project_list_state.selected() { + if selected + 1 < state.projects.len() { + let new_index = selected + 1; + project_list_state.select(Some(new_index)); + state.selected_project = new_index; + } + } + } + _ => {} } - _ => {} - }, + } AppEvent::Mouse(mouse) => { if mouse.kind == crossterm::event::MouseEventKind::ScrollUp { state.output_scroll = state.output_scroll.saturating_sub(1); diff --git a/src/lib.rs b/src/lib.rs index 53c6f8b..8d926b6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,10 +3,10 @@ use ratatui::crossterm::event; use rhai::{AST, Dynamic, Engine, Scope}; use std::collections::HashMap; use std::error::Error; -use std::fs::{File, create_dir_all, read_dir, read_to_string, remove_dir_all}; -use std::io::{self, Write}; +use std::fs::{self, File, create_dir_all, read_dir, read_to_string, remove_dir_all}; +use std::io::{self, BufRead, BufReader, Write}; use std::path::PathBuf; -use std::process::Command; +use std::process::{Command, Stdio}; use std::sync::mpsc::Receiver; use std::sync::{Arc, mpsc::Sender, mpsc::channel}; @@ -28,6 +28,7 @@ pub struct AppState { pub curent_intput: String, pub module_loader: ModuleLoader, pub output_scroll: u16, + pub prompt: Prompt, } impl AppState { @@ -53,6 +54,11 @@ impl AppState { curent_intput: String::new(), module_loader: ModuleLoader::new(), output_scroll: 0, + prompt: Prompt { + action: None, + responses: Vec::new(), + completed: false, + }, }, main_rx, ) @@ -99,6 +105,24 @@ impl AppState { let mut new_project = Project::new(); new_project.config_folder(conf_file.path()); new_project.load_config()?; + let template_box = self.config.get("template_box").unwrap(); + let tools = self.config.get("tools").unwrap(); + let db = DistroBox { + name: format!( + "{}-{}-{}", + template_box, + new_project.org_name.clone(), + new_project.name.clone() + ), + volumes: vec![ + new_project.files.display().to_string(), + new_project.notes.display().to_string(), + tools.clone(), + ], + created: false, + template: template_box.clone(), + }; + new_project.db = Some(db); self.projects.push(new_project); } } @@ -165,6 +189,10 @@ impl AppState { "notes: {}", &self.projects[self.selected_project].notes.display() )); + out_vec.push(format!( + "config_folder: {}", + &self.projects[self.selected_project].config_folder.display() + )); if let Some(db) = self.projects[self.selected_project].db.clone() { out_vec.push(format!("distrobox: {}", db.name)); } @@ -178,23 +206,60 @@ impl AppState { } } "promote_project" | "pp" => { - match self.promote_project() { - Ok(out) => { - let _ = self.main_tx.send(ToolMessage::Output(out)); - } - Err(e) => { - let _ = self - .main_tx - .send(ToolMessage::Output(format!("error promoting project! {e}"))); - } - }; + self.promote_project()?; + } + "remove_project_confirm" => { + let project = self.projects[self.selected_project].clone(); + let mut config_file = project.config_folder.clone(); + config_file.pop(); + if let Err(e) = remove_dir_all(&project.config_folder) { + let _ = self.main_tx.send(ToolMessage::Output(format!( + "failed to delete config: {e} on {}", + &project.config_folder.display() + ))); + } + if let Err(e) = remove_dir_all(&project.notes) { + let _ = self.main_tx.send(ToolMessage::Output(format!( + "failed to delete notes: {e} on {}", + &project.notes.display() + ))); + } + if let Err(e) = remove_dir_all(&project.files) { + let _ = self.main_tx.send(ToolMessage::Output(format!( + "failed to delete files: {e} on {}", + &project.files.display() + ))); + } + self.projects.remove(self.selected_project); + self.prompt.action = None; + self.prompt.responses.clear(); + let _ = self.main_tx.send(ToolMessage::Output(format!( + "{}-{} was sucessfully removed!", + project.org_name, project.name + ))); + } + "remove_project" | "rp" => { + let project = self.projects[self.selected_project].clone(); + if project.name == "default" && project.org_name == "default" { + let _ = self.main_tx.send(ToolMessage::Output( + "The default project must remain. Canceling.".to_string(), + )); + self.prompt.action = None; + } else { + self.prompt.action = Some(ToolMessage::RemoveProject); + let _ = self.main_tx.send(ToolMessage::Output(format!( + "{}, {} and all contents will be deleted. Continue? (y/N)", + project.files.display(), + project.notes.display() + ))); + } } _ => { let ast = self .module_loader .asts .get(command_name) - .ok_or_else(|| format!("AST not found for command: {}", command_name))? + .ok_or_else(|| format!("Command not found: {}", command_name))? .clone(); let engine = Arc::clone(&self.module_loader.engine); @@ -284,33 +349,98 @@ impl AppState { pub fn new_project(&mut self, org: String, name: String) -> Result<(), Box> { let template_box = self.config.get("template_box").unwrap(); - let tools_dir = self.config.get("tools").unwrap(); let project_files = PathBuf::from(self.config.get("upcoming_files").unwrap()) .join(format!("{}/{}", org, name)); let project_notes = PathBuf::from(self.config.get("upcoming_notes").unwrap()) .join(format!("{}/{}", org, name)); - create_dir_all(&project_files)?; - create_dir_all(&project_notes)?; + let mut template_path = self.config_file.clone(); + let mut project_conf_folder = self.config_file.clone(); + let tools = self.config.get("tools").unwrap(); + template_path.pop(); + template_path.push("note_templates"); + project_conf_folder.pop(); + project_conf_folder.push(format!("projects/{}-{}", org, name)); + let mut options = CopyOptions::new(); + options.overwrite = true; + match create_dir_all(&project_files) { + Ok(_) => {} + Err(e) => { + let _ = self.main_tx.send(ToolMessage::Output(format!( + "error making project files! {e} on {}", + project_files.display() + ))); + } + } + match create_dir_all(&project_notes) { + Ok(_) => {} + Err(e) => { + let _ = self.main_tx.send(ToolMessage::Output(format!( + "error making project notes! {e} on {}", + project_notes.display() + ))); + } + } + match create_dir_all(&project_conf_folder) { + Ok(_) => {} + Err(e) => { + let _ = self.main_tx.send(ToolMessage::Output(format!( + "error making config folder! {e} on {}", + project_conf_folder.display() + ))); + } + } + for entry in read_dir(template_path)? { + let entry = entry?; + let template_name_os = entry.file_name(); + let template_name = template_name_os.to_string_lossy(); + if name.clone().contains(template_name.as_ref()) { + for entry in read_dir(entry.path())? { + let entry = entry?; + let file_name_os = entry.file_name(); + let file_name = file_name_os.to_string_lossy(); + if entry.file_type()?.is_dir() { + let dest = project_notes.join(file_name.as_ref()); + match copy(entry.path(), dest, &options) { + Ok(_) => {} + Err(e) => { + let _ = self.main_tx.send(ToolMessage::Output(format!( + "error copying note template! {e}" + ))); + } + } + } else { + match fs::copy(entry.path(), project_notes.join(file_name.as_ref())) { + Ok(_) => {} + Err(e) => { + let _ = self.main_tx.send(ToolMessage::Output(format!( + "error copying note template! {e}" + ))); + } + } + } + } + } + } let db = DistroBox { name: format!("{}-{}-{}", template_box, org, name), volumes: vec![ project_files.display().to_string(), project_notes.display().to_string(), - tools_dir.clone(), + tools.to_string(), ], created: false, - template: template_box.clone(), + template: template_box.to_string(), + }; + let mut new_project = Project { + org_name: org, + name: name, + notes: project_notes, + files: project_files, + hosts: Vec::new(), + config_folder: project_conf_folder, + current: false, + db: Some(db), }; - let mut new_project = Project::new(); - let mut project_conf_folder = self.config_file.clone(); - project_conf_folder.pop(); - project_conf_folder.push(format!("projects/{}-{}", org, name)); - new_project.config_folder(project_conf_folder); - new_project.name = name; - new_project.db = Some(db); - new_project.org_name(org); - new_project.files(project_files); - new_project.notes(project_notes); new_project.save_config()?; self.projects.push(new_project); return Ok(()); @@ -324,6 +454,8 @@ impl AppState { let new_notes_path = PathBuf::from(self.config.get("current_notes").unwrap()) .join(format!("{}/{}", new_project.org_name, new_project.name)); let mut options = CopyOptions::new(); + create_dir_all(&new_files_path)?; + create_dir_all(&new_notes_path)?; options.overwrite = true; copy(new_project.files.clone(), new_files_path.clone(), &options)?; copy(new_project.notes.clone(), new_notes_path.clone(), &options)?; @@ -347,8 +479,8 @@ impl AppState { new_notes_path.display().to_string(), self.config.get("tools").unwrap().to_string(), ]; - db.create()?; new_project.db = Some(db); + let _ = self.main_tx.send(ToolMessage::RebuildDB); } self.projects[self.selected_project] = new_project; return Ok(return_string); @@ -563,7 +695,9 @@ pub enum ToolMessage { Input(String), Output(String), UpdateHost(Host), - UpdateProject(Project), + RebuildDB, + UpdateProject(usize, Project), + RemoveProject, } #[derive(Clone, Debug)] @@ -747,7 +881,7 @@ pub struct DistroBox { } impl DistroBox { - pub fn create(&mut self) -> Result<(), Box> { + pub fn create(&mut self, tx: Sender) -> Result<(), Box> { let mut create_command = Command::new("distrobox"); create_command .arg("create") @@ -762,8 +896,39 @@ impl DistroBox { .arg("--volume") .arg(format!("{}:/{}:rw", volume, folder_name)); } - create_command.status()?; + create_command.stdout(Stdio::piped()); + create_command.stderr(Stdio::piped()); + let mut child = create_command.spawn()?; + if let Some(stdout) = child.stdout.take() { + let tx_out = tx.clone(); + std::thread::spawn(move || { + let reader = BufReader::new(stdout); + for line in reader.lines().flatten() { + let _ = tx_out.send(ToolMessage::Output(line)); + } + }); + } + if let Some(stderr) = child.stderr.take() { + let tx_out = tx.clone(); + std::thread::spawn(move || { + let reader = BufReader::new(stderr); + for line in reader.lines().flatten() { + let _ = tx_out.send(ToolMessage::Output(line)); + } + }); + } + let status = child.wait()?; + if !status.success() { + let _ = tx.send(ToolMessage::Output(format!("error creating distrobox!"))); + } + self.created = true; return Ok(()); } } + +struct Prompt { + action: Option, + responses: Vec, + completed: bool, +}