diff --git a/Cargo.lock b/Cargo.lock index ce33459..209990f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -190,6 +190,18 @@ dependencies = [ "typenum", ] +[[package]] +name = "dns-lookup" +version = "3.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e39034cee21a2f5bbb66ba0e3689819c4bb5d00382a282006e802a7ffa6c41d" +dependencies = [ + "cfg-if", + "libc", + "socket2", + "windows-sys 0.60.2", +] + [[package]] name = "generic-array" version = "0.14.9" @@ -217,6 +229,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "inout" version = "0.1.4" @@ -258,6 +276,16 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "once_cell_polyfill" version = "1.70.2" @@ -416,6 +444,8 @@ dependencies = [ "chacha20poly1305", "clap", "colored", + "dns-lookup", + "num_cpus", "tokio", "walkdir", ] diff --git a/Cargo.toml b/Cargo.toml index db871dc..8596eb5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,5 +7,7 @@ edition = "2024" chacha20poly1305 = "0.10.1" clap = { version = "4.5.51", features = ["derive"] } colored = "3.0.0" +dns-lookup = "3.0.1" +num_cpus = "1.17.0" tokio = { version = "1.48.0", features = ["full"] } walkdir = "2.5.0" diff --git a/src/cli.rs b/src/cli.rs index 2553b5b..4b333b0 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,184 +1,294 @@ -use crate::commands; -use crate::commands::ToolArgument; -use crate::commands::ToolCommand; -use crate::lib::Message; -use crate::load_projects; -use crate::load_settings; -use crate::network; -use crate::print_error; -use crate::{get_user_input, lib::Destination, lib::Project, lib::Table, print_success}; -use std::io::Read; -use std::path::PathBuf; -use tokio; -use tokio::sync::mpsc::{Receiver, Sender, channel}; +use crate::{ + commands::{ToolArgument, build_args, build_tools}, + lib::{Destination, Message, Project}, + load_projects, load_settings, print_error, print_success, +}; +use std::{path::PathBuf, thread::sleep, time::Duration}; +use tokio::{ + self, + io::AsyncBufReadExt, + sync::mpsc::{Receiver, Sender, channel}, +}; -pub async fn rec_message(mut rx: Receiver) { +pub async fn rec_message(mut console_rx: Receiver, cli_tx: Sender) { + print_success("Console output initialized!"); + let prompt = String::from("what is thy bidding my master?"); + let prompt_message = Message { + source: Destination::Console, + destination: Destination::Console, + content: String::from("PROMPT"), + }; let mut display = true; + let mut exit = false; + let mut interactive = false; + let mut acked = false; loop { - let rx_res = rx.try_recv(); + let mut output = String::new(); + let rx_res = console_rx.try_recv(); if rx_res.is_ok() { let message = rx_res.unwrap(); - if message.content.to_lowercase().contains("error") { - print_error(&message.content, None); - } else { - print_success(&message.content); + if message.source == Destination::Control { + match message.content.as_str() { + "init" => { + sleep(Duration::from_secs(5)); + output.push_str("\nconsole output started successfully!\n"); + display = true; + } + "exit" => { + output.push_str("\ngood bye!\n"); + display = true; + exit = true; + } + "status" => { + if console_rx.is_closed() { + output.push_str("\nerror! the console_rx channel is closed!!!\n"); + display = true; + } else { + output.push_str("\nthe console_rx channel is still open!!\n"); + display = true; + } + } + "noninteractive" => { + println!("we got the non-interactive message!"); + interactive = false; + display = true; + let done_message = Message { + source: Destination::Console, + destination: Destination::Console, + content: String::from("DONE"), + }; + cli_tx.send(done_message).await.unwrap(); + } + "ack" => { + display = false; + interactive = true; + cli_tx.send(message.clone()).await.unwrap(); + acked = true; + } + _ => { + output.push_str(&format!("{}", message.content)); + display = true; + } + } + } else if message.source == Destination::Server { + output.push_str(&format!( + "\n{} recieved from server! (functionality coming soon!)\n", + message.content + )); + display = true; + } else if message.source == Destination::Console { + output.push_str(&format!("\n{}\n", &message.content)); + display = true; } - display = true; } if display { - println!("\n\ncommand?"); - display = false; + match interactive { + false => { + if !exit { + output.push_str(&prompt); + if output.contains("error") { + print_error(&output, None); + } else { + print_success(&output); + } + } + } + true => { + if acked { + cli_tx.send(prompt_message.clone()).await.unwrap(); + println!("{}", &output); + } + } + } + } + display = false; + if exit { + break; } } } +pub async fn console_user_input() -> String { + let mut reader = tokio::io::BufReader::new(tokio::io::stdin()); + let mut response = Vec::new(); + let _fut = reader.read_until(b'\n', &mut response).await.unwrap(); + let user_input = match str::from_utf8(&response) { + Ok(s) => s.to_string(), + Err(_e) => "".to_string(), + }; + return user_input; +} + pub async fn cli( - mut projects: Vec, + console_tx: Sender, + console_rx: Receiver, main_tx: Sender, - tool_tx: Sender, - server_address: String, - config_path: PathBuf, + config: PathBuf, + runtime: tokio::runtime::Handle, ) { + let (cli_tx, mut cli_rx) = channel(1); + let handle = runtime.spawn(rec_message(console_rx, cli_tx)); print_success("started the CLI!"); - let mut commands = commands::build_tools(); + print_success("happy hacking!"); loop { - let mut valid_command = false; - let settings = load_settings(&config_path, false); - projects = load_projects(&config_path, false); - if tool_tx.is_closed() { - println!("tool tx closed at the start of this loop!"); - } - let command = get_user_input(""); - let mut command_name = String::new(); - let mut args = Vec::new(); - if command.contains(" ") { - let mut command_vec: Vec<&str> = command.split(" ").collect(); - command_name = command_vec[0].to_string(); - for arg in &mut command_vec[1..] { - args.push(arg.to_string()); - } + let input_handle = tokio::spawn(console_user_input()); + let mut settings = load_settings(&config, false); + let mut projects = load_projects(&config, false); + let tool_args = build_args(&projects, &config, &settings); + let tool_commands = build_tools(console_tx.clone()); + let user_input = input_handle.await.unwrap().trim().to_string(); + let mut user_command_name = String::new(); + let mut user_command_args = Vec::new(); + if user_input.contains(" ") { + let input_vec: Vec<&str> = user_input.split(" ").collect(); + user_command_name = input_vec[0].to_string(); + user_command_args = input_vec[1..].to_vec(); } else { - command_name = command; + user_command_name = user_input; } - if command_name == String::from("exit") { + if user_command_name == String::from("exit") { + let message = Message { + source: Destination::Control, + destination: Destination::Console, + content: String::from("exit"), + }; + console_tx.send(message).await.unwrap(); + handle.abort(); let message = Message { source: Destination::Console, destination: Destination::Control, content: String::from("exit"), }; - valid_command = true; main_tx.send(message).await.unwrap(); + break; } - for toolcommand in &mut commands { - let mut ready = true; - if toolcommand.name == command_name { + let mut valid_command = false; + let mut command_to_run = tool_commands[0].clone(); + for command in &tool_commands { + if command.name == user_command_name { valid_command = true; - if toolcommand.req_args.len() > 0 { - let mut args_vec = Vec::new(); - for req_arg in toolcommand.req_args.clone() { - let mut new_arg = ToolArgument::default(); - new_arg.name = req_arg; - args_vec.push(new_arg); - } - if args.len() < toolcommand.user_args.len() { - ready = false; - print_error("not enough arguments provided!", None); - let mut command_usage = format!("{}", toolcommand.name.clone()); - for arg in toolcommand.user_args.clone() { - command_usage.push_str(&format!(" {}", &arg)); - } - print_error("usage:", Some(command_usage)); - } - if ready { - let mut position_count = 0; - for user_arg in toolcommand.user_args.clone() { - println!("enough args supplied, building arguments list"); - let mut new_arg = ToolArgument::default(); - new_arg.name = user_arg; - new_arg.user_supplied = true; - new_arg.position = Some(position_count); - args_vec.push(new_arg); - position_count += 1; - } - for given_arg in &args { - if given_arg.contains("=") { - for optional_arg in toolcommand.optional_args.clone() { - let arg_vec: Vec<&str> = given_arg.split("=").collect(); - if arg_vec[0].to_string() == optional_arg { - println!("arg found {}", optional_arg); - let mut new_arg = ToolArgument::default(); - new_arg.name = optional_arg; - new_arg.string = Some(arg_vec[1].trim().to_string()); - new_arg.optional = true; - args_vec.push(new_arg); - } - } - } - } - } - if ready { - for arg in &mut args_vec { - let iargs = args.clone(); - if arg.user_supplied { - arg.string = Some(iargs[arg.position.unwrap()].clone()); - } else if !arg.optional { - match arg.name.as_str() { - "projects" => arg.projects = Some(projects.clone()), - "upcoming_files" => { - arg.path = - Some(PathBuf::from(settings["upcoming_files"].clone())) - } - "upcoming_notes" => { - arg.path = - Some(PathBuf::from(settings["upcoming_notes"].clone())) - } - "files" => { - arg.path = - Some(PathBuf::from(settings["current_files"].clone())) - } - "notes" => { - arg.path = - Some(PathBuf::from(settings["current_notes"].clone())) - } - "tools" => { - arg.path = Some(PathBuf::from(settings["tools"].clone())) - } - "template" => { - arg.string = - Some(String::from(settings["templatebox"].clone())) - } - "config" => arg.path = Some(config_path.clone()), - "terminal" => arg.string = Some(settings["terminal"].clone()), - _ => print_error( - &format!("unkown arg requested! {}", arg.name), - None, - ), - } - } - } - toolcommand.args = Some(args_vec); - } - } - let mut message = Message { - source: Destination::Console, - destination: Destination::Console, - content: String::new(), - }; - if ready { - message.content = toolcommand.execute(); - } else { - message.content = String::from("error in command!"); - } - tool_tx.send(message).await.unwrap(); + command_to_run = command.clone(); } } - if !valid_command { + if valid_command == false { let message = Message { source: Destination::Console, destination: Destination::Console, - content: String::from("error: command not found!"), + content: String::from("error! invalid command!"), }; - tool_tx.send(message).await.unwrap(); + console_tx.send(message).await.unwrap(); + continue; + } + let mut command_to_run_arg_vec = Vec::new(); + for arg in tool_args { + if command_to_run.req_args.contains(&arg.name) { + command_to_run_arg_vec.push(arg.clone()); + println!("{} added!", &arg.name); + } + } + let mut correct_args = false; + for arg in user_command_args.clone() { + if !arg.contains("=") { + let mut new_arg = ToolArgument::default(); + new_arg.name = String::from(arg); + new_arg.string = Some(String::from(arg)); + command_to_run_arg_vec.push(new_arg); + } + } + if command_to_run_arg_vec.len() + == command_to_run.req_args.len() + command_to_run.user_args.len() + { + correct_args = true; + } + for arg in &user_command_args { + if arg.contains("=") { + let arg_vec: Vec<&str> = arg.split("=").collect(); + let mut new_arg = ToolArgument::default(); + new_arg.name = String::from(arg_vec[0]); + new_arg.string = Some(String::from(arg_vec[1])); + command_to_run_arg_vec.push(new_arg); + } + } + if correct_args == false { + println!("{}", command_to_run.help); + print_error( + "wrong number of arguments supplied!\n please read the above help message.", + None, + ); + print_error( + "", + Some(format!( + "args reqired: {}\nargs given: {}", + command_to_run.req_args.len() + command_to_run.user_args.len(), + command_to_run_arg_vec.len() + )), + ); + continue; + } + if command_to_run_arg_vec.len() > 0 { + command_to_run.args = Some(command_to_run_arg_vec.clone()); + } + if command_to_run.optionally_interactive { + if command_to_run_arg_vec.len() + == command_to_run.req_args.len() + + command_to_run.user_args.len() + + command_to_run.optional_args.len() + { + command_to_run.interactive = false; + } else { + command_to_run.interactive = true; + } + } + if command_to_run.interactive { + println!("we got to the interactive section!"); + let (command_tx, command_rx) = channel(1); + runtime.spawn(command_to_run.execute(Some(command_rx), runtime.clone())); + let init_message = Message { + source: Destination::Control, + destination: Destination::Console, + content: String::from("ack"), + }; + let mut inited = false; + loop { + if !inited { + command_tx.send(init_message.clone()).await.unwrap(); + } + let rx_res = cli_rx.try_recv(); + if rx_res.is_ok() { + let message = rx_res.unwrap(); + match message.content.as_str() { + "ack" => { + inited = true; + command_tx + .send(Message { + source: Destination::Console, + destination: Destination::Console, + content: String::from("ready"), + }) + .await + .unwrap(); + } + "PROMPT" => { + let input_handle = tokio::spawn(console_user_input()); + let response = input_handle.await.unwrap(); + command_tx + .send(Message { + source: Destination::Console, + destination: Destination::Console, + content: response.trim().to_string(), + }) + .await + .unwrap(); + } + "DONE" => { + break; + } + _ => {} + } + } + } + } else { + runtime.spawn(command_to_run.execute(None, runtime.clone())); } } } diff --git a/src/commands.rs b/src/commands.rs index ce9cbd6..4f63d87 100644 --- a/src/commands.rs +++ b/src/commands.rs @@ -5,14 +5,22 @@ 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::sync::mpsc::{Receiver, Sender}; +use tokio::spawn; +use tokio::sync::mpsc::{Receiver, Sender, channel}; #[derive(Clone)] pub struct ToolCommand { @@ -20,38 +28,82 @@ pub struct ToolCommand { pub help: String, pub req_args: Vec, pub user_args: Vec, + pub console_tx: Sender, pub optional_args: Vec, pub args: Option>, - pub func: fn(Option>) -> String, + pub interactive: bool, + pub optionally_interactive: bool, + pub func: fn( + Option>, + Sender, + Option>, + Option>, + ), } impl ToolCommand { - pub fn new(name: String, help: String, func: fn(Option>) -> String) -> Self { + 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: func, } } - pub fn execute(&self) -> String { - if self.req_args.len() > 0 { - if self.args.is_none() { - return String::from("Error: no arguments given, but arguments are required!"); - } else { - let min_args = self.req_args.len() + self.user_args.len(); - let max_args = min_args + self.optional_args.len(); - let args_provided = self.args.clone().unwrap().len(); - if args_provided > max_args || args_provided < min_args { - return String::from("Error: the wrong number of args were supplied!"); + 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) }); } - return (self.func)(self.args.clone()); } } @@ -66,62 +118,147 @@ pub struct ToolArgument { pub project: Option, pub projects: Option>, pub boolean: Option, + pub console_tx: Option>, + pub tx: Option>, } -pub fn build_tools() -> Vec { +pub fn build_tools(tx: Sender) -> Vec { let mut tool_commands = Vec::new(); - let mut listproject = ToolCommand::new( + let mut list_projects = ToolCommand::new( "list_projects".to_string(), - "coming soon".to_string(), + "lists the currently tracked projects".to_string(), + tx.clone(), list_projects, ); - listproject.req_args = vec![String::from("projects")]; - tool_commands.push(listproject); - let mut createproject = ToolCommand::new( - "create_project".to_string(), - "creates a new project and saves it to the projects file to be reloaded on the next loop - usage: - create_project project_name" - .to_string(), - new_project, - ); - createproject.req_args = vec![ + 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("template"), + String::from("templatebox"), ]; - createproject.user_args = vec![String::from("name")]; - tool_commands.push(createproject); - let mut promoteproject = ToolCommand::new( + tool_commands.push(new_project_command); + let mut promote_project_command = ToolCommand::new( "promote_project".to_string(), - "promotes a project from upcoming to current, and sets up the distrobox. -Optional arguments: home=/path/to/distrobox/home name=project_name_to_promote" - .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, ); - promoteproject.req_args = vec![ + promote_project_command.req_args = vec![ String::from("projects"), - String::from("tools"), - String::from("files"), - String::from("notes"), - String::from("template"), + String::from("current_files"), + String::from("current_notes"), + String::from("templatebox"), ]; - promoteproject.optional_args = vec![String::from("home"), String::from("name")]; - tool_commands.push(promoteproject); - let mut removeproject = ToolCommand::new( - "remove_project".to_string(), - "removes a project from the tool, deletes the files and notes, and deletes the distrobox." - .to_string(), - remove_project, - ); - removeproject.optional_args = vec!["project".to_string()]; - removeproject.req_args = vec!["projects".to_string()]; - tool_commands.push(removeproject); + 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 list_projects(args: Option>) -> String { +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 { @@ -139,16 +276,37 @@ pub fn list_projects(args: Option>) -> String { } let mut table = Table::default(); table.build(lines); - return table.get_table(); + 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>) -> String { +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" => { @@ -170,40 +328,92 @@ pub fn new_project(args: Option>) -> String { } } if name.len() == 0 { - return String::from( - "usage: newproject projectname\nprevious command was missing project name!", - ); + 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(); } - println!("gatered the data! name: {}", name); - print_success("arguments parsed correctly!"); - println!("setting up files..."); 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() { - return format!( - "Error failure to create project config file!\n{}\n{}", - file_create_res.err().unwrap(), - &project_path.display().to_string() - ); + 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() { - return (format!( - "Error failure to create project files folder!\n{}", - files_dir_res.err().unwrap() - )); + 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() { - return (format!( - "Error failure to create project files folder!\n{}", - files_dir_res.err().unwrap() - )); + 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(); @@ -216,10 +426,27 @@ pub fn new_project(args: Option>) -> String { new_project.save_project(); print_success("folder structure and config file created successfully!"); println!("setting up default notes..."); - return new_project.generate_default_notes(&config_path); + 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>) -> String { +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(); @@ -228,37 +455,52 @@ pub fn promote_project(args: Option>) -> String { 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(), + "name" => { + project = arg.string.clone().unwrap(); + interactive = false; + } "projects" => projects = arg.projects.clone().unwrap(), - "files" => files = arg.path.unwrap(), - "notes" => notes = arg.path.unwrap(), - "template" => template = arg.string.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(), _ => {} } } - let project_set = project.len() > 1; - if !project_set { - println!("{} : {}", project, project.len()); - println!("{}", project.len() > 1); - let mut current_index = 0; - for existing_project in &projects { - println!("{}.) {}", current_index, existing_project.name); - current_index += 1; + 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 response: usize = get_user_input("which project would you like to promote?") - .parse() - .unwrap(); - if response >= projects.len() { - return format!("error invalid project selection!"); - } - project = projects[response].name.clone(); + 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 { @@ -269,12 +511,15 @@ pub fn promote_project(args: Option>) -> String { &tools, home.clone(), ); - if !promote_res.to_lowercase().contains("success") { - return format!("Error promoting project!\n{}", promote_res); - } + tx.blocking_send(Message { + source: Destination::Console, + destination: Destination::Console, + content: promote_res, + }) + .unwrap(); } } - return String::from("Success!"); + deinitialize_interactive(tx.clone()); } pub fn remove_project(args: Option>) -> String { @@ -304,3 +549,170 @@ pub fn remove_project(args: Option>) -> String { } 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(); +} diff --git a/src/lib.rs b/src/lib.rs index 8658c73..88b5d94 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -138,6 +138,7 @@ pub struct Project { pub files: PathBuf, pub notes: PathBuf, pub current: bool, + pub active: bool, pub boxname: String, pub config: PathBuf, } @@ -288,10 +289,10 @@ impl Project { let mut conf_file = conf_open_create_res.unwrap(); let mut config_string = format!( "name|{}\nfiles|{}\nnotes|{}\nboxname|{}\nconfig|{}\n", - self.name, + self.name.trim(), self.files.display(), self.notes.display(), - self.boxname, + self.boxname.trim(), self.config.display(), ); if self.current { @@ -352,6 +353,7 @@ impl Project { .arg("stop") .arg("--root") .arg(&self.boxname) + .arg("--yes") .status(); if pdb_stop_res.is_err() { return format!( @@ -363,6 +365,7 @@ impl Project { .arg("stop") .arg("--root") .arg(&template) + .arg("--yes") .status(); if tdb_stop_res.is_err() { return format!( @@ -376,6 +379,7 @@ impl Project { .arg("rm") .arg("--root") .arg(&self.boxname) + .arg("--yes") .status(); if db_remove_res.is_err() { return format!( @@ -395,6 +399,7 @@ impl Project { .arg(format!("{}:/tools:rw", tools.display())) .arg("--name") .arg(&self.boxname) + .arg("--yes") .status(); if db_create_res.is_err() { return format!( diff --git a/src/main.rs b/src/main.rs index 9a99cd1..8851226 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,6 +4,7 @@ use chacha20poly1305::aead::{Aead, AeadCore, KeyInit, OsRng}; use chacha20poly1305::{ChaCha20Poly1305, Key}; use clap::Parser; use colored::Colorize; +use num_cpus; use std::collections::HashMap; use std::fs::{self, File, OpenOptions, read_dir, read_to_string}; use std::io::{Read, Write}; @@ -14,6 +15,7 @@ use std::thread::sleep; use std::time::Duration; use std::{thread, time}; use tokio; +use tokio::io::AsyncBufReadExt; use tokio::sync::mpsc::{Receiver, Sender, channel}; mod cli; @@ -103,31 +105,33 @@ pub fn load_projects(path: &PathBuf, display: bool) -> Vec { if conf_string_res.is_ok() { let conf_string = conf_string_res.unwrap(); for line in conf_string.lines() { - let line_vec: Vec<&str> = line.split("|").collect(); - match line_vec[0] { - "name" => { - new_project.name = line_vec[1].trim().to_string(); - } - "stage" => { - if line_vec[1].contains("current") { - new_project.current = true; - } else { - new_project.current = false; + if line.len() > 2 { + let line_vec: Vec<&str> = line.split("|").collect(); + match line_vec[0] { + "name" => { + new_project.name = line_vec[1].trim().to_string(); + } + "stage" => { + if line_vec[1].contains("current") { + new_project.current = true; + } else { + new_project.current = false; + } + } + "files" => { + new_project.files = PathBuf::from(line_vec[1]); + } + "notes" => { + new_project.notes = PathBuf::from(line_vec[1]); + } + "boxname" => new_project.boxname = String::from(line_vec[1]), + "config" => new_project.config = PathBuf::from(line_vec[1]), + _ => { + print_error( + "unknown setting discoverd in project config file!", + Some(String::from(line_vec[0])), + ); } - } - "files" => { - new_project.files = PathBuf::from(line_vec[1]); - } - "notes" => { - new_project.notes = PathBuf::from(line_vec[1]); - } - "boxname" => new_project.boxname = String::from(line_vec[1]), - "config" => new_project.config = PathBuf::from(line_vec[1]), - _ => { - print_error( - "unknown setting discoverd in project config file!", - None, - ); } } } @@ -272,37 +276,35 @@ async fn main() { } let (main_tx, mut main_rx) = channel(1024); let (console_tx, console_rx) = channel(1024); - if !args.gui { - let input_handle = tokio::spawn(cli::cli( - projects, - main_tx.clone(), - console_tx.clone(), - server_address, - config_path, + let cpu_count = num_cpus::get_physical(); + let thread_count = cpu_count / 2; + let runtime = tokio::runtime::Builder::new_multi_thread() + .thread_name("root") + .worker_threads(thread_count / 2) + .max_blocking_threads(thread_count / 2) + .build() + .unwrap(); + if args.gui { + println!("coming soon!"); + } else { + let rt_handle = runtime.handle(); + let handle = runtime.spawn(cli::cli( + console_tx, + console_rx, + main_tx, + config_path.clone(), + rt_handle.clone(), )); - thread::sleep(Duration::from_secs(1)); - let output_handle = tokio::spawn(cli::rec_message(console_rx)); loop { - sleep(Duration::from_secs(1)); - let rx_rex = main_rx.try_recv(); - if rx_rex.is_ok() { - let message = rx_rex.unwrap(); - if message.destination == lib::Destination::Control { - match message.content.as_str() { - "exit" => { - input_handle.abort(); - output_handle.abort(); - exit(0); - } - _ => { - println!("unknown message recieved!"); - println!("{}", message.content); - } - } + let rx_res = main_rx.try_recv(); + if rx_res.is_ok() { + let message = rx_res.unwrap(); + if message.content == String::from("exit") { + handle.abort(); + runtime.shutdown_background(); + break; } } } - } else { - println!("gui coming soon!"); } }