use fs_extra::dir::{CopyOptions, copy}; use ratatui::crossterm::event; use rhai::{AST, Dynamic, Engine, Scope}; use std::collections::HashMap; 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, Write}; use std::path::PathBuf; use std::process::{Command, Stdio}; use std::sync::mpsc::Receiver; use std::sync::{Arc, mpsc::Sender, mpsc::channel}; pub mod funcs; 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: Option, pub selected_project: usize, pub curent_intput: String, pub module_loader: ModuleLoader, pub output_scroll: u16, prompt: Prompt, } impl AppState { pub fn new() -> (Self, Receiver) { let (main_tx, main_rx) = 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: None, selected_project: 0, curent_intput: String::new(), module_loader: ModuleLoader::new(), output_scroll: 0, prompt: Prompt { action: None, responses: Vec::new(), execute_command: String::new(), num_responses: 0, }, }, main_rx, ) } pub fn load_config(&mut self, file: PathBuf) -> 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" => { let addresses: Vec<&str> = parts[1].trim().split(", ").collect(); for address in addresses { let new_server = Server { address: address.trim().to_string(), connected: false, }; self.servers.push(new_server); } self.config .insert("servers".to_string(), line.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()?; 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); } } } } } 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 { 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, ) -> Result<(), Box> { let tx = self.main_tx.clone(); match command_name { "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()) { tx.send(ToolMessage::Output(format!( "Error making {}-{}: {e}", org, name )))?; } else { tx.send(ToolMessage::Output(format!( "{}-{} created successfully!", org, name )))?; } } "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")); } for line in out_vec { let _ = self.main_tx.send(ToolMessage::Output(line)); } } "promote_project" | "pp" => { self.promote_project()?; } "remove_project_confirm" => { self.prompt.action = None; self.prompt.num_responses = 0; if self.prompt.responses.len() > 0 { if self.prompt.responses[0] == "y".to_string() { let project = self.projects[self.selected_project].clone(); let mut config_file = project.config_folder.clone(); config_file.pop(); if let Err(e) = remove_dir_all(&project.config_folder) { let _ = self.main_tx.send(ToolMessage::Output(format!( "failed to delete config: {e} on {}", &project.config_folder.display() ))); } if let Err(e) = remove_dir_all(&project.notes) { let _ = self.main_tx.send(ToolMessage::Output(format!( "failed to delete notes: {e} on {}", &project.notes.display() ))); } 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(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(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(format!( "failed to delete empty files parent: {e} on {}", parent.display() ))); } } } } self.projects.remove(self.selected_project); let _ = self.main_tx.send(ToolMessage::Output(format!( "{}-{} was sucessfully removed!", project.org_name, project.name ))); } } self.prompt.responses.clear(); } "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"); let _ = self.main_tx.send(ToolMessage::Output(format!( "{}, {} and all contents will be deleted. Continue? (y/N)", project.files.display(), project.notes.display() ))); } _ => { let ast = self .module_loader .asts .get(command_name) .ok_or_else(|| format!("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.as_str() { "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); } _ 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); } _ => {} } } 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() { result .into_string() .unwrap_or_else(|_| "failed to parse string".into()) } else { format!("{:?}", result) }; if result_str.contains("\n") { result_str.lines().into_iter().for_each(|line| { let _ = tx.send(ToolMessage::Output(line.to_string())); }); } else { let _ = tx.send(ToolMessage::Output(result_str)); } } Err(err) => { let _ = tx.send(ToolMessage::Output(format!( "Script Error in execution: {}", err ))); } } }); } } Ok(()) } pub fn new_project(&mut self, org: String, name: String) -> Result<(), Box> { 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 project_notes = PathBuf::from(self.config.get("upcoming_notes").unwrap()) .join(format!("{}/{}", org, name)); let mut template_path = self.config_file.clone(); let mut project_conf_folder = self.config_file.clone(); let tools = self.config.get("tools").unwrap(); template_path.pop(); template_path.push("note_templates"); project_conf_folder.pop(); project_conf_folder.push(format!("projects/{}-{}", org, name)); let mut options = CopyOptions::new(); options.overwrite = true; match create_dir_all(&project_files) { Ok(_) => {} Err(e) => { let _ = self.main_tx.send(ToolMessage::Output(format!( "error making project files! {e} on {}", project_files.display() ))); } } match create_dir_all(&project_notes) { Ok(_) => {} Err(e) => { let _ = self.main_tx.send(ToolMessage::Output(format!( "error making project notes! {e} on {}", project_notes.display() ))); } } match create_dir_all(&project_conf_folder) { Ok(_) => {} Err(e) => { let _ = self.main_tx.send(ToolMessage::Output(format!( "error making config folder! {e} on {}", project_conf_folder.display() ))); } } for entry in read_dir(template_path)? { let entry = entry?; let template_name_os = entry.file_name(); let template_name = template_name_os.to_string_lossy(); if name.clone().contains(template_name.as_ref()) { for entry in read_dir(entry.path())? { let entry = entry?; let file_name_os = entry.file_name(); let file_name = file_name_os.to_string_lossy(); if entry.file_type()?.is_dir() { let dest = project_notes.join(file_name.as_ref()); match copy(entry.path(), dest, &options) { Ok(_) => {} Err(e) => { let _ = self.main_tx.send(ToolMessage::Output(format!( "error copying note template! {e}" ))); } } } else { match fs::copy(entry.path(), project_notes.join(file_name.as_ref())) { Ok(_) => {} Err(e) => { let _ = self.main_tx.send(ToolMessage::Output(format!( "error copying note template! {e}" ))); } } } } } } let db = DistroBox { name: format!("{}-{}-{}", template_box, org, name), volumes: vec![ project_files.display().to_string(), project_notes.display().to_string(), tools.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_conf_folder, current: false, db: Some(db), }; new_project.save_config()?; self.projects.push(new_project); return Ok(()); } pub fn promote_project(&mut self) -> Result> { let mut return_string = String::from("Project promoted successfully!"); 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)?; let cleanup_res = Command::new("rm") .arg("-rf") .arg(new_project.files.display().to_string()) .arg(new_project.notes.display().to_string()) .status(); if cleanup_res.is_err() { return_string = format!( "Error deleting {} and {}, please clean up manually. Otherwise Project promotion succeeded!", new_files_path.display(), new_notes_path.display() ); } new_project.files(new_files_path.clone()); new_project.notes(new_notes_path.clone()); 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.projects[self.selected_project] = new_project; return Ok(return_string); } } #[derive(Clone)] pub struct Server { pub address: String, pub connected: bool, } #[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, } 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, } } 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, host: Host) { self.hosts.push(host); } pub fn config_folder(&mut self, path: PathBuf) { self.config_folder = path; } pub fn load_config(&mut self) -> 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; } } _ => {} } } } println!("{} | {} loaded!", self.org_name, self.name); return Ok(()); } pub fn save_config(&mut self) -> Result<(), Box> { create_dir_all(&self.config_folder)?; let mut conf_file = self.config_folder.clone(); conf_file.push("project.conf"); let mut file = File::create(conf_file)?; let mut out_string = format!( "org_name: {}\nname: {}\nnotes: {}\nfiles: {}\n", self.org_name, self.name, self.notes.display(), self.files.display() ); if self.current { out_string.push_str("stage: current"); } else { out_string.push_str("stage: upcoming"); } file.write_all(out_string.as_bytes())?; 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: usize) { 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 enum Destination { Server, Victim, Attacker, Control, } #[derive(Clone, Debug)] pub struct User { pub name: String, pub password: Option, pub hash: Option, pub ticket: Option, pub compromised: bool, } #[derive(Clone, Debug)] pub enum ToolMessage { Input(String), Output(String), UpdateHost(Host), RebuildDB, UpdateProject(usize, Project), RemoveProject, } #[derive(Clone, Debug)] pub enum ToolArg { Project(Project), Projects(Vec), Host(Host), Hosts(Vec), Config(HashMap), Path(PathBuf), } #[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 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_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) }); 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(); 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(); } _ => {} } } } 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, }; Some((command, ast)) } } enum AppEvent { Key(event::KeyEvent), Worker(ToolMessage), Mouse(event::MouseEvent), } #[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) -> Result<(), Box> { let mut create_command = Command::new("distrobox"); create_command .arg("create") .arg("--root") .arg("--clone") .arg(self.template.clone()) .arg("--name") .arg(self.name.clone()); for volume in &self.volumes { let folder_name = volume.split("/").last().unwrap(); create_command .arg("--volume") .arg(format!("{}:/{}:rw", volume, folder_name)); } create_command.stdout(Stdio::piped()); create_command.stderr(Stdio::piped()); let mut child = create_command.spawn()?; if let Some(stdout) = child.stdout.take() { let tx_out = tx.clone(); std::thread::spawn(move || { let reader = BufReader::new(stdout); for line in reader.lines().flatten() { let _ = tx_out.send(ToolMessage::Output(line)); } }); } if let Some(stderr) = child.stderr.take() { let tx_out = tx.clone(); std::thread::spawn(move || { let reader = BufReader::new(stderr); for line in reader.lines().flatten() { let _ = tx_out.send(ToolMessage::Output(line)); } }); } let status = child.wait()?; if !status.success() { let _ = tx.send(ToolMessage::Output(format!("error creating distrobox!"))); } self.created = true; return Ok(()); } } #[derive(Clone)] struct Prompt { action: Option, responses: Vec, execute_command: String, num_responses: usize, }