pyro/tetanus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added a remove project function, and layed the framework for prompting

Browse Source 
for user interaction.
This commit is contained in:
pyro
2026-05-21 13:15:44 -05:00
parent 2bdf100ed6
commit 9a43f4c378
2 changed files with 379 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/funcs.rs
+181 -88
View File
@@ -31,7 +31,7 @@ pub fn run_tui(
cursor::Hide, cursor::Hide,
EnableMouseCapture EnableMouseCapture
)?; )?;
let backend = CrosstermBackend::new(stdout); let backend = CrosstermBackend::new(&stdout);
let mut terminal = Terminal::new(backend)?; let mut terminal = Terminal::new(backend)?;
let (event_tx, event_rx) = channel::<AppEvent>(); let (event_tx, event_rx) = channel::<AppEvent>();
let input_tx = event_tx.clone(); let input_tx = event_tx.clone();
@@ -142,114 +142,207 @@ pub fn run_tui(
state.output.push(txt); state.output.push(txt);
state.output_scroll = u16::MAX; 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 { AppEvent::Key(key) => {
KeyCode::Esc => break, match key.code {
KeyCode::Enter => { KeyCode::Esc => break,
let trimmed = state.curent_intput.trim().to_string(); KeyCode::Enter => {
if !trimmed.is_empty() { let trimmed = state.curent_intput.trim().to_string();
state.history.push(trimmed.clone()); if !trimmed.is_empty() {
state.output.push("\n".to_string()); state.history.push(trimmed.clone());
state.output.push(format!("[user input] > {}", trimmed)); state.output.push("\n".to_string());
state.output.push("\n".to_string()); state.output.push(format!("[user input] > {}", trimmed));
let (command, args) = trimmed.split_once(' ').unwrap_or((&trimmed, "")); state.output.push("\n".to_string());
match command { if let Some(action) = state.prompt.action.clone() {
"exit" | "quit" => break, match action {
"reload-modules" => { ToolMessage::RemoveProject => {
state.initialize_modules(); if trimmed.to_lowercase().contains("y") {
state.output.push("Reloading module paths...".into()); state.execute_command(
} "remove_project_confirm",
"help" => { None,
state.output.push("Available Modules:".into()); )?;
for name in state.module_loader.commands.keys() { }
state.output.push(format!(" - {}", name)); }
_ => {}
} }
} }
"new_project" | "np" => { let (command, args) =
if args.split_once(' ').is_some() { trimmed.split_once(' ').unwrap_or((&trimmed, ""));
let _ = match command {
state.execute_command(command, Some(args.to_string())); "exit" | "quit" => break,
} else { "reload-modules" => {
state.output.push("Error: USAGE -> np <org> <name>".into()); state.initialize_modules();
state.output.push("Reloading module paths...".into());
} }
} "help" => {
"current_project" | "cp" => { state.output.push("Available Modules:".into());
let _ = state.execute_command(command, None); for name in state.module_loader.commands.keys() {
} state.output.push(format!(" - {}", name));
command_name => { }
if state.module_loader.commands.contains_key(command_name) { }
state.output.push(format!( "new_project" | "np" => {
"[Worker] Executing script '{}'...", if args.split_once(' ').is_some() {
command_name let _ = state
)); .execute_command(command, Some(args.to_string()));
if let Err(e) = state.execute_command(command_name, None) { } else {
state state
.output .output
.push(format!("[Error] Pipeline fail: {}", e)); .push("Error: USAGE -> np <org> <name>".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::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::Backspace => {
KeyCode::Down if key.modifiers.is_empty() => { state.curent_intput.pop();
if history_index < state.history.len() { }
history_index += 1; KeyCode::Up if key.modifiers.is_empty() => {
if history_index == state.history.len() { if !state.history.is_empty() && history_index > 0 {
state.curent_intput.clear(); // Clear back to a fresh prompt history_index -= 1;
} else {
state.curent_intput = state.history[history_index].clone(); state.curent_intput = state.history[history_index].clone();
} }
} }
} KeyCode::Down if key.modifiers.is_empty() => {
KeyCode::Up if history_index < state.history.len() {
if key history_index += 1;
.modifiers if history_index == state.history.len() {
.contains(crossterm::event::KeyModifiers::CONTROL) => state.curent_intput.clear(); // Clear back to a fresh prompt
{ } else {
if let Some(selected) = project_list_state.selected() { state.curent_intput = state.history[history_index].clone();
if selected > 0 { }
let new_index = selected - 1;
project_list_state.select(Some(new_index));
state.selected_project = new_index;
} }
} }
} KeyCode::Up
KeyCode::Down if key
if key .modifiers
.modifiers .contains(crossterm::event::KeyModifiers::CONTROL) =>
.contains(crossterm::event::KeyModifiers::CONTROL) => {
{ if let Some(selected) = project_list_state.selected() {
if let Some(selected) = project_list_state.selected() { if selected > 0 {
if selected + 1 < state.projects.len() { let new_index = selected - 1;
let new_index = selected + 1; project_list_state.select(Some(new_index));
project_list_state.select(Some(new_index)); state.selected_project = 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) => { AppEvent::Mouse(mouse) => {
if mouse.kind == crossterm::event::MouseEventKind::ScrollUp { if mouse.kind == crossterm::event::MouseEventKind::ScrollUp {
state.output_scroll = state.output_scroll.saturating_sub(1); state.output_scroll = state.output_scroll.saturating_sub(1);
src/lib.rs
+198 -33
View File
@@ -3,10 +3,10 @@ use ratatui::crossterm::event;
use rhai::{AST, Dynamic, Engine, Scope}; use rhai::{AST, Dynamic, Engine, Scope};
use std::collections::HashMap; use std::collections::HashMap;
use std::error::Error; use std::error::Error;
use std::fs::{File, create_dir_all, read_dir, read_to_string, remove_dir_all}; use std::fs::{self, File, create_dir_all, read_dir, read_to_string, remove_dir_all};
use std::io::{self, Write}; use std::io::{self, BufRead, BufReader, Write};
use std::path::PathBuf; use std::path::PathBuf;
use std::process::Command; use std::process::{Command, Stdio};
use std::sync::mpsc::Receiver; use std::sync::mpsc::Receiver;
use std::sync::{Arc, mpsc::Sender, mpsc::channel}; use std::sync::{Arc, mpsc::Sender, mpsc::channel};
@@ -28,6 +28,7 @@ pub struct AppState {
pub curent_intput: String, pub curent_intput: String,
pub module_loader: ModuleLoader, pub module_loader: ModuleLoader,
pub output_scroll: u16, pub output_scroll: u16,
pub prompt: Prompt,
} }
impl AppState { impl AppState {
@@ -53,6 +54,11 @@ impl AppState {
curent_intput: String::new(), curent_intput: String::new(),
module_loader: ModuleLoader::new(), module_loader: ModuleLoader::new(),
output_scroll: 0, output_scroll: 0,
prompt: Prompt {
action: None,
responses: Vec::new(),
completed: false,
},
}, },
main_rx, main_rx,
) )
@@ -99,6 +105,24 @@ impl AppState {
let mut new_project = Project::new(); let mut new_project = Project::new();
new_project.config_folder(conf_file.path()); new_project.config_folder(conf_file.path());
new_project.load_config()?; 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); self.projects.push(new_project);
} }
} }
@@ -165,6 +189,10 @@ impl AppState {
"notes: {}", "notes: {}",
&self.projects[self.selected_project].notes.display() &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() { if let Some(db) = self.projects[self.selected_project].db.clone() {
out_vec.push(format!("distrobox: {}", db.name)); out_vec.push(format!("distrobox: {}", db.name));
} }
@@ -178,23 +206,60 @@ impl AppState {
} }
} }
"promote_project" | "pp" => { "promote_project" | "pp" => {
match self.promote_project() { self.promote_project()?;
Ok(out) => { }
let _ = self.main_tx.send(ToolMessage::Output(out)); "remove_project_confirm" => {
} let project = self.projects[self.selected_project].clone();
Err(e) => { let mut config_file = project.config_folder.clone();
let _ = self config_file.pop();
.main_tx if let Err(e) = remove_dir_all(&project.config_folder) {
.send(ToolMessage::Output(format!("error promoting project! {e}"))); 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 let ast = self
.module_loader .module_loader
.asts .asts
.get(command_name) .get(command_name)
.ok_or_else(|| format!("AST not found for command: {}", command_name))? .ok_or_else(|| format!("Command not found: {}", command_name))?
.clone(); .clone();
let engine = Arc::clone(&self.module_loader.engine); 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<dyn Error>> { pub fn new_project(&mut self, org: String, name: String) -> Result<(), Box<dyn Error>> {
let template_box = self.config.get("template_box").unwrap(); 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()) let project_files = PathBuf::from(self.config.get("upcoming_files").unwrap())
.join(format!("{}/{}", org, name)); .join(format!("{}/{}", org, name));
let project_notes = PathBuf::from(self.config.get("upcoming_notes").unwrap()) let project_notes = PathBuf::from(self.config.get("upcoming_notes").unwrap())
.join(format!("{}/{}", org, name)); .join(format!("{}/{}", org, name));
create_dir_all(&project_files)?; let mut template_path = self.config_file.clone();
create_dir_all(&project_notes)?; 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 { let db = DistroBox {
name: format!("{}-{}-{}", template_box, org, name), name: format!("{}-{}-{}", template_box, org, name),
volumes: vec![ volumes: vec![
project_files.display().to_string(), project_files.display().to_string(),
project_notes.display().to_string(), project_notes.display().to_string(),
tools_dir.clone(), tools.to_string(),
], ],
created: false, 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()?; new_project.save_config()?;
self.projects.push(new_project); self.projects.push(new_project);
return Ok(()); return Ok(());
@@ -324,6 +454,8 @@ impl AppState {
let new_notes_path = PathBuf::from(self.config.get("current_notes").unwrap()) let new_notes_path = PathBuf::from(self.config.get("current_notes").unwrap())
.join(format!("{}/{}", new_project.org_name, new_project.name)); .join(format!("{}/{}", new_project.org_name, new_project.name));
let mut options = CopyOptions::new(); let mut options = CopyOptions::new();
create_dir_all(&new_files_path)?;
create_dir_all(&new_notes_path)?;
options.overwrite = true; options.overwrite = true;
copy(new_project.files.clone(), new_files_path.clone(), &options)?; copy(new_project.files.clone(), new_files_path.clone(), &options)?;
copy(new_project.notes.clone(), new_notes_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(), new_notes_path.display().to_string(),
self.config.get("tools").unwrap().to_string(), self.config.get("tools").unwrap().to_string(),
]; ];
db.create()?;
new_project.db = Some(db); new_project.db = Some(db);
let _ = self.main_tx.send(ToolMessage::RebuildDB);
} }
self.projects[self.selected_project] = new_project; self.projects[self.selected_project] = new_project;
return Ok(return_string); return Ok(return_string);
@@ -563,7 +695,9 @@ pub enum ToolMessage {
Input(String), Input(String),
Output(String), Output(String),
UpdateHost(Host), UpdateHost(Host),
UpdateProject(Project), RebuildDB,
UpdateProject(usize, Project),
RemoveProject,
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
@@ -747,7 +881,7 @@ pub struct DistroBox {
} }
impl DistroBox { impl DistroBox {
pub fn create(&mut self) -> Result<(), Box<dyn Error>> { pub fn create(&mut self, tx: Sender<ToolMessage>) -> Result<(), Box<dyn Error>> {
let mut create_command = Command::new("distrobox"); let mut create_command = Command::new("distrobox");
create_command create_command
.arg("create") .arg("create")
@@ -762,8 +896,39 @@ impl DistroBox {
.arg("--volume") .arg("--volume")
.arg(format!("{}:/{}:rw", volume, folder_name)); .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; self.created = true;
return Ok(()); return Ok(());
} }
} }
struct Prompt {
action: Option<ToolMessage>,
responses: Vec<String>,
completed: bool,
}