use crate::get_user_input; use crate::lib::Destination; use crate::lib::Message; use crate::lib::Project; use crate::lib::Table; use crate::print_error; use crate::print_success; use dns_lookup::lookup_host; use std::collections::HashMap; use std::fmt::Arguments; use std::fs::create_dir; use std::fs::create_dir_all; use std::fs::{File, OpenOptions, ReadDir, read_to_string}; use std::io::Write; use std::path::PathBuf; use std::process::Output; use std::process::exit; use std::result; use std::thread::sleep; use std::time::Duration; use tokio; use tokio::spawn; use tokio::sync::mpsc::{Receiver, Sender, channel}; #[derive(Clone)] pub struct ToolCommand { pub name: String, pub help: String, pub req_args: Vec, pub user_args: Vec, pub console_tx: Sender, pub optional_args: Vec, pub args: Option>, pub interactive: bool, pub optionally_interactive: bool, pub func: fn( Option>, Sender, Option>, Option>, ), } impl ToolCommand { pub fn new( name: String, help: String, tx: Sender, func: fn( Option>, Sender, Option>, Option>, ), ) -> Self { Self { name, help, req_args: Vec::new(), user_args: Vec::new(), console_tx: tx, optional_args: Vec::new(), interactive: false, optionally_interactive: false, args: None, func: func, } } pub async fn execute( mut self, execute_rx: Option>, runtime: tokio::runtime::Handle, ) { if self.interactive { let message = Message { source: Destination::Console, destination: Destination::Control, content: String::from("interactive"), }; self.console_tx.send(message).await.unwrap(); let (command_tx, mut command_rx) = channel(1); let console_tx = self.console_tx.clone(); runtime.spawn_blocking(move || { (self.func)(self.args, self.console_tx, Some(command_tx), execute_rx) }); loop { let rx_res = command_rx.try_recv(); if rx_res.is_ok() { let message = rx_res.unwrap(); match message.content.as_str() { "noninteractive" => { self.interactive = false; console_tx.send(message).await.unwrap(); } "finished" => { console_tx.send(message).await.unwrap(); break; } _ => {} } } } } else { runtime.spawn(async move { (self.func)(self.args, self.console_tx, None, None) }); } } } #[derive(Clone, Default)] pub struct ToolArgument { pub name: String, pub user_supplied: bool, pub optional: bool, pub position: Option, pub path: Option, pub string: Option, pub project: Option, pub projects: Option>, pub boolean: Option, pub console_tx: Option>, pub tx: Option>, } pub fn build_tools(tx: Sender) -> Vec { let mut tool_commands = Vec::new(); let mut list_projects = ToolCommand::new( "list_projects".to_string(), "lists the currently tracked projects".to_string(), tx.clone(), list_projects, ); list_projects.req_args = vec![String::from("projects")]; list_projects.interactive = false; tool_commands.push(list_projects); let mut new_project_command = ToolCommand::new("new_project".to_string(), "creates a new project. you can give it a project name as an argument, or it will prompt you for the name.".to_string(), tx.clone(), new_project); new_project_command.req_args = vec![ String::from("templatebox"), String::from("upcoming_notes"), String::from("upcoming_files"), String::from("config"), ]; new_project_command.optional_args = vec![String::from("name")]; new_project_command.optionally_interactive = true; new_project_command.req_args = vec![ String::from("config"), String::from("upcoming_files"), String::from("upcoming_notes"), String::from("templatebox"), ]; tool_commands.push(new_project_command); let mut promote_project_command = ToolCommand::new( "promote_project".to_string(), "promote a project to be promoted from upcoming to curent. Optionally takes a name= argument and a home= argument.".to_string(), tx.clone(), promote_project, ); promote_project_command.req_args = vec![ String::from("projects"), String::from("current_files"), String::from("current_notes"), String::from("templatebox"), ]; promote_project_command.optionally_interactive = true; promote_project_command.optional_args = vec![String::from("name"), String::from("home")]; tool_commands.push(promote_project_command); return tool_commands; } pub fn build_args( projects: &Vec, config: &PathBuf, settings: &HashMap, ) -> Vec { let mut args = Vec::new(); for project in projects { if project.active { let mut project_arg = ToolArgument::default(); project_arg.name = String::from("active"); project_arg.project = Some(project.clone()); args.push(project_arg); } } let mut projects_arg = ToolArgument::default(); projects_arg.name = String::from("projects"); projects_arg.projects = Some(projects.clone()); args.push(projects_arg); let mut config_arg = ToolArgument::default(); config_arg.name = String::from("config"); config_arg.path = Some(config.clone()); args.push(config_arg); for setting in settings.keys() { match setting.as_str() { "templatebox" => { let mut template_arg = ToolArgument::default(); template_arg.name = String::from("templatebox"); let boxname = settings.get("templatebox").unwrap().clone(); template_arg.string = Some(boxname); args.push(template_arg); } "current_notes" => { let mut current_note_arg = ToolArgument::default(); current_note_arg.name = String::from("current_notes"); let path = PathBuf::from(settings.get("current_notes").unwrap()); current_note_arg.path = Some(path); args.push(current_note_arg); } "current_files" => { let mut new_arg = ToolArgument::default(); new_arg.name = String::from("current_files"); let path = PathBuf::from(settings.get("current_files").unwrap()); new_arg.path = Some(path); args.push(new_arg); } "upcoming_files" => { let mut new_arg = ToolArgument::default(); new_arg.name = String::from("upcoming_files"); let path = PathBuf::from(settings.get("upcoming_files").unwrap()); new_arg.path = Some(path); args.push(new_arg); } "upcoming_notes" => { let mut new_arg = ToolArgument::default(); new_arg.name = String::from("upcoming_notes"); let path = PathBuf::from(settings.get("upcoming_notes").unwrap()); new_arg.path = Some(path); args.push(new_arg); } "tools" => { let mut new_arg = ToolArgument::default(); new_arg.name = String::from("tools"); let path = PathBuf::from(settings.get("tools").unwrap()); new_arg.path = Some(path); args.push(new_arg); } "distrobox" => { let mut new_arg = ToolArgument::default(); new_arg.name = String::from("distrobox"); if settings.get("distrobox").unwrap().contains("yes") { new_arg.boolean = Some(true); } else { new_arg.boolean = Some(false); } args.push(new_arg); } _ => {} } } return args; } pub async fn send_command_output(tx: Sender, message: Message) { tx.send(message.clone()).await.unwrap(); } pub fn list_projects( args: Option>, tx: Sender, _command_tx: Option>, _rx: Option>, ) { let given_args = args.unwrap(); let mut lines = vec![String::from("name|stage|boxname")]; for arg in given_args { if arg.name == String::from("projects") { for project in arg.projects.unwrap() { if project.current { let line = format!("{}|{}|{}", project.name, "current", project.boxname); lines.push(line); } else { let line = format!("{}|{}|{}", project.name, "upcomming", project.boxname); lines.push(line); } } } } let mut table = Table::default(); table.build(lines); let message = Message { source: Destination::Console, destination: Destination::Console, content: table.get_table(), }; tokio::spawn(send_command_output(tx, message.clone())); } pub fn new_project( args: Option>, tx: Sender, comand_tx: Option>, rx: Option>, ) { let mut interactive = false; let given_args = args.unwrap(); let mut config_path = PathBuf::new(); let mut name = String::new(); let mut files_path = PathBuf::new(); let mut notes_path = PathBuf::new(); let mut template_box = String::new(); let non_interactive_message = Message { source: Destination::Control, destination: Destination::Console, content: String::from("noninteractive"), }; let mut interactive_message = Message { source: Destination::Control, destination: Destination::Console, content: String::from("interactive"), }; for arg in given_args { match arg.name.as_str() { "config" => { config_path = arg.path.unwrap(); } "name" => { name = arg.string.unwrap(); } "upcoming_files" => { files_path = arg.path.unwrap(); } "upcoming_notes" => { notes_path = arg.path.unwrap(); } "template" => { template_box = arg.string.unwrap(); } _ => {} } } if name.len() == 0 { interactive = true; } if interactive { let message = Message { source: Destination::Console, destination: Destination::Console, content: String::from("project_name?"), }; let mut rx = initialize_interactive(rx, tx.clone()); let mut name_asked = false; loop { if !name_asked { tx.blocking_send(message.clone()).unwrap(); } let response = rx.blocking_recv().unwrap(); name = response.content; break; } comand_tx .clone() .unwrap() .blocking_send(non_interactive_message.clone()) .unwrap(); } let mut project_path = config_path.clone(); project_path.pop(); project_path.push("projects"); project_path.push(format!("{}.conf", &name)); let file_create_res = File::create(&project_path); if file_create_res.is_err() { let message = Message { source: Destination::Console, destination: Destination::Console, content: format!("error! unable to create {}", &project_path.display()), }; tokio::spawn(send_command_output(tx.clone(), message)); if interactive { comand_tx .clone() .unwrap() .blocking_send(non_interactive_message.clone()) .unwrap(); } return; } files_path.push(&name); notes_path.push(&name); let files_dir_res = create_dir_all(&files_path); let notes_dir_res = create_dir_all(¬es_path); if files_dir_res.is_err() { let message = Message { source: Destination::Console, destination: Destination::Console, content: format!( "Error failure to create project files folder!\n{}", files_dir_res.err().unwrap() ), }; tokio::spawn(send_command_output(tx.clone(), message)); if interactive { comand_tx .clone() .unwrap() .blocking_send(non_interactive_message.clone()) .unwrap(); } return; } if notes_dir_res.is_err() { let message = Message { source: Destination::Console, destination: Destination::Console, content: format!( "Error failure to create project files folder!\n{}", files_dir_res.err().unwrap() ), }; tokio::spawn(send_command_output(tx.clone(), message)); if interactive { comand_tx .clone() .unwrap() .blocking_send(non_interactive_message.clone()) .unwrap(); } return; } let mut new_project = Project::default(); new_project.name = name.clone(); new_project.files = files_path; new_project.notes = notes_path; new_project.current = false; new_project.boxname = format!("{}_{}", template_box, name); new_project.config = project_path; println!("{}", new_project.config.display()); new_project.save_project(); print_success("folder structure and config file created successfully!"); println!("setting up default notes..."); let res = new_project.generate_default_notes(&config_path); let message = Message { source: Destination::Console, destination: Destination::Console, content: res, }; tokio::spawn(send_command_output(tx.clone(), message)); let message = Message { source: Destination::Console, destination: Destination::Console, content: String::from("finished"), }; tokio::spawn(send_command_output(tx.clone(), message)); } pub fn promote_project( args: Option>, tx: Sender, comand_tx: Option>, rx: Option>, ) { let mut project = String::new(); let mut projects = Vec::new(); let mut files = PathBuf::new(); let mut notes = PathBuf::new(); let mut tools = PathBuf::new(); let mut template = String::new(); let mut home = None; let mut given_args = Vec::new(); let mut interactive = true; if args.is_some() { given_args = args.unwrap(); } for arg in given_args { match arg.name.as_str() { "name" => { project = arg.string.clone().unwrap(); interactive = false; } "projects" => projects = arg.projects.clone().unwrap(), "current_files" => files = arg.path.unwrap(), "current_notes" => notes = arg.path.unwrap(), "templatebox" => template = arg.string.unwrap(), "home" => home = arg.path, "tools" => tools = arg.path.unwrap(), _ => {} } } if interactive { let rx = initialize_interactive(rx, tx.clone()); let mut lines = vec![String::from("id|name|status")]; let mut id = 0; for project in &projects { let mut line = format!("{}|", id); id += 1; line.push_str(&project.name); if project.current { line.push_str("|current"); } else { line.push_str("|upcoming"); } lines.push(line); } let mut project_table = Table::default(); project_table.build(lines); let projects_string = project_table.get_table(); let table_message = Message { source: Destination::Console, destination: Destination::Console, content: projects_string, }; tx.blocking_send(table_message).unwrap(); let (gotten_name, mut rx) = prompt_interactive(rx, tx.clone(), "project to promote?"); let selection: usize = gotten_name.parse().unwrap(); project = format!("{}", &projects[selection].name); } for mut existing_project in projects { if existing_project.name == project { let promote_res = existing_project.promote_project( &files, ¬es, template.clone(), &tools, home.clone(), ); tx.blocking_send(Message { source: Destination::Console, destination: Destination::Console, content: promote_res, }) .unwrap(); } } deinitialize_interactive(tx.clone()); } pub fn remove_project(args: Option>) -> String { let mut project = Project::default(); let mut projects = Vec::new(); if args.is_some() { let given_args = args.unwrap(); for arg in given_args { match arg.name.as_str() { "project" => project = arg.project.unwrap(), "projects" => projects = arg.projects.unwrap(), _ => {} } } } if !project.name.len() > 1 { let mut current_index = 0; for project in &projects { println!("{}.) {}", current_index, project.name); current_index += 1; } let response: usize = get_user_input("project to remove?").parse().unwrap(); if response >= projects.len() { return format!("error invalid project selection!"); } project = projects[response].clone() } return project.remove_project(); } pub fn subdomain_brute(args: Option>) -> String { let mut result_string = String::from("subdmain bruteforce failed!!"); let mut target_domain = String::new(); let mut dict = PathBuf::new(); let mut project = Project::default(); if args.is_some() { let iargs = args.unwrap(); for arg in iargs { if arg.user_supplied { match arg.name.as_str() { "wordlist" => dict = PathBuf::from(arg.string.unwrap()), "domain" => target_domain = arg.string.unwrap(), _ => {} } } else if arg.name == "project" { project = arg.project.unwrap(); } } } if target_domain.len() > 1 { target_domain = get_user_input("target domain to bruteforce?"); } if !dict.exists() { dict = PathBuf::from(get_user_input("path to wordlist?")); } print_success("arguments parsed successfully!"); println!("starting bruteforce..."); let wordlist_read_res = read_to_string(dict); if wordlist_read_res.is_err() { print_error( "error reading wordlist!", Some(wordlist_read_res.err().unwrap().to_string()), ); return result_string; } let wordlist = wordlist_read_res.unwrap(); let mut results = Vec::new(); for sub in wordlist.lines() { let mut found_ips = Vec::new(); let domain_name = format!("{}.{}", sub, target_domain); let lookup_res = lookup_host(&domain_name); if lookup_res.is_ok() { let ips = lookup_res.unwrap(); for ip in ips { found_ips.push(ip); } } if found_ips.len() > 0 { let mut ip_string = String::new(); for ip in found_ips { ip_string.push_str(&format!("{},", ip)); } results.push(format!("{} | {}", domain_name, ip_string)); } } let mut enumeration_path = project.notes.clone(); enumeration_path.push("enumeration.md"); let note_open_res = OpenOptions::new() .create_new(true) .append(true) .open(enumeration_path); if note_open_res.is_err() { print_error( "error opening enumeration notes file!", Some(note_open_res.err().unwrap().to_string()), ); } else { let mut note_file = note_open_res.unwrap(); let mut note_text = String::from("\n\n# Subdomain Bruteforce\n"); note_text.push_str(&format!("## {}", target_domain)); note_text.push_str("\n| domain name | IP Addresses |\n"); note_text.push_str("| ----------- | ------------ |\n"); for res in results { note_text.push_str(&format!("| {} |\n", res)); } let write_res = write!(note_file, "{}", note_text); if write_res.is_ok() { write_res.unwrap(); result_string = String::from("Subdomain bruteforcing completed successfully!"); print_success("Subdomain brueforcing completed successfully!"); } } return result_string; } pub fn activate_project(args: Option>) -> String { let mut result_string = String::from("Failed to activeate project!"); let mut projects = Vec::new(); let mut activate_target = String::new(); for arg in args.unwrap() { if arg.user_supplied { if arg.name == "project" { activate_target = arg.string.unwrap(); } } match arg.name.as_str() { "projects" => projects = arg.projects.unwrap().clone(), _ => {} } } for mut project in projects { if project.name == activate_target { project.active = true; result_string = String::from("Project Activated Successfully!"); } else { project.active = false; } project.save_project(); } return result_string; } pub fn initialize_interactive( rx: Option>, tx: Sender, ) -> Receiver { let mut rx = rx.unwrap(); let mut acked = false; loop { let response = rx.blocking_recv().unwrap(); match response.content.as_str() { "ack" => { if !acked { tx.blocking_send(response.clone()).unwrap(); acked = true; } } "ready" => { return rx; } _ => {} } } } pub fn prompt_interactive( mut rx: Receiver, tx: Sender, prompt: &str, ) -> (String, Receiver) { let message = Message { source: Destination::Console, destination: Destination::Console, content: String::from(prompt), }; tx.blocking_send(message); let mut return_string = String::new(); loop { let rx_res = rx.try_recv(); if rx_res.is_ok() { let response = rx_res.unwrap(); return_string = response.content; break; } } return (return_string, rx); } pub fn deinitialize_interactive(tx: Sender) { tx.blocking_send(Message { source: Destination::Console, destination: Destination::Console, content: String::from("noninteractive"), }) .unwrap(); }