use clipboard::ClipboardContext; use clipboard::ClipboardProvider; use fs_extra::dir::{CopyOptions, copy}; use ipnet::IpNet; use keyring::Entry; use ratatui::crossterm::event; use rhai::{AST, Dynamic, Engine, Scope}; use rustls::{ClientConfig, ClientConnection, RootCertStore, StreamOwned, pki_types::ServerName}; use std::collections::HashMap; use std::env; use std::error::Error; use std::fs::{self, File, create_dir_all, read_dir, read_to_string, remove_dir, remove_dir_all}; use std::io::{self, BufRead, BufReader, Read, Write}; use std::net::IpAddr; use std::net::TcpStream; use std::path::PathBuf; use std::process::{Command, Stdio}; use std::sync::mpsc::Receiver; use std::sync::{Arc, Mutex, mpsc::Sender, mpsc::channel}; use std::time::Duration; use std::time::{self, Instant}; use walkdir::WalkDir; use crate::rust_tools::rustwitness; pub mod funcs; pub mod rust_tools; pub mod server; pub mod victim; enum AppEvent { Key(event::KeyEvent), Worker(ToolMessage), Mouse(event::MouseEvent), } #[derive(Clone, Debug)] pub enum Destination { Server, Victim, Attacker, Control, } #[derive(Clone, Debug)] pub enum ToolMessage { Input((usize, String)), Output((usize, String)), UpdateHost(Host), RebuildDB, DestroyDB(DistroBox), StopDB(DistroBox), StopTemplate(DistroBox), UpdateProject(usize, Project), RemoveProject, AddServer, EndPrompt, ConnectServer, SelectServer, SendServer((usize, String)), ServerBrokerExit, AppStateExit, DisconnectServer, DisconnectAllServers, } #[derive(Clone, Debug)] pub enum ToolArg { Project(Project), Projects(Vec), Host(Host), Hosts(Vec), Config(HashMap), Path(PathBuf), } pub struct AppState { pub projects: Vec, pub servers: Vec>>, pub config: HashMap, pub config_file: PathBuf, pub workers: rayon::ThreadPool, pub worker_txes: Vec>, pub main_tx: Sender, pub history: Vec, pub log: Vec, pub output: Vec, pub selected_server: usize, pub selected_project: usize, pub curent_intput: String, pub module_loader: ModuleLoader, pub output_scroll: u16, pub output_follow: bool, prompt: Prompt, pub info_scroll: u16, pub server_broker_running: bool, pub app_state_running: bool, pub remoting: bool, pub help: Vec, pub name: String, } impl AppState { pub fn new() -> Self { let (main_tx, _) = channel(); Self { projects: Vec::new(), servers: Vec::new(), config: HashMap::new(), config_file: PathBuf::new(), workers: rayon::ThreadPoolBuilder::new() .num_threads(4) .build() .unwrap(), worker_txes: Vec::new(), main_tx, history: Vec::new(), log: Vec::new(), output: Vec::new(), selected_server: 0, selected_project: 0, curent_intput: String::new(), module_loader: ModuleLoader::new(), output_scroll: 0, output_follow: true, prompt: Prompt { action: None, responses: Vec::new(), execute_command: String::new(), num_responses: 0, prompts: Vec::new(), last_prompted: None, }, info_scroll: 0, server_broker_running: true, app_state_running: true, remoting: false, help: Vec::new(), name: String::new(), } } pub fn load_config(&mut self, file: PathBuf, display: bool) -> Result<(), Box> { self.config_file = file.clone(); let config_contents = read_to_string(file)?; for line in config_contents.lines() { let parts: Vec<&str> = line.split(": ").collect(); if parts.len() == 2 { match parts[0].trim() { "servers" => { for server_data in parts[1].trim().split(", ").collect::>() { if let Some((address, name)) = server_data.split_once("|") { let new_server = Server { address: address.trim().to_string(), connected: false, timer: Duration::from_secs(60), last_check: time::Instant::now(), message_que: Vec::new(), action_que: Vec::new(), client_id: 0, selected_client: 0, config: None, logged_in: false, password: String::new(), name: name.to_string(), cert_text: String::new(), }; self.servers.push(Arc::new(Mutex::new(new_server))); } } self.config .insert("servers".to_string(), parts[1].trim().to_string()); } _ => { if parts[0].len() > 1 { self.config .insert(parts[0].trim().to_string(), parts[1].trim().to_string()); } } } } } let mut projet_folder_path = self.config_file.clone(); projet_folder_path.pop(); projet_folder_path.push("projects"); let project_dir_reses = read_dir(projet_folder_path)?; for res in project_dir_reses { if let Ok(project_folder) = res { let project_folder_file_reses = read_dir(project_folder.path())?; for res in project_folder_file_reses { if let Ok(conf_file) = res { if conf_file.file_name().to_string_lossy() == "project.conf".to_string() { let mut new_project = Project::new(); new_project.config_folder(conf_file.path()); new_project.load_config(display)?; 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.help.push("help\nThat's MEEEEEEE\n".to_string()); self.help.push("new_project\nadd a project to the upcoming project pool and create default notes based on templates\nusage:\nnew_project organization_name, project_name\n".to_string()); self.help.push( "current_project\nshow information on the currently selected project\n".to_string(), ); self.help .push("promote_project\npromote a project from upcoming to current\n".to_string()); self.help.push( "remove_project\nremove a project, delete its folders, and destroy its distrobox\n" .to_string(), ); self.help.push("new_terminal\nopen a new terminal in the distrobox of the currently selected project\n".to_string()); self.help .push("new_server\ncreate a new server object\n".to_string()); self.help .push("list_servers\nlist all servers currently tracked by the tool\n".to_string()); self.help .push("connect_server\nconnect to a tracked server\n".to_string()); self.help .push("save_all\nsave all projects and settings\n".to_string()); self.help .push("save_projects\nsave just the projects\n".to_string()); self.help .push("test_server\ntest the server connection\n".to_string()); self.help.push( "add_scope\nadd a host, cidr range, or ip address to the current project's scope\n" .to_string(), ); self.help.push("parse_scope\nparse the projects general.md notes file for the scope copied from the workbook\n".to_string()); self.help.push("parse_portscan | pps\nParse a services.tsv file exported from cobalt stirke in your current project's files directory".to_string()); self.help .push("stop_db\nrestart this project's distrobox".to_string()); self.help.push("rustwitness\nRun a scan on the hosts in the urls.txt file in yoru project files directory that captures screenshots of those urls, optionally through a proxy.".to_string()); self.help.push("exit\nquit the tool\n".to_string()); self.initialize_modules(); return Ok(()); } pub fn initialize_modules(&mut self) { if let Some(base_path_str) = self.config.get("module_path") { let base_path = PathBuf::from(base_path_str); println!("Loading Tetanus modules from: {:?}", base_path); if let Err(e) = self.module_loader.load_all(&base_path) { eprintln!("Failed to load modules: {}", e); } else { for module in &self.module_loader.commands { self.help.push(format!( "{}\n{}\n", module.1.name.clone(), module.1.help.clone() )); } println!( "Successfully loaded {} modules.", self.module_loader.commands.len() ); } } else { eprintln!("Warning: 'module_path' is missing from AppState config map!"); } } pub fn execute_command( &mut self, command_name: &str, command_args: Option, rid: usize, ) -> Result<(), Box> { let tx = self.main_tx.clone(); match command_name { "remote" => { if let Some(server) = self.servers.get(self.selected_server) { if let Ok(mut server) = server.lock() { if server.selected_client != 0 { if let Some(args) = command_args.clone() { let args = args.trim(); if args.contains(" ") { let (cmd, iargs) = args.split_once(" ").unwrap(); server.message_que.push(format!("CMD|{} {}", cmd, iargs)); } else { server.message_que.push(format!("CMD|{}", args)); } let _ = self.main_tx.send(ToolMessage::Output(( 0, format!("Tasked {} to run {}", server.selected_client, args), ))); } } } } } "new_project" | "np" => { let args = command_args.unwrap(); let (org, name) = args.split_once(" ").unwrap(); if let Err(e) = self.new_project(org.to_string(), name.to_string(), rid) { tx.send(ToolMessage::Output(( rid, format!("[error] making {}-{}: {e}", org, name), )))?; } else { tx.send(ToolMessage::Output(( rid, format!("[success] {}-{} created!", org, name), )))?; } } "generate_victim" => { let mut success = false; if let Some(args) = command_args { if args.split(" ").into_iter().count() == 1 { let mut address = String::new(); let mut password = String::new(); let config_file = self.config_file.clone(); let files = self.projects[self.selected_project].files.clone(); let mut cert_text = String::new(); if let Ok(lock) = self.servers[self.selected_server].lock() { address = lock.address.clone(); password = lock.password.clone(); cert_text = lock.cert_text.clone(); } match victim::generate_code( address, password, config_file, files, args, self.main_tx.clone(), cert_text, ) { Ok(_) => { let _ = self.main_tx.send(ToolMessage::Output(( rid, "[Success] Victim code generated! Compilation thread spawned!" .to_string(), ))); success = true; } Err(e) => { let _ = self.main_tx.send(ToolMessage::Output(( 0, format!("[error] generating victim code: {e}"), ))); } } } } if !success { let _ = self.main_tx.send(ToolMessage::Output(( 0, format!("Usage: generate_victim (windows/linux)"), ))); } } "current_project" | "cp" => { let mut out_vec = Vec::new(); out_vec.push(format!( "org: {}", &self.projects[self.selected_project].org_name )); out_vec.push(format!( "name: {}", &self.projects[self.selected_project].name )); out_vec.push(format!( "files: {}", &self.projects[self.selected_project].files.display() )); out_vec.push(format!( "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)); } if self.projects[self.selected_project].current { out_vec.push(format!("status: current")); } else { out_vec.push(format!("status: upcoming")); } out_vec.push(format!( "Num hosts: {}", &self.projects[self.selected_project].hosts.iter().count() )); for line in out_vec { let _ = self.main_tx.send(ToolMessage::Output((rid, line))); } } "rebuild_distro_box" | "rdb" => { if let Some(db) = self.projects[self.selected_project].db.clone() { if db.created { let _ = self.main_tx.send(ToolMessage::DestroyDB(db.clone())); } let _ = self.main_tx.send(ToolMessage::RebuildDB); let _ = self.main_tx.send(ToolMessage::Output(( rid, "[success] Distorbox made!".to_string(), ))); } } "promote_project" | "pp" => { self.main_tx.send(ToolMessage::Output(( rid, String::from("promote_project command recieved!"), )))?; match self.promote_project(rid) { Ok(_) => { self.main_tx.send(ToolMessage::Output(( rid, String::from("[success] project promoted!"), )))?; } Err(e) => { self.main_tx.send(ToolMessage::Output(( rid, format!("[error] promoting project: {}", e), )))?; } } } "remove_project_confirm" => { if self.prompt.responses.len() > 0 { if self.prompt.responses[0] .response .to_lowercase() .contains("y") { 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(config_file) { let _ = self.main_tx.send(ToolMessage::Output(( rid, 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(( rid, format!( "failed to delete notes: {e} on {}", &project.notes.display() ), ))); } else if let Some(parent) = project.notes.parent() { if let Ok(mut entries) = read_dir(parent) { if entries.next().is_none() { if let Err(e) = remove_dir(parent) { let _ = self.main_tx.send(ToolMessage::Output((rid, format!("failed to delete the empty client notes folder: {e} on {}", parent.display())))); } } } } if let Err(e) = remove_dir_all(&project.files) { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!( "failed to delete files: {e} on {}", &project.files.display() ), ))); } else if let Some(parent) = project.files.parent() { if let Ok(mut entries) = read_dir(parent) { if entries.next().is_none() { if let Err(e) = remove_dir(parent) { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!( "failed to delete empty files parent: {e} on {}", parent.display() ), ))); } } } } let remove_id = self.selected_project.clone(); let project = self.projects[remove_id].clone(); if let Some(db) = project.db { self.main_tx.send(ToolMessage::DestroyDB(db))?; } self.selected_project -= 1; self.projects.remove(remove_id); self.prompt.responses.clear(); self.prompt.action = None; let _ = self.main_tx.send(ToolMessage::Output(( rid, format!( "{}-{} was sucessfully removed!", project.org_name, project.name ), ))); } } } "remove_project" | "rp" => { let project = self.projects[self.selected_project].clone(); self.prompt.action = Some(ToolMessage::RemoveProject); self.prompt.num_responses = 1; self.prompt.execute_command = String::from("remove_project_confirm"); self.prompt.prompts.push(PromptQuery { id: 0, query: format!( "{}, {} and all contentes wil lbe deleted. Contineue? (y/N", project.files.display(), project.notes.display() ), }); let _ = self .main_tx .send(ToolMessage::Input((rid, "prompt".to_string()))); } "new_terminal" | "nt" => { let project = self.projects[self.selected_project].clone(); if let Some(db) = project.db { if let Some(cmd) = self.config.get("term_cmd") { db.launch_terminal(cmd.to_string()); } } } "new_server" | "ns" => { self.prompt.num_responses = 2; self.prompt.action = Some(ToolMessage::AddServer); self.prompt.execute_command = String::from("add_server"); self.prompt.prompts.push(PromptQuery { id: 0, query: "What is the IP of the server?".to_string(), }); self.prompt.prompts.push(PromptQuery { id: 1, query: "What is the port for the server?".to_string(), }); let _ = self .main_tx .send(ToolMessage::Input((0, "prompt".to_string()))); } "add_server" => { let mut ip = String::new(); let mut port = String::new(); self.prompt.responses.iter().for_each(|res| { if res.query.id == 0 { ip = res.response.clone(); } else if res.query.id == 1 { port = res.response.clone(); } }); let new_server = Server { address: format!("{}:{}", ip, port), connected: false, timer: Duration::from_secs(5), last_check: time::Instant::now(), message_que: Vec::new(), action_que: Vec::new(), client_id: 0, selected_client: 0, config: None, logged_in: false, password: String::new(), name: String::new(), cert_text: String::new(), }; self.servers.push(Arc::new(Mutex::new(new_server))); self.prompt.reset(); } "connect_server" => { let _ = self.main_tx.send(ToolMessage::ConnectServer); } "prompt" => { if self.prompt.responses.len() == self.prompt.num_responses { let _ = self .main_tx .send(ToolMessage::Input((0, self.prompt.execute_command.clone()))); let _ = self.main_tx.send(ToolMessage::EndPrompt); } else { for query in &self.prompt.prompts { if !self .prompt .responses .iter() .any(|res| res.query.id == query.id) { self.prompt.last_prompted = Some(query.clone()); let _ = self .main_tx .send(ToolMessage::Output((rid, query.query.clone()))); break; } } } } "list_servers" => { self.servers.iter().enumerate().for_each(|(id, m)| { if let Ok(s) = m.lock() { if id == self.selected_server { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!( "Server ID:{} Name:{} Address:{} ClientID:{} Sleep:{} Last:{} Loggedin: {} - Currently Selected", id, s.name, s.address, s.client_id, s.timer.as_secs(), s.last_check.elapsed().as_secs(), s.logged_in ), ))); } else { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!( "Server ID:{} Name:{} Address:{} ClientID:{} Sleep:{} Last:{}, loggedin: {}", id, s.name, s.address, s.client_id, s.timer.as_secs(), s.last_check.elapsed().as_secs(), s.logged_in ), ))); } } }); } "remove_server" => { let _ = self.main_tx.send(ToolMessage::DisconnectServer); } "server_remove_confim" => { let args = command_args.unwrap(); let id = args.trim().parse::().unwrap(); self.servers.remove(id); if self.selected_server == id && self.selected_server >= 1 { let _ = self.main_tx.send(ToolMessage::Output(( 0, format!("[success] server {} removed!", self.selected_server), ))); self.selected_server = self.selected_server - 1; } } "select_server" => { self.prompt.action = Some(ToolMessage::SelectServer); self.prompt.num_responses = 1; self.prompt.prompts.push(PromptQuery { id: 0, query: String::from("Selection?"), }); self.prompt.execute_command = String::from("server_selected"); self.servers.iter().enumerate().for_each(|(id, servermut)| { if let Ok(server) = servermut.lock() { if id == self.selected_server { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!("{}:{} - Currently Selected", id, server.address.clone()), ))); } else { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!("{}: {}", id, server.address.clone()), ))); } } }); } "server_selected" => { if self.prompt.responses.len() == 1 { if let Ok(id) = self.prompt.responses[0].response.parse::() { self.selected_server = id; if let Ok(server) = self.servers[id].lock() { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!("{} Selected!", server.address), ))); } } else { let _ = self.main_tx.send(ToolMessage::Output(( rid, String::from("Invalid selection."), ))); } } else { let _ = self.main_tx.send(ToolMessage::Output(( rid, String::from("[error] selecting server: no prompt responses found!"), ))); } let _ = self.main_tx.send(ToolMessage::EndPrompt); } "test_server" => { if let Some(server) = self.servers.get(self.selected_server) { if let Ok(mut server) = server.lock() { server.message_que.push("TEST|NONE".to_string()); } else { let _ = self.main_tx.send(ToolMessage::Output(( rid, "[error] getting server lock!".to_string(), ))); } } else { self.output .push(String::from("[error] getting selected server!")); } } "server_sleep" => { if let Some(args) = command_args { if let Ok(seconds) = args.trim().parse::() { if let Some(server) = self.servers.get(self.selected_server) { if let Ok(mut server) = server.lock() { server.timer = Duration::from_secs(seconds); self.output .push(format!("Server timer set to {} seconds", seconds)); } else { self.output .push(String::from("[error] locking server object!")); } } else { self.output .push(String::from("[error] not a valid server selected!")); } } else { self.output .push(String::from("invalid second count provided!")); } } else { let _ = self.main_tx.send(ToolMessage::Output(( rid, "[error] no seconds provided! please use server_sleep number_of_seconds" .to_string(), ))); } } "save_all" => { self.save_all()?; } "save_projects" => { self.save_projects()?; } "list_clients" => { if let Some(server) = self.servers.get(self.selected_server) { if let Ok(mut server) = server.lock() { server.message_que.push("LIST_CLIENTS|NONE".to_string()); } else { let _ = self.main_tx.send(ToolMessage::Output(( rid, "[error] getting server lock!".to_string(), ))); } } else { self.output .push(String::from("[error] getting selected server!")); } } "select_client" => { if let Some(arg) = command_args { if let Some(server) = self.servers.get(self.selected_server) { if let Ok(mut server) = server.lock() { if let Ok(id) = arg.parse::() { server.message_que.push(format!("STOP_CONTROL")); server.message_que.push(format!("CONTROL|{}", id.clone())); server.selected_client = id; let _ = self.main_tx.send(ToolMessage::Output(( 0, format!("client {} selected!", id), ))); } else { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!("[error] parsing client id from arg: {}", arg), ))); } } } } else { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!("No client identified. Usage: select_client client_id"), ))); } } "deselect_client" => { if let Some(server) = self.servers.get(self.selected_server) { if let Ok(mut server) = server.lock() { server.message_que.push(format!("STOP_CONTROL")); server.selected_client = 0; let _ = self .main_tx .send(ToolMessage::Output((rid, format!("client deselected!")))); } } } "add_scope" => { if let Some(args) = command_args { args.split_whitespace().into_iter().for_each(|host| { self.projects[self.selected_project].add_scope(host.to_string()); }); } } "parse_scope" => { let mut general_path = self.projects[self.selected_project].notes.clone(); general_path.push("general.md"); let general_string = read_to_string(general_path)?; if let Some((_, scope_section)) = general_string.split_once("# Scope") { if let Some((scope_section, _)) = scope_section.split_once("#") { scope_section.lines().into_iter().for_each(|line| { if line.contains("|") { if !line.contains("Subnet") { if !line.contains("--") { let host = line.split("|").collect::>()[1].to_string(); self.projects[self.selected_project].add_scope(host); } } } }); } } } "show_config" | "sc" => { let _ = self .main_tx .send(ToolMessage::Output((rid, "Current Config:".to_string()))); self.config.iter().for_each(|setting| { let _ = self.main_tx.send(ToolMessage::Output(( 0, format!("{}:{}", setting.0, setting.1), ))); }); } "set" => { let mut ready = false; let mut key = String::new(); let mut value = String::new(); if let Some(args) = command_args.clone() { let args_vec: Vec<&str> = args.split(" ").collect(); if args_vec.len() == 2 { key = args_vec[0].to_string(); value = args_vec[1].to_string(); if self.config.contains_key(&key) { ready = true; } else { let _ = self.main_tx.send(ToolMessage::Output(( 0, "[error] Unknown setting! {key}, {value}".to_string(), ))); return Ok(()); } } } if ready { self.config.insert(key, value); self.save_config()?; let _ = self.main_tx.send(ToolMessage::Output(( 0, "[success]Setting changed!".to_string(), ))); } else { let _ = self.main_tx.send(ToolMessage::Output(( 0, "[error] malformed command!".to_string(), ))); let _ = self.main_tx.send(ToolMessage::Output(( 0, "Usag: set setting value".to_string(), ))); } } "list_modules" => { self.module_loader.commands.iter().for_each(|module| { let _ = self .main_tx .send(ToolMessage::Output((rid, format!("{}", module.0)))); }); } "server_login" => { if let Some(pass) = command_args { let pass = pass.trim().to_string(); if let Some(server) = self.servers.get(self.selected_server) { if let Ok(mut lock) = server.lock() { lock.login(pass)?; let _ = self.main_tx.send(ToolMessage::Output(( rid, "[success] Logged into Server!".to_string(), ))); } } } } "new_config" => { self.prompt.num_responses = 8; self.prompt.action = Some(ToolMessage::Input((0, "add_config".to_string()))); self.prompt.execute_command = String::from("add_config"); self.prompt.prompts.push(PromptQuery { id: 0, query: "name for the new config?".to_string(), }); self.prompt.prompts.push(PromptQuery { id: 1, query: "Path to save the config folder?".to_string(), }); self.prompt.prompts.push(PromptQuery { id: 2, query: "Path for upcoming project files (not notes)?".to_string(), }); self.prompt.prompts.push(PromptQuery { id: 3, query: "Path for current project files (not notes)?".to_string(), }); self.prompt.prompts.push(PromptQuery { id: 4, query: "Path for upcoming notes?".to_string(), }); self.prompt.prompts.push(PromptQuery { id: 5, query: "Path for current notes?".to_string(), }); self.prompt.prompts.push(PromptQuery { id: 6, query: "Path to tool directory?".to_string(), }); self.prompt.prompts.push(PromptQuery { id: 7, query: "Template box name?".to_string(), }); let _ = self .main_tx .send(ToolMessage::Input((0, "prompt".to_string()))); } "add_config" => { self.prompt.responses.iter().for_each(|res| { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!("{} | {}", res.query.query, res.response), ))); }); self.new_config()?; self.save_config()?; } "list_configs" => { if let Some(config_line) = self.config.get("alternate_configs") { config_line.split(",").into_iter().for_each(|entry| { if let Some((name, path)) = entry.split_once("|") { let _ = self .main_tx .send(ToolMessage::Output((rid, format!("{}: {}", name, path)))); } }); } } "collab_with" => { if let Some(args) = command_args { if let Ok(id) = args.trim().parse::() { if let Some(server) = self.servers.get(self.selected_server) { if let Ok(mut lock) = server.lock() { lock.action_que.push(format!("COLLAB|{}", id)); } } } } } "collab" => { if let Some(args) = command_args { if let Some(server) = self.servers.get(self.selected_server) { if let Ok(mut lock) = server.lock() { lock.action_que.push(format!("CMD|{}", args)); } } } } "parse_port_scan" | "pps" => { self.parse_cs_portscan()?; } "stop_db" => { if let Some(project) = self.projects.get_mut(self.selected_project) { if let Some(db) = project.db.as_mut() { db.stop(self.main_tx.clone(), rid)?; } } } "rustwitness" => { if let Some(project) = self.projects.get(self.selected_project) { if let Some(args) = command_args { if let Some((outfile, proxy)) = args.split_once(" ") { let mut proxy_string = None; if let Some((ip, port)) = proxy.split_once(" ") { proxy_string = Some(format!("socks5://{}:{}", ip, port)); } WalkDir::new(&project.files) .into_iter() .filter(|res| res.is_ok()) .for_each(|res| { if let Ok(entry) = res { if entry.file_name().to_string_lossy().contains("urls.txt") { let input = PathBuf::from(entry.path()); let mut output = project.files.clone(); output.push(outfile.trim()); let proxy_intput = proxy_string.clone(); let tx = self.main_tx.clone(); self.workers.spawn(move || { rustwitness(input, output, proxy_intput, tx) }); } } }); } } else { let _ = self.main_tx.send(ToolMessage::Output(( 0, "[error] incorrect command usage!".to_string(), ))); let _ = self.main_tx.send(ToolMessage::Output(( 0, "usage: rustwitness output_file_name optional_proxy_ip optional_proxy_port".to_string(), ))); } } } "exit" => { self.save_all()?; let _ = self.main_tx.send(ToolMessage::AppStateExit); let _ = self .main_tx .send(ToolMessage::Output((rid, String::from("Good bye!")))); } _ => { if self.module_loader.asts.get(command_name).is_none() { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!("[error] command not found: {}", command_name), ))); return Ok(()); } let ast = self .module_loader .asts .get(command_name) .ok_or_else(|| format!("[error] Command not found: {}", command_name))? .clone(); let engine = Arc::clone(&self.module_loader.engine); let cmd_meta = self .module_loader .commands .get(command_name) .ok_or_else(|| format!("Metadata not found for command: {}", command_name))? .clone(); let mut scope = Scope::new(); for arg_requirement in &cmd_meta.args { match arg_requirement.trim() { "project" => { scope.push("project", self.projects[self.selected_project].clone()); } "projects" => { let mut arr = rhai::Array::new(); for proj in &self.projects { arr.push(rhai::Dynamic::from(proj.clone())); } scope.push("projects", arr); } "host" => { scope.push("host", "todo"); } "hosts" => { scope.push("hosts", self.projects[self.selected_project].clone()); } "config" => { let mut map = rhai::Map::new(); for (k, v) in &self.config { map.insert(k.clone().into(), v.clone().into()); } scope.push("config", map); } "distrobox" => { if let Some(db) = self.projects[self.selected_project].db.clone() { scope.push("distrobox", db); } } _ if arg_requirement.starts_with("string") => { let input_val = command_args .clone() .unwrap_or_else(|| self.curent_intput.clone()); scope.push("input_string", input_val); } _ if arg_requirement.starts_with("bool") => { let mut input_val = false; if command_args .clone() .unwrap_or_else(|| self.curent_intput.clone()) .contains("true") { input_val = true; } scope.push("input_bool", input_val); } _ => {} } } let worker_tx = self.main_tx.clone(); self.workers.spawn(move || { match engine.eval_ast_with_scope::(&mut scope, &ast) { Ok(result) => { let result_str = if result.is_array() { if let Ok(arr) = result.into_array() { let lines: Vec = arr .into_iter() .map(|item| { item.into_string().unwrap_or_else(|_| "".into()) }) .collect(); lines.join("\n") } else { "failed to process array output".into() } } else if result.is_string() { let out_string = result .into_string() .unwrap_or_else(|_| "failed to parse string".into()); if cmd_meta.calls == "clipboard" { let mut ctx: ClipboardContext = ClipboardProvider::new().unwrap(); ctx.set_contents(out_string.clone()).unwrap(); } else if cmd_meta.calls != "none" { let _ = worker_tx.send(ToolMessage::Input((0, out_string.clone()))); } out_string } else { format!("{:?}", result) }; if result_str.contains("\n") { result_str.lines().into_iter().for_each(|line| { let _ = tx.send(ToolMessage::Output((rid, line.to_string()))); }); } else { let _ = tx.send(ToolMessage::Output((rid, result_str))); } } Err(err) => { let _ = tx.send(ToolMessage::Output(( rid, format!("[error] in script execution: {}", err), ))); } } }); } } Ok(()) } pub fn edit_config(&mut self, key: String, value: String) { if self.config.contains_key(&key) { self.config.insert(key.clone(), value.clone()); } let _ = self.main_tx.send(ToolMessage::Output(( 0, format!("Setting saved! {key}:{value}"), ))); } pub fn new_project( &mut self, org: String, name: String, rid: usize, ) -> Result<(), Box> { let mut main_config_file = self.config_file.clone(); main_config_file.pop(); let _ = self.main_tx.send(ToolMessage::Output(( rid, format!("{} is the main config folder.", main_config_file.display()), ))); let template_box = self.config.get("template_box").unwrap(); let project_files = PathBuf::from(self.config.get("upcoming_files").unwrap()) .join(format!("{}/{}", org, name)); let _ = self.main_tx.send(ToolMessage::Output(( rid, format!("{} is the project_files folder", project_files.display()), ))); let project_notes = PathBuf::from(self.config.get("upcoming_notes").unwrap()) .join(format!("{}/{}", org, name)); let _ = self.main_tx.send(ToolMessage::Output(( rid, format!("{} is the project_notes folder", project_notes.display()), ))); let mut template_path = main_config_file.clone(); let mut project_conf_folder = main_config_file.clone(); let tools = self.config.get("tools").unwrap(); template_path.push("note_templates"); let _ = self.main_tx.send(ToolMessage::Output(( rid, format!("{} is the note template folder", template_path.display()), ))); project_conf_folder.push(format!("projects/{}-{}", org, name)); let _ = self.main_tx.send(ToolMessage::Output(( rid, format!( "{} is the project config folder", project_conf_folder.display() ), ))); let mut options = CopyOptions::new(); options.overwrite = true; if !project_files.exists() { match create_dir_all(&project_files) { Ok(_) => {} Err(e) => { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!( "[error] making project files! {e} on {}", project_files.display() ), ))); } } } if !project_notes.exists() {} match create_dir_all(&project_notes) { Ok(_) => {} Err(e) => { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!( "[error] making project notes! {e} on {}", project_notes.display() ), ))); } } if !project_conf_folder.exists() { match create_dir_all(&project_conf_folder) { Ok(_) => {} Err(e) => { let _ = self.main_tx.send(ToolMessage::Output(( rid, format!( "[error] making config folder! {e} on {}", project_conf_folder.display() ), ))); } } } let mut project_config_file = project_conf_folder.clone(); project_config_file.push("project.conf"); 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(( rid, 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(( rid, 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.to_string(), ], created: false, 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_config_file, current: false, db: Some(db), scope: Vec::new(), }; match new_project.save_config() { Ok(_) => {} Err(e) => { let _ = self.main_tx.send(ToolMessage::Output(( 0, format!("[error] saving project! {e}"), ))); } } new_project.save_config()?; self.projects.push(new_project); return Ok(()); } pub fn promote_project(&mut self, rid: usize) -> Result> { self.main_tx.send(ToolMessage::Output(( rid, String::from("promoting project..."), )))?; let return_string = String::from("[success] Project promoted!"); let mut new_project = self.projects[self.selected_project].clone(); let new_files_path = PathBuf::from(self.config.get("current_files").unwrap()) .join(format!("{}/{}", new_project.org_name, new_project.name)); 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)?; self.main_tx.send(ToolMessage::Output(( rid, String::from("[success] folders copied!"), )))?; remove_dir_all(&new_project.files)?; remove_dir_all(&new_project.notes)?; new_project.files(new_files_path.clone()); new_project.notes(new_notes_path.clone()); new_project.current = true; self.main_tx.send(ToolMessage::Output(( rid, String::from("upcoming folders cleaned up!"), )))?; self.main_tx.send(ToolMessage::Output(( rid, String::from("creating distrobox..."), )))?; if let Some(mut db) = new_project.db { db.volumes = vec![ new_files_path.display().to_string(), new_notes_path.display().to_string(), self.config.get("tools").unwrap().to_string(), ]; new_project.db = Some(db); let _ = self.main_tx.send(ToolMessage::RebuildDB); self.main_tx.send(ToolMessage::Output(( rid, String::from("distrobox sent to rebuild function..."), )))?; } new_project.save_config()?; self.projects[self.selected_project] = new_project; return Ok(return_string); } pub fn save_projects(&mut self) -> Result<(), Box> { for mut project in self.projects.clone() { project.save_config()?; } return Ok(()); } pub fn save_config(&mut self) -> Result<(), Box> { let mut config_vec = Vec::new(); for key in self.config.keys() { if let Some(value) = self.config.get(key) { config_vec.push(format!("{}: {}", key, value)); } } if config_vec.len() > 0 { fs::write(self.config_file.clone(), config_vec.join("\n"))?; } else { return Err("No config loaded!".into()); } return Ok(()); } pub fn save_all(&mut self) -> Result<(), Box> { self.save_config()?; self.save_projects()?; return Ok(()); } pub fn new_server(&mut self) -> Result<(), Box> { let mut ip = String::new(); let mut port = String::new(); self.prompt.responses.iter().for_each(|res| { if res.query.id == 0 { ip = res.response.clone(); } else if res.query.id == 1 { port = res.response.clone(); } }); let new_server = Server { address: format!("{}:{}", ip, port), connected: false, timer: Duration::from_secs(60), last_check: time::Instant::now(), message_que: Vec::new(), action_que: Vec::new(), client_id: 0, selected_client: 0, config: None, logged_in: false, password: String::new(), name: String::new(), cert_text: String::new(), }; self.servers.push(Arc::new(Mutex::new(new_server))); self.prompt.reset(); Ok(()) } pub fn new_config(&mut self) -> Result<(), Box> { let mut name = String::new(); let mut config_path = PathBuf::new(); let mut upcoming_files_path = PathBuf::new(); let mut current_files_path = PathBuf::new(); let mut upcoming_notes_path = PathBuf::new(); let mut current_notes_path = PathBuf::new(); let mut tools_folder = PathBuf::new(); let mut template_box = String::new(); let term_cmd = self.config.get("term_cmd").unwrap().clone(); self.prompt .responses .iter() .for_each(|res| match res.query.id { 0 => name = res.response.clone(), 1 => config_path = PathBuf::from(res.response.clone()), 2 => upcoming_files_path = PathBuf::from(res.response.clone()), 3 => current_files_path = PathBuf::from(res.response.clone()), 4 => upcoming_notes_path = PathBuf::from(res.response.clone()), 5 => current_notes_path = PathBuf::from(res.response.clone()), 6 => tools_folder = PathBuf::from(res.response.clone()), 7 => template_box = res.response.clone(), _ => {} }); let mut existing = false; if let Some(line) = self.config.get("alternate_configs") { line.split(",").into_iter().for_each(|line| { if let Some((ename, _)) = line.split_once("|") { if ename.trim().to_string() == name { existing = true; } } }); } if existing { let _ = self.main_tx.send(ToolMessage::Output(( 0, format!("{} already exists!", name), ))); return Ok(()); } let root_config_folder_path = config_path.clone(); let mut projects_path = config_path.clone(); projects_path.push("projects"); let mut module_path = config_path.clone(); module_path.push("modules"); let mut note_templates_path = config_path.clone(); note_templates_path.push("note_templates"); create_dir_all(&module_path)?; config_path.push("client.conf"); let client_conf_path = config_path.clone(); projects_path.push("default"); match create_dir_all(&projects_path) { Ok(_) => { let _ = self.main_tx.send(ToolMessage::Output(( 0, format!( "[success] create projects path at {}", projects_path.display() ), ))); } Err(e) => { let _ = self.main_tx.send(ToolMessage::Output(( 0, format!( "[error] couldn't create projects path at {}: {e}", projects_path.display() ), ))); return Ok(()); } } projects_path.push("project.conf"); let mut client_config_file = File::create(&config_path)?; client_config_file.write(format!("projects: {}\n", projects_path.display()).as_bytes())?; client_config_file.write("servers: 127.0.0.1:31337\n".as_bytes())?; let mut default_project_file = File::create(projects_path)?; config_path.pop(); let mut server_path = config_path.clone(); server_path.push("server"); create_dir_all(&server_path)?; server_path.push("server.conf"); let mut server_config_file = File::create(&server_path)?; server_config_file.write("address: 127.0.0.1:31337\n".as_bytes())?; server_config_file.write("running: false\n".as_bytes())?; default_project_file.write("org_name: default\n".as_bytes())?; default_project_file.write("name: default\n".as_bytes())?; default_project_file .write(format!("notes: {}\n", current_files_path.display()).as_bytes())?; default_project_file .write(format!("files: {}\n", current_notes_path.display()).as_bytes())?; default_project_file.write("stage: current".as_bytes())?; client_config_file .write(format!("current_files: {}\n", current_files_path.display()).as_bytes())?; client_config_file .write(format!("current_notes: {}\n", current_notes_path.display()).as_bytes())?; client_config_file .write(format!("upcoming_files: {}\n", upcoming_files_path.display()).as_bytes())?; client_config_file .write(format!("upcoming_notes: {}\n", upcoming_notes_path.display()).as_bytes())?; client_config_file.write(format!("module_path: {}\n", module_path.display()).as_bytes())?; client_config_file.write(format!("template_box: {}\n", template_box).as_bytes())?; client_config_file.write(format!("term_cmd: {}\n", term_cmd).as_bytes())?; client_config_file.write(format!("tools: {}\n", tools_folder.display()).as_bytes())?; create_dir_all("./temp")?; env::set_current_dir("./temp")?; let git_clone_output = Command::new("git") .arg("clone") .arg("https://git.pyro.monster/pyro/tetanus.git") .output()?; let git_clone_status = git_clone_output.status.success(); if git_clone_status { let mut options = CopyOptions::new(); options.overwrite = true; options.copy_inside = true; copy( "./tetanus/note_templates", root_config_folder_path.clone(), &options, )?; copy( "./tetanus/default-modules", module_path.join("default"), &options, )?; copy( "./tetanus/victim_templates", root_config_folder_path.join("victim"), &options, )?; env::set_current_dir("../")?; remove_dir_all("./temp")?; let _ = self.main_tx.send(ToolMessage::Output(( 0, "New config structure created!".to_string(), ))); } if let Some(original_value) = self.config.get("alternate_configs") { let new_value = format!( "{}|{},{}", name, client_conf_path.display().to_string(), original_value, ); self.config .insert("alternate_configs".to_string(), new_value); } else { self.config.insert( "alternate_configs".to_string(), format!("{}|{},", name, client_conf_path.display().to_string()), ); } self.prompt.responses.clear(); self.prompt.prompts.clear(); self.prompt.action = None; self.prompt.execute_command.clear(); Ok(()) } pub fn parse_cs_portscan(&mut self) -> Result<(), Box> { if let Some(project) = self.projects.get_mut(self.selected_project) { let walk = WalkDir::new(project.files.clone()); walk.into_iter().for_each(|res| { if let Ok(entry) = res { if entry.file_name().to_string_lossy().contains("services.tsv") { if let Ok(scan) = read_to_string(entry.path()) { let mut url_lines = Vec::new(); let mut web_host_lines = Vec::new(); let mut rdp_host_lines = Vec::new(); let mut ftp_host_lines = Vec::new(); let mut telnet_host_lines = Vec::new(); let mut winrm_host_lines = Vec::new(); let mut mysql_host_lines = Vec::new(); let mut mssql_host_lines = Vec::new(); let mut snmp_host_lines = Vec::new(); let mut kerberos_host_lines = Vec::new(); let mut ldap_host_lines = Vec::new(); let mut smb_host_lines = Vec::new(); let mut unknonw_host_lines = Vec::new(); let mut ssh_host_lines = Vec::new(); scan.lines().into_iter().for_each(|line| { if line.len() > 0 { let data: Vec<&str> = line.split_whitespace().collect(); match data[1].trim() { "80" | "443" | "8080" | "8443" | "8000" | "4433" => { let http_line = format!("http://{}:{}", data[0], data[1]); let https_line = format!("https://{}:{}", data[0], data[1]); url_lines.push(http_line); url_lines.push(https_line); web_host_lines.push(data[0].to_string()); let mut existing = false; project .hosts .iter_mut() .filter(|h| h.ip == data[0].to_string()) .for_each(|h| { h.add_port(data[1]); existing = true; }); if !existing { let mut new_host = Host::new(); new_host.ip = data[0].trim().to_string(); new_host.add_port(data[1]); project.add_host(new_host); } } "3389" => { rdp_host_lines.push(data[0].to_string()); let mut existing = false; project .hosts .iter_mut() .filter(|h| h.ip == data[0].to_string()) .for_each(|h| { h.add_port(data[1]); existing = true; }); if !existing { let mut new_host = Host::new(); new_host.ip = data[0].trim().to_string(); new_host.add_port(data[1]); project.add_host(new_host); } } "21" => { ftp_host_lines.push(data[0].to_string()); let mut existing = false; project .hosts .iter_mut() .filter(|h| h.ip == data[0].to_string()) .for_each(|h| { h.add_port(data[1]); existing = true; }); if !existing { let mut new_host = Host::new(); new_host.ip = data[0].trim().to_string(); new_host.add_port(data[1]); project.add_host(new_host); } } "23" => { telnet_host_lines.push(data[0].to_string()); let mut existing = false; project .hosts .iter_mut() .filter(|h| h.ip == data[0].to_string()) .for_each(|h| { h.add_port(data[1]); existing = true; }); if !existing { let mut new_host = Host::new(); new_host.ip = data[0].trim().to_string(); new_host.add_port(data[1]); project.add_host(new_host); } } "22" => { ssh_host_lines.push(data[0].to_string()); let mut existing = false; project .hosts .iter_mut() .filter(|h| h.ip == data[0].to_string()) .for_each(|h| { h.add_port(data[1]); existing = true; }); if !existing { let mut new_host = Host::new(); new_host.ip = data[0].trim().to_string(); new_host.add_port(data[1]); project.add_host(new_host); } } "5985" | "5986" => { winrm_host_lines.push(data[0].to_string()); let mut existing = false; project .hosts .iter_mut() .filter(|h| h.ip == data[0].to_string()) .for_each(|h| { h.add_port(data[1]); existing = true; }); if !existing { let mut new_host = Host::new(); new_host.ip = data[0].trim().to_string(); new_host.add_port(data[1]); project.add_host(new_host); } } "3306" => { mysql_host_lines.push(data[0].to_string()); let mut existing = false; project .hosts .iter_mut() .filter(|h| h.ip == data[0].to_string()) .for_each(|h| { h.add_port(data[1]); existing = true; }); if !existing { let mut new_host = Host::new(); new_host.ip = data[0].trim().to_string(); new_host.add_port(data[1]); project.add_host(new_host); } } "1433" => { mssql_host_lines.push(data[0].to_string()); let mut existing = false; project .hosts .iter_mut() .filter(|h| h.ip == data[0].to_string()) .for_each(|h| { h.add_port(data[1]); existing = true; }); if !existing { let mut new_host = Host::new(); new_host.ip = data[0].trim().to_string(); new_host.add_port(data[1]); project.add_host(new_host); } } "161" => { snmp_host_lines.push(data[0].to_string()); let mut existing = false; project .hosts .iter_mut() .filter(|h| h.ip == data[0].to_string()) .for_each(|h| { h.add_port(data[1]); existing = true; }); if !existing { let mut new_host = Host::new(); new_host.ip = data[0].trim().to_string(); new_host.add_port(data[1]); project.add_host(new_host); } } "88" => { kerberos_host_lines.push(data[0].to_string()); let mut existing = false; project .hosts .iter_mut() .filter(|h| h.ip == data[0].to_string()) .for_each(|h| { h.add_port(data[1]); existing = true; }); if !existing { let mut new_host = Host::new(); new_host.ip = data[0].trim().to_string(); new_host.add_port(data[1]); project.add_host(new_host); } } "389" | "636" => { ldap_host_lines.push(data[0].to_string()); let mut existing = false; project .hosts .iter_mut() .filter(|h| h.ip == data[0].to_string()) .for_each(|h| { h.add_port(data[1]); existing = true; }); if !existing { let mut new_host = Host::new(); new_host.ip = data[0].trim().to_string(); new_host.add_port(data[1]); project.add_host(new_host); } } "445" | "139" => { smb_host_lines.push(data[0].to_string()); let mut existing = false; project .hosts .iter_mut() .filter(|h| h.ip == data[0].to_string()) .for_each(|h| { h.add_port(data[1]); existing = true; }); if !existing { let mut new_host = Host::new(); new_host.ip = data[0].trim().to_string(); new_host.add_port(data[1]); project.add_host(new_host); } } _ => { unknonw_host_lines.push(data[0].to_string()); let mut existing = false; project .hosts .iter_mut() .filter(|h| h.ip == data[0].to_string()) .for_each(|h| { h.add_port(data[1]); existing = true; }); if !existing { let mut new_host = Host::new(); new_host.ip = data[0].trim().to_string(); new_host.add_port(data[1]); project.add_host(new_host); } } } } }); if !url_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("urls.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", url_lines.join("\n")); } } if !web_host_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("web_hosts.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", web_host_lines.join("\n")); } } if !rdp_host_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("rdp_hosts.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", rdp_host_lines.join("\n")); } } if !ftp_host_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("ftp_hosts.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", ftp_host_lines.join("\n")); } } if !telnet_host_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("telnet_hosts.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", telnet_host_lines.join("\n")); } } if !winrm_host_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("winrm_hosts.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", winrm_host_lines.join("\n")); } } if !mysql_host_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("mysql_hosts.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", mysql_host_lines.join("\n")); } } if !mssql_host_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("mssql_hosts.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", mssql_host_lines.join("\n")); } } if !snmp_host_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("snmp_hosts.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", snmp_host_lines.join("\n")); } } if !kerberos_host_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("kerberos_hosts.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", kerberos_host_lines.join("\n")); } } if !ldap_host_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("ldap_hosts.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", ldap_host_lines.join("\n")); } } if !smb_host_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("smb_hosts.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", smb_host_lines.join("\n")); } } if !unknonw_host_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("unknown_hosts.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", unknonw_host_lines.join("\n")); } } if !ssh_host_lines.is_empty() { let mut file_path = project.files.clone(); file_path.push("ssh_hosts.txt"); if let Ok(mut file) = File::create(file_path) { let _ = write!(file, "{}", ssh_host_lines.join("\n")); } } } } } }); project.save_config()?; } Ok(()) } } #[derive(Clone, Debug)] pub struct Server { pub address: String, pub connected: bool, pub timer: Duration, pub config: Option>, pub last_check: time::Instant, pub message_que: Vec, pub action_que: Vec, pub client_id: usize, pub selected_client: usize, pub logged_in: bool, pub password: String, pub name: String, pub cert_text: String, } impl Server { pub fn donwload_certificate(&self) -> Result> { let mut stream = TcpStream::connect(self.address.clone())?; stream.write_all("CERT_REQ".as_bytes())?; let mut buf = [0u8; 8192]; let bytes_read = stream.read(&mut buf)?; let response = String::from_utf8_lossy(&buf[..bytes_read]); if let Some((_, cert)) = response.split_once("|") { return Ok(cert.to_string()); } Err("Server did not send certificate".into()) } pub fn connect(&mut self) -> Result> { if let Ok(cert) = self.donwload_certificate() { self.config = Some(self.build_tls_config(cert.as_str())?); self.cert_text = cert; } let mut stream = self.tls_connect()?; stream.write("HELLO|attack".as_bytes())?; let mut buf = [0; 8192]; let bytes_read = stream.read(&mut buf)?; let response = String::from_utf8_lossy(&buf[..bytes_read]); if let Some((_, id)) = response.split_once("|") { self.client_id = id.trim().parse()?; } self.connected = true; self.last_check = Instant::now(); return Ok(self.client_id); } pub fn login(&mut self, pass: String) -> Result<(), Box> { let mut stream = self.tls_connect()?; stream.write(format!("NONE**{}|||LOGIN|{}||", self.client_id, pass).as_bytes())?; let mut buf = [0; 8192]; let bytes_read = stream.read(&mut buf)?; let response = String::from_utf8_lossy(&buf[..bytes_read]); if response.contains("successful") { self.password = pass; self.last_check = Instant::now(); self.logged_in = true; self.message_que.push(format!("SET_NAME|{}", self.name)); return Ok(()); } else { return Err("[error] logging in!".into()); } } pub fn checkin(&mut self) { if self.logged_in { if let Ok(mut stream) = self.tls_connect() { self.last_check = Instant::now(); let payload = format!( "{}**{}|||{}\n", self.password.clone(), self.client_id.clone(), self.message_que.join("||") ); if let Err(e) = stream.write_all(payload.as_bytes()) { self.action_que .push(format!("Error|sending reponse output to server! {e}")); } self.message_que.clear(); let mut buffer = [0; 4096]; if let Ok(bytes_read) = stream.read(&mut buffer) { if bytes_read > 0 { let response = String::from_utf8_lossy(&buffer[..bytes_read]); response.split("||").into_iter().for_each(|action| { self.action_que.push(action.trim().to_string()); if action.contains("Not authenticated") { self.logged_in = false; } }); } } } } } pub fn build_tls_config(&self, pem: &str) -> Result, Box> { let mut roots = RootCertStore::empty(); let mut reader = BufReader::new(pem.as_bytes()); let certs = rustls_pemfile::certs(&mut reader).collect::, _>>()?; for cert in certs { roots.add(cert)?; } let config = ClientConfig::builder() .with_root_certificates(roots) .with_no_client_auth(); Ok(Arc::new(config)) } fn tls_connect(&self) -> Result, Box> { let config = self.config.clone().ok_or("TLS config not loaded")?; let tcp = TcpStream::connect(&self.address)?; let (host, _) = self.address.split_once(':').ok_or("Invalid address")?; let server_name = ServerName::try_from(host.to_string())?; let connection = ClientConnection::new(config, server_name)?; let tls_stream = StreamOwned::new(connection, tcp); Ok(tls_stream) } fn disconnect(&mut self) -> Result<(), Box> { let mut stream = self.tls_connect()?; stream.write(format!("{}**{}|||DISCONNECT||", self.password, self.client_id).as_bytes())?; self.connected = false; Ok(()) } } #[derive(Clone, Debug)] pub struct Project { pub org_name: String, pub config_folder: PathBuf, pub name: String, pub notes: PathBuf, pub files: PathBuf, pub hosts: Vec, pub current: bool, pub db: Option, pub scope: Vec, } impl Project { pub fn new() -> Self { Self { org_name: String::new(), config_folder: PathBuf::new(), name: String::new(), notes: PathBuf::new(), files: PathBuf::new(), hosts: Vec::new(), current: false, db: None, scope: Vec::new(), } } pub fn org_name(&mut self, name: String) { self.org_name = name; } pub fn notes(&mut self, path: PathBuf) { self.notes = path; } pub fn files(&mut self, path: PathBuf) { self.files = path; } pub fn hosts(&mut self, hosts: Vec) { self.hosts = hosts; } pub fn add_host(&mut self, mut host: Host) { let mut new_id = 0; self.hosts.iter().for_each(|h| { if h.id > new_id { new_id = h.id.clone(); } }); new_id += 1; if !self.hosts.iter().any(|h| h.ip == host.ip) { host.id = new_id; self.hosts.push(host); } } pub fn config_folder(&mut self, path: PathBuf) { self.config_folder = path; } pub fn add_scope(&mut self, host: String) { if let Ok(ip) = host.parse::() { self.scope.push(ip.to_string()) } else if let Ok(net) = host.parse::() { net.hosts().into_iter().for_each(|ip| { self.scope.push(ip.to_string()); }); } else { self.scope.push(host); } } pub fn load_config(&mut self, display: bool) -> Result<(), Box> { let config_contents = read_to_string(&self.config_folder)?; for line in config_contents.lines() { let parts: Vec<&str> = line.split(": ").collect(); if parts.len() == 2 { match parts[0].trim() { "org_name" => self.org_name(parts[1].trim().to_string()), "name" => self.name = parts[1].trim().to_string(), "notes" => self.notes(PathBuf::from(parts[1].trim())), "files" => self.files(PathBuf::from(parts[1].trim())), "stage" => { if parts[1].trim() == "current" { self.current = true; } } "scope" => { let ips: Vec<&str> = parts[1].split(",").collect(); ips.iter().for_each(|ip| { self.add_scope(ip.trim().to_string()); }); } _ => {} } } } if display { println!("{} | {} config loaded!", self.org_name, self.name); println!("loading hosts..."); } let mut conf_folder = self.config_folder.clone(); conf_folder.pop(); let mut hosts_folder = conf_folder.clone(); let mut users_folder = conf_folder.clone(); hosts_folder.push("hosts"); users_folder.push("users"); let mut users = HashMap::new(); if users_folder.exists() { for entry in read_dir(users_folder)? { let entry = entry?; let user_config_string = read_to_string(entry.path())?; let mut new_user = User::new(); for line in user_config_string.lines() { if line.contains(": ") { let (setting, data) = line.split_once(": ").unwrap(); match setting.trim() { "username" => new_user.name = data.trim().to_string(), "password" => new_user.password = Some(data.trim().to_string()), "hash" => new_user.hash = Some(data.trim().to_string()), "compromised" => new_user.compromised = data.trim().parse().unwrap(), "ticket" => new_user.ticket = Some(PathBuf::from(data.trim())), _ => {} } } } users.insert(new_user.name.clone(), new_user); } } if hosts_folder.exists() { read_dir(hosts_folder)?.into_iter().for_each(|res| { if let Ok(entry) = res { let host_path = entry.path(); let mut new_host = Host::new(); println!("loading {}...", host_path.display()); if let Ok(host_conf_text) = read_to_string(host_path) { host_conf_text.lines().into_iter().for_each(|line| { if line.contains(": ") { let (setting, data) = line.split_once(": ").unwrap(); match setting.trim() { "ip" => new_host.ip = data.trim().to_string(), "hostname" => new_host.hostname = data.trim().to_string(), "control_port" => { new_host.control_port = data.trim().parse().unwrap(); } "pwned" => { new_host.pwned = data.trim().parse().unwrap(); } "id" => new_host.id = data.trim().parse().unwrap(), "findings" => { new_host.findings = data .trim() .split(",") .map(|finding| finding.to_string()) .collect::>() } "open_ports" => { new_host.open_ports = data .trim() .split(",") .into_iter() .map(|port| port.parse().unwrap_or_default()) .collect::>() } "users" => { let names: Vec<&str> = data.trim().split(",").collect(); names.into_iter().for_each(|name| { if let Some(user) = users.get(name) { new_host.users.push(user.clone()); } }); } _ => {} } } }); self.add_host(new_host); } } }); } if display { println!("{} | {} loaded!", self.org_name, self.name); } return Ok(()); } pub fn save_config(&mut self) -> Result<(), Box> { let mut conf_folder = self.config_folder.clone(); conf_folder.pop(); if !conf_folder.exists() { match create_dir_all(&conf_folder) { Ok(_) => {} Err(e) => { return Err( format!("save_config: couldn't create project conf folder! {e}").into(), ); } } } let mut hosts_folder = conf_folder.clone(); hosts_folder.push("hosts"); let mut users_folder = conf_folder.clone(); users_folder.push("users"); if !hosts_folder.exists() { create_dir_all(&hosts_folder)?; } if !users_folder.exists() { create_dir_all(&users_folder)?; } let mut file = File::create(&self.config_folder)?; let mut out_string = format!( "org_name: {}\nname: {}\nnotes: {}\nfiles: {}\nscope: {}\n", self.org_name, self.name, self.notes.display(), self.files.display(), self.scope.join(","), ); if self.current { out_string.push_str("stage: current"); } else { out_string.push_str("stage: upcoming"); } file.write_all(out_string.as_bytes())?; self.hosts .iter_mut() .try_for_each(|h| -> Result<(), Box> { let mut host_path = hosts_folder.clone(); host_path.push(format!("{}.conf", h.id)); h.save_config_file(host_path) })?; return Ok(()); } } impl std::fmt::Display for Project { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{} | {}", self.org_name, self.name) } } impl PartialEq for Project { fn eq(&self, other: &Self) -> bool { format!("{} | {}", self.org_name, self.name) == format!("{} | {}", other.org_name, other.name) } } impl Eq for Project {} #[derive(Clone, Debug)] pub struct Host { pub ip: String, pub hostname: String, pub open_ports: Vec, pub control_port: usize, pub users: Vec, pub pwned: bool, pub id: usize, pub findings: Vec, } impl Host { pub fn new() -> Self { Self { ip: String::new(), hostname: String::new(), open_ports: Vec::new(), control_port: 0, users: Vec::new(), pwned: false, id: 0, findings: Vec::new(), } } pub fn ip(&mut self, ip: String) { self.ip = ip; } pub fn hostname(&mut self, name: String) { self.hostname = name; } pub fn open_ports(&mut self, ports: Vec) { self.open_ports = ports; } pub fn control_port(&mut self, port: usize) { self.control_port = port; } pub fn users(&mut self, users: Vec) { self.users = users; } pub fn pwnd(&mut self) { self.pwned = true; } pub fn id(&mut self, id: usize) { self.id = id; } pub fn findings(&mut self, findings: Vec) { self.findings = findings; } pub fn add_port(&mut self, port: &str) { if let Ok(port) = port.trim().parse() { if !self.open_ports.contains(&port) { self.open_ports.push(port); } } } pub fn add_user(&mut self, user: User) { self.users.push(user); } pub fn add_finding(&mut self, finding: String) { self.findings.push(finding); } pub fn save_config_file(&mut self, path: PathBuf) -> Result<(), Box> { let mut host_file = File::create(path)?; let mut out_string = format!( " ip: {} hostname: {} control_port: {} pwned: {} id: {}", self.ip, self.hostname, self.control_port, self.pwned, self.id ); if !self.findings.is_empty() { out_string.push_str(&format!("\nfindings: ")); out_string.push_str(&self.findings.join(",")); } if !self.open_ports.is_empty() { out_string.push_str("\nopen_ports: "); out_string.push_str( &self .open_ports .iter() .map(|port| port.to_string()) .collect::>() .join(","), ); } write!(host_file, "{}", out_string)?; host_file.sync_all()?; Ok(()) } } #[derive(Clone, Debug)] pub struct User { pub name: String, pub password: Option, pub hash: Option, pub ticket: Option, pub compromised: bool, } impl User { pub fn new() -> Self { Self { name: String::new(), password: None, hash: None, ticket: None, compromised: false, } } } #[derive(Clone, Debug)] pub struct ToolCommand { pub name: String, pub path: PathBuf, pub help: String, pub args: Vec, pub output_type: String, pub finished: bool, pub result: bool, pub calls: String, } pub struct ModuleLoader { pub engine: Arc, pub commands: HashMap, pub asts: HashMap, } impl ModuleLoader { pub fn new() -> Self { let mut engine = Engine::new(); engine .register_type_with_name::("PathBuf") .register_get_set( "display", |p: &mut PathBuf| p.to_string_lossy().into_owned(), |p: &mut PathBuf, s: String| *p = PathBuf::from(s), ); engine .register_type_with_name::("Host") .register_get_set( "ip", |h: &mut Host| h.ip.clone(), |h: &mut Host, val: String| h.ip = val, ) .register_get_set( "hostname", |h: &mut Host| h.hostname.clone(), |h: &mut Host, val: String| h.hostname = val, ) .register_get_set( "pwned", |h: &mut Host| h.pwned, |h: &mut Host, val: bool| h.pwned = val, ) .register_fn("add_port", Host::add_port) .register_fn("add_finding", Host::add_finding); engine .register_type_with_name::("Project") .register_get_set( "name", |p: &mut Project| p.name.clone(), |p: &mut Project, val: String| p.name = val, ) .register_get_set( "org_name", |p: &mut Project| p.org_name.clone(), |p: &mut Project, val: String| p.org_name = val, ) .register_get_set( "current", |p: &mut Project| p.current.clone(), |p: &mut Project, val: bool| p.current = val, ) .register_get_set( "notes", |p: &mut Project| p.notes.clone(), |p: &mut Project, val: PathBuf| p.notes = val, ) .register_get_set( "files", |p: &mut Project| p.files.clone(), |p: &mut Project, val: PathBuf| p.files = val, ) .register_fn("add_host", Project::add_host); engine.register_fn("get_host", |hosts: &mut Vec, index: i64| { hosts.get(index as usize).cloned().unwrap_or_else(Host::new) }); engine.register_fn("new_host", Host::new); engine.register_fn("open_file", |path: &str| { std::fs::read_to_string(path).unwrap_or_else(|_| String::new()) }); engine.register_fn("add_port", |host: &mut Host, port: i64| { host.add_port(format!("{}", port).as_str()); }); engine.register_fn("expand_target", |s: &str| -> rhai::Array { let mut results = rhai::Array::new(); if let Ok(net) = s.trim().parse::() { for host in net.hosts() { results.push(host.to_string().into()); } } else if let Ok(ip) = s.trim().parse::() { results.push(ip.to_string().into()); } else if s.contains(".") { results.push(s.to_string().into()); } results }); engine.register_fn("filter_ips", |s: &str| -> rhai::Array { let mut results = rhai::Array::new(); if let Ok(net) = s.trim().parse::() { results.push(net.addr().to_string().trim().into()); } else if let Ok(ip) = s.trim().parse::() { results.push(ip.to_string().trim().into()); } else if s.contains(".") { results.push(s.to_string().trim().into()); } results }); engine.register_fn("write_file", |path: &str, contents: &str| { std::fs::write(path, contents).is_ok() }); engine.register_fn("find_file", |path: PathBuf, filename: String| -> String { for entry in walkdir::WalkDir::new(path).max_depth(50) { if let Ok(entry) = entry { if entry .file_name() .to_string_lossy() .contains(filename.as_str()) { return entry.path().display().to_string(); } } } String::new() }); engine.register_fn("execute", |cmd: &str| -> String { let out_string; if let Ok(output) = Command::new("sh").arg("-c").arg(cmd).output() { if output.status.success() { out_string = format!("[success] {}", String::from_utf8_lossy(&output.stdout)); } else { out_string = format!("[error] {}", String::from_utf8_lossy(&output.stderr)); } } else { out_string = "[error] failed to run command".to_string(); } return out_string; }); engine.register_fn("db_run", |p: &mut Project, command: String| -> String { let mut out_string = format!("No distrobox configured for selected project."); if let Some(db) = p.db.clone() { out_string = db.run_command(&command); } return out_string; }); Self { engine: Arc::new(engine), commands: HashMap::new(), asts: HashMap::new(), } } pub fn load_all(&mut self, base_path: &PathBuf) -> Result<(), std::io::Error> { self.commands.clear(); self.asts.clear(); self.load_from_dir(&base_path.join("default"))?; self.load_from_dir(&base_path.join("custom"))?; Ok(()) } fn load_from_dir(&mut self, dir: &PathBuf) -> Result<(), std::io::Error> { if !dir.exists() { return Ok(()); } for entry in read_dir(dir)? { let entry = entry?; let path = entry.path(); if path.is_dir() { if let Some((command, ast)) = self.parse_module_directory(&path) { self.asts.insert(command.name.clone(), ast); self.commands.insert(command.name.clone(), command); } } } Ok(()) } fn parse_module_directory(&self, dir: &PathBuf) -> Option<(ToolCommand, AST)> { let config_path = dir.join("config.conf"); let help_path = dir.join("help.txt"); let script_path = dir.join("script.rhai"); if !config_path.exists() || !help_path.exists() || !script_path.exists() { return None; } let config_content = read_to_string(&config_path).ok()?; let mut name = String::new(); let mut output_type = String::new(); let mut args = Vec::new(); let mut calls = String::new(); for line in config_content.lines() { if let Some((key, val)) = line.split_once(':') { match key.trim() { "name" => name = val.trim().to_string(), "output_type" => output_type = val.trim().to_string(), "args" => { args = val .split(',') .map(|s| s.trim().to_string()) .filter(|s| !s.is_empty()) .collect(); } "calls" => { calls = val.trim().to_string(); } _ => {} } } } if name.is_empty() { return None; } let help = read_to_string(&help_path).unwrap_or_default(); let ast = match self.engine.compile_file(script_path.clone()) { Ok(compiled_ast) => compiled_ast, Err(e) => { eprintln!("[error] compiling script in {:?}: {}", script_path, e); return None; } }; let command = ToolCommand { name, path: script_path, help, args, output_type, finished: false, result: false, calls, }; Some((command, ast)) } } #[derive(Debug, Clone)] pub struct DistroBox { pub name: String, pub volumes: Vec, pub created: bool, pub template: String, } impl DistroBox { pub fn create(&mut self, tx: Sender, rid: usize) -> Result<(), Box> { if let Err(e) = self.stop_template(tx.clone(), rid) { tx.send(ToolMessage::Output(( rid, format!("[error] stopping template box! {e}"), )))?; return Ok(()); } let mut create_command = Command::new("distrobox"); create_command .arg("create") .arg("--clone") .arg(self.template.clone()) .arg("--name") .arg(self.name.clone()); let mut added_volumes = Vec::new(); for volume in &self.volumes { let folder_name = volume.split("/").last().unwrap(); if !added_volumes.contains(&String::from(folder_name)) { create_command .arg("--volume") .arg(format!("{}:/{}:rw", volume, folder_name)); added_volumes.push(folder_name.to_string()); } } 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((rid, line.clone()))); println!("{}", 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((rid, line.clone()))); println!("{}", line); } }); } let status = child.wait()?; if !status.success() { let _ = tx.send(ToolMessage::Output(( rid, format!("[error] creating distrobox!"), ))); } self.created = true; return Ok(()); } pub fn stop_template( &mut self, tx: Sender, rid: usize, ) -> Result<(), Box> { let mut stop_command = Command::new("distrobox"); stop_command .arg("stop") .arg(self.template.clone()) .arg("--yes"); stop_command.stdin(Stdio::piped()); stop_command.stdout(Stdio::piped()); stop_command.stderr(Stdio::piped()); let mut child = stop_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((rid, line.clone()))); println!("{}", 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((rid, line.clone()))); println!("{}", line); } }); } let status = child.wait()?; if !status.success() { let _ = tx.send(ToolMessage::Output(( rid, format!("[error] stopping template distrobox!"), ))); } return Ok(()); } pub fn stop(&mut self, tx: Sender, rid: usize) -> Result<(), Box> { let mut stop_command = Command::new("distrobox"); stop_command.arg("stop").arg(self.name.clone()).arg("--yes"); stop_command.stdin(Stdio::piped()); stop_command.stdout(Stdio::piped()); stop_command.stderr(Stdio::piped()); let mut child = stop_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((rid, line.clone()))); } }); } 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((rid, line.clone()))); } }); } let status = child.wait()?; if !status.success() { let _ = tx.send(ToolMessage::Output(( rid, format!("[error] stopping distrobox!"), ))); } return Ok(()); } pub fn destroy(&mut self, tx: Sender, rid: usize) -> Result<(), Box> { let mut destroycmd = Command::new("distrobox"); destroycmd.arg("rm").arg(self.name.clone()).arg("-f"); let res = destroycmd.status()?; if !res.success() { let _ = tx.send(ToolMessage::Output(( rid, format!("[error] destroying distrobox!"), ))); let _ = tx.send(ToolMessage::Output(( rid, String::from("pleas try destroying manually:"), ))); let _ = tx.send(ToolMessage::Output(( rid, format!("distrobox rm {} -f", self.name), ))); } self.created = false; return Ok(()); } pub fn launch_terminal(&self, cmd: String) { if let Some((terminal_cmd, args)) = cmd.split_once(" ") { let mut term_cmd = Command::new(terminal_cmd); args.split(" ").into_iter().for_each(|arg| { if arg == "|||COMMAND|||" { term_cmd .arg("distrobox") .arg("enter") .arg(self.name.clone()); } else if arg.contains("ENV_NAME=") { let (_, env_name) = arg.split_once("=").unwrap(); unsafe { std::env::set_var(env_name, self.name.clone()) }; } else { term_cmd.arg(arg); } }); let _ = term_cmd.spawn(); } } pub fn run_command(&self, cmd: &str) -> String { let output = Command::new("distrobox") .arg("enter") .arg("-e") .arg(format!("\"{}\"", cmd)) .output(); match output { Ok(output) => { if output.status.success() { return String::from_utf8_lossy(&output.stdout).to_string(); } else { return format!("[error] {}", String::from_utf8_lossy(&output.stdout)); } } Err(e) => { return format!("[error] {e}"); } } } } #[derive(Clone)] struct Prompt { action: Option, responses: Vec, execute_command: String, num_responses: usize, prompts: Vec, last_prompted: Option, } impl Prompt { pub fn reset(&mut self) { self.action = None; self.responses.clear(); self.execute_command.clear(); self.num_responses = 0; self.prompts.clear(); self.last_prompted = None; } } #[derive(Clone)] struct PromptResponse { query: PromptQuery, response: String, } #[derive(Clone)] struct PromptQuery { id: usize, query: String, }