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

re-wrote how console input and output is handled, this should lead to

Browse Source 
easier integration with a future gui, and makes everything nice and
fast... still tracking down a problem with distrobox creation and
figuring out how to handle interactive input for other processes.
This commit is contained in:
pyro57000
2025-12-03 16:48:48 -06:00
parent 5053113dc0
commit 921a2f5923
6 changed files with 861 additions and 300 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
Cargo.lock generated
View File

@@ -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",
]

2
Cargo.toml
View File

@@ -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"

406
src/cli.rs
View File

@@ -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<Message>) {
pub async fn rec_message(mut console_rx: Receiver<Message>, cli_tx: Sender<Message>) {
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;
}
if display {
println!("\n\ncommand?");
"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;
}
}
if display {
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<Project>,
console_tx: Sender<Message>,
console_rx: Receiver<Message>,
main_tx: Sender<Message>,
tool_tx: Sender<Message>,
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);
command_to_run = command.clone();
}
}
}
}
}
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();
}
}
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()));
}
}
}

602
src/commands.rs
View File

@@ -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,39 +28,83 @@ pub struct ToolCommand {
pub help: String,
pub req_args: Vec<String>,
pub user_args: Vec<String>,
pub console_tx: Sender<Message>,
pub optional_args: Vec<String>,
pub args: Option<Vec<ToolArgument>>,
pub func: fn(Option<Vec<ToolArgument>>) -> String,
pub interactive: bool,
pub optionally_interactive: bool,
pub func: fn(
Option<Vec<ToolArgument>>,
Sender<Message>,
Option<Sender<Message>>,
Option<Receiver<Message>>,
),
}
impl ToolCommand {
pub fn new(name: String, help: String, func: fn(Option<Vec<ToolArgument>>) -> String) -> Self {
pub fn new(
name: String,
help: String,
tx: Sender<Message>,
func: fn(
Option<Vec<ToolArgument>>,
Sender<Message>,
Option<Sender<Message>>,
Option<Receiver<Message>>,
),
) -> 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!");
pub async fn execute(
mut self,
execute_rx: Option<Receiver<Message>>,
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 {
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!");
runtime.spawn(async move { (self.func)(self.args, self.console_tx, None, None) });
}
}
}
return (self.func)(self.args.clone());
}
}
#[derive(Clone, Default)]
@@ -66,62 +118,147 @@ pub struct ToolArgument {
pub project: Option<Project>,
pub projects: Option<Vec<Project>>,
pub boolean: Option<bool>,
pub console_tx: Option<Sender<String>>,
pub tx: Option<Sender<String>>,
}
pub fn build_tools() -> Vec<ToolCommand> {
pub fn build_tools(tx: Sender<Message>) -> Vec<ToolCommand> {
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<Vec<ToolArgument>>) -> String {
pub fn build_args(
projects: &Vec<Project>,
config: &PathBuf,
settings: &HashMap<String, String>,
) -> Vec<ToolArgument> {
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: Message) {
tx.send(message.clone()).await.unwrap();
}
pub fn list_projects(
args: Option<Vec<ToolArgument>>,
tx: Sender<Message>,
_command_tx: Option<Sender<Message>>,
_rx: Option<Receiver<Message>>,
) {
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<Vec<ToolArgument>>) -> 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<Vec<ToolArgument>>) -> String {
pub fn new_project(
args: Option<Vec<ToolArgument>>,
tx: Sender<Message>,
comand_tx: Option<Sender<Message>>,
rx: Option<Receiver<Message>>,
) {
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<Vec<ToolArgument>>) -> 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(&notes_path);
if files_dir_res.is_err() {
return (format!(
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!(
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<Vec<ToolArgument>>) -> 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<Vec<ToolArgument>>) -> String {
pub fn promote_project(
args: Option<Vec<ToolArgument>>,
tx: Sender<Message>,
comand_tx: Option<Sender<Message>>,
rx: Option<Receiver<Message>>,
) {
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<Vec<ToolArgument>>) -> 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");
}
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!");
lines.push(line);
}
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<Vec<ToolArgument>>) -> 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<Vec<ToolArgument>>) -> String {
@@ -304,3 +549,170 @@ pub fn remove_project(args: Option<Vec<ToolArgument>>) -> String {
}
return project.remove_project();
}
pub fn subdomain_brute(args: Option<Vec<ToolArgument>>) -> 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<Vec<ToolArgument>>) -> 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<Receiver<Message>>,
tx: Sender<Message>,
) -> Receiver<Message> {
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<Message>,
tx: Sender<Message>,
prompt: &str,
) -> (String, Receiver<Message>) {
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<Message>) {
tx.blocking_send(Message {
source: Destination::Console,
destination: Destination::Console,
content: String::from("noninteractive"),
})
.unwrap();
}

9
src/lib.rs
View File

@@ -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!(

66
src/main.rs
View File

@@ -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,6 +105,7 @@ pub fn load_projects(path: &PathBuf, display: bool) -> Vec<lib::Project> {
if conf_string_res.is_ok() {
let conf_string = conf_string_res.unwrap();
for line in conf_string.lines() {
if line.len() > 2 {
let line_vec: Vec<&str> = line.split("|").collect();
match line_vec[0] {
"name" => {
@@ -126,11 +129,12 @@ pub fn load_projects(path: &PathBuf, display: bool) -> Vec<lib::Project> {
_ => {
print_error(
"unknown setting discoverd in project config file!",
None,
Some(String::from(line_vec[0])),
);
}
}
}
}
if new_project.boxname.len() > 0 {
projects.push(new_project);
} else {
@@ -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,
));
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 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 {
println!("gui coming soon!");
let rt_handle = runtime.handle();
let handle = runtime.spawn(cli::cli(
console_tx,
console_rx,
main_tx,
config_path.clone(),
rt_handle.clone(),
));
loop {
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;
}
}
}
}
}