did some refinement on the client/server protocol

Cargo.toml

@@ -8,8 +8,10 @@ clap = { version = "4.6.1", features = ["derive"] }
 crossterm = "0.29.0"
 fs_extra = "1.3.0"
 iced = { version = "0.14.0", features = ["advanced", "tokio"] }
+ipnet = "2.12.0"
 ratatui = "0.30.0"
 rayon = "1.12.0"
 rhai = { version = "1.24.0", features = ["metadata", "sync"] }
 rustc-hash = "2.1.2"
+tokio = { version = "1.52.3", features = ["full"] }
 walkdir = "2.5.0"

src/funcs.rs

@@ -1,17 +1,31 @@
-use crate::*;
+use crate::{ToolMessage::Input, *};
 use crossterm::{
     cursor,
     event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode, KeyEventKind},
     execute,
     terminal::{EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode},
 };
-use iced::alignment::Vertical::Bottom;
 use ratatui::{
     prelude::*,
     widgets::{Block, Borders, List, ListItem, ListState, Paragraph},
 };
-use std::error::Error;
+use std::sync::Arc;
+use std::sync::Mutex;
 use std::time::Duration;
+use std::{error::Error, time::Instant};
+use std::{
+    io::{Read, Write},
+    usize,
+};
+use std::{net::TcpStream, thread};
+
+pub enum ServerBrokerCmd {
+    ConnectServer(usize),
+    RegisterActionOutput((Server, String)),
+    Disconnect(Server),
+    DisconnectAll,
+    Exit,
+}
 
 pub fn get_user_input(prompt: &str) -> Result<String, Box<dyn Error>> {
     println!("{}", prompt);
@@ -20,6 +34,158 @@ pub fn get_user_input(prompt: &str) -> Result<String, Box<dyn Error>> {
     return Ok(response.trim().to_string());
 }
 
+pub fn server_broker(
+    rx: Receiver<ServerBrokerCmd>,
+    tx: Sender<ToolMessage>,
+    servers: Arc<Mutex<Vec<Server>>>,
+) {
+    loop {
+        if let Ok(msg) = rx.try_recv() {
+            match msg {
+                ServerBrokerCmd::ConnectServer(id) => {
+                    let _ = tx.send(ToolMessage::Output((
+                        0,
+                        format!("attempting to connect to server {}", id),
+                    )));
+                    if let Ok(mut lock) = servers.lock() {
+                        lock.iter_mut().for_each(|s| {
+                            let _ =
+                                tx.send(ToolMessage::Output((0, format!("got server list lock!"))));
+                            if s.id == id {
+                                let _ = tx.send(ToolMessage::Output((
+                                    0,
+                                    format!("Found the server for {}: {}", s.id, s.address),
+                                )));
+                                if let Ok(mut stream) = TcpStream::connect(s.address.clone()) {
+                                    let _ = tx.send(ToolMessage::Output((
+                                        0,
+                                        String::from("TCP successful to server!"),
+                                    )));
+                                    let _ = stream.write(format!("HELLO|{}\n", s.id).as_bytes());
+                                    let _ = tx.send(ToolMessage::Output((
+                                        0,
+                                        String::from("Hello message sent!"),
+                                    )));
+                                    let mut buf = [0; 1024];
+                                    if let Ok(bytes_read) = stream.read(&mut buf) {
+                                        let response = String::from_utf8_lossy(&buf[..bytes_read]);
+                                        if let Some((_, id)) = response.split_once("|") {
+                                            if let Ok(id) = id.trim().parse::<usize>() {
+                                                s.client_id = id;
+                                                s.connected = true;
+                                                let _ = tx.send(ToolMessage::Output((
+                                                    0,
+                                                    String::from("got client ID from server!"),
+                                                )));
+                                            } else {
+                                                let _ = tx.send(ToolMessage::Output((
+                                                    0,
+                                                    String::from(
+                                                        "failed to get response from server!",
+                                                    ),
+                                                )));
+                                            }
+                                        }
+                                    }
+                                }
+                            }
+                        });
+                    }
+                }
+                ServerBrokerCmd::Disconnect(server) => {
+                    if let Ok(mut lock) = servers.lock() {
+                        let mut rid = 0;
+                        let mut found = false;
+                        lock.iter().enumerate().for_each(|(id, s)| {
+                            if s.id == server.id {
+                                rid = id;
+                                found = true;
+                            }
+                        });
+                        if found {
+                            lock.remove(rid);
+                        }
+                    }
+                }
+                ServerBrokerCmd::DisconnectAll => {
+                    println!("todo");
+                }
+                ServerBrokerCmd::Exit => {
+                    println!("todo");
+                    break;
+                }
+                ServerBrokerCmd::RegisterActionOutput(text) => {
+                    if let Ok(mut lock) = servers.lock() {
+                        lock.iter_mut().for_each(|s| {
+                            if s.id == text.0.id {
+                                s.message_que.push(text.1.clone());
+                            }
+                        });
+                    }
+                }
+            }
+        }
+        if let Ok(mut lock) = servers.try_lock() {
+            lock.iter_mut().for_each(|s| {
+                if s.connected {
+                    if Instant::now().duration_since(s.last_check) >= s.timer {
+                        if let Ok(mut stream) = TcpStream::connect(s.address.clone()) {
+                            s.last_check = Instant::now();
+                            let mut success_ids = Vec::new();
+                            if s.message_que.len() > 0 {
+                                s.message_que.iter().enumerate().for_each(|(id, msg)| {
+                                    if let Ok(_) = stream.write(
+                                        format!("{}|{}\n", s.client_id.clone(), msg).as_bytes(),
+                                    ) {
+                                        success_ids.push(id);
+                                    }
+                                });
+                            } else {
+                                let _ = stream.write(format!("{}|PING", s.client_id).as_bytes());
+                            }
+                            for id in success_ids {
+                                s.message_que.remove(id);
+                            }
+                            let mut buffer = [0; 1024];
+                            if let Ok(bytes_read) = stream.read(&mut buffer) {
+                                let response =
+                                    String::from_utf8_lossy(&buffer[..bytes_read]).to_string();
+                                let (_, action_string) =
+                                    response.split_once("|").unwrap_or(("ACTIONS", "NONE"));
+                                let actions: Vec<&str> = action_string.split(",").collect();
+                                for action in actions {
+                                    let (action, data) =
+                                        action.split_once("|").unwrap_or(("PING", "NONE"));
+                                    match action {
+                                        "OUTPUT" => {
+                                            let _ = tx.send(ToolMessage::Output((
+                                                0,
+                                                format!("FROM SERVER: {}", data),
+                                            )));
+                                        }
+                                        "CMD" => {
+                                            let (rid_str, cmd) =
+                                                data.split_once("|").unwrap_or(("0", "NONE"));
+                                            if let Ok(rid) = rid_str.parse::<usize>() {
+                                                let _ = tx.send(ToolMessage::Input((
+                                                    rid,
+                                                    cmd.to_string(),
+                                                )));
+                                            }
+                                        }
+                                        _ => {}
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            });
+        }
+        thread::sleep(Duration::from_millis(10));
+    }
+}
+
 pub fn run_tui(
     mut state: AppState,
     main_rx: Receiver<ToolMessage>,
@@ -32,6 +198,12 @@ pub fn run_tui(
         cursor::Hide,
         EnableMouseCapture
     )?;
+    let (server_tx, server_rx) = channel();
+    let main_tx = state.main_tx.clone();
+    let server_clone = state.servers.clone();
+    std::thread::spawn(move || {
+        server_broker(server_rx, main_tx, server_clone);
+    });
     let backend = CrosstermBackend::new(&stdout);
     let mut terminal = Terminal::new(backend)?;
     let (event_tx, event_rx) = channel::<AppEvent>();
@@ -43,6 +215,9 @@ pub fn run_tui(
     }
     std::thread::spawn(move || {
         loop {
+            if !state.server_broker_running && !state.app_state_running {
+                break;
+            }
             if event::poll(Duration::from_millis(100)).unwrap_or(false) {
                 match event::read() {
                     Ok(Event::Key(key)) => {
@@ -222,10 +397,24 @@ pub fn run_tui(
         if let Ok(event) = event_rx.recv() {
             match event {
                 AppEvent::Worker(msg) => match msg {
-                    ToolMessage::Output(txt) => {
+                    ToolMessage::EndPrompt => {
+                        state.prompt.reset();
+                    }
+                    ToolMessage::Input(cmd) => {
+                        state.log.push(cmd.1.clone());
+                    }
+                    ToolMessage::Output((rid, txt)) => {
                         state.log.push(txt.clone());
-                        state.output.push(txt);
+                        state.output.push(txt.clone());
                         state.output_scroll = u16::MAX;
+                        if rid != 0 {
+                            if let Some(server) = state.selected_server.clone() {
+                                let _ = server_tx.send(ServerBrokerCmd::RegisterActionOutput((
+                                    server,
+                                    format!("OUTPUT|{}|{}", rid, txt),
+                                )));
+                            }
+                        }
                     }
                     ToolMessage::UpdateProject(index, project) => {
                         if index < state.projects.len() {
@@ -233,8 +422,16 @@ pub fn run_tui(
                         }
                     }
                     ToolMessage::RebuildDB => {
+                        state.main_tx.send(ToolMessage::Output((
+                            0,
+                            String::from("rebuilddb message recieved..."),
+                        )))?;
                         let project = state.projects[state.selected_project].clone();
-                        if let Some(db) = project.db {
+                        if let Some(mut db) = project.db.clone() {
+                            state.main_tx.send(ToolMessage::Output((
+                                0,
+                                String::from("distrobox detected..."),
+                            )))?;
                             disable_raw_mode()?;
                             execute!(
                                 &stdout,
@@ -243,67 +440,74 @@ pub fn run_tui(
                                 DisableMouseCapture
                             )?;
                             std::io::stdout().flush()?;
-                            let mut status = Command::new("distrobox");
+                            if db.created {
-                            let template_box = state.config.get("template_box").unwrap();
+                                if let Err(_) = db.destroy(state.main_tx.clone(), 0) {
-                            status
+                                    state.main_tx.send(ToolMessage::Output((
-                                .arg("create")
+                                        0,
-                                .arg("--root")
+                                        String::from("failed to destroy the project box!"),
-                                .arg("--clone")
+                                    )))?;
-                                .arg(&template_box)
-                                .arg("--name")
-                                .arg(format!(
-                                    "{}-{}-{}",
-                                    template_box, project.org_name, project.name
-                                ));
-                            for volume in db.volumes {
-                                let mut dir_name = volume.split("/").last().unwrap();
-                                if volume.contains("files") {
-                                    dir_name = "/pentest";
                                 }
-                                if volume.contains("/notes") {
-                                    dir_name = "notes";
-                                }
-                                status
-                                    .arg("--volume")
-                                    .arg(format!("{}:{}:rw", volume, dir_name));
                             }
-                            status
+                            if let Err(_) = db.create(state.main_tx.clone(), 0) {
-                                .stdin(Stdio::inherit())
+                                state.main_tx.send(ToolMessage::Output((
-                                .stdout(Stdio::inherit())
+                                    0,
-                                .stderr(Stdio::inherit());
+                                    String::from("failed to create the project box!"),
-                            match status.status() {
+                                )))?;
-                                Ok(status) if !status.success() => {
-                                    println!(
-                                        "\nCommand failed with exit code {status}, press enter to return to tetanus."
-                                    );
-                                    let mut discard = String::new();
-                                    let _ = std::io::stdin().read_line(&mut discard);
-                                }
-                                Err(e) => {
-                                    println!(
-                                        "\nFailed to execure distrobox: {e}, press enter to return to tetanus."
-                                    );
-                                    let mut discard = String::new();
-                                    let _ = std::io::stdin().read_line(&mut discard);
-                                }
-                                _ => {}
                             }
                             enable_raw_mode()?;
                             execute!(
                                 &stdout,
                                 EnterAlternateScreen,
-                                DisableMouseCapture,
+                                EnableMouseCapture,
                                 cursor::Hide
                             )?;
                             terminal.clear()?;
                         }
                     }
-
+                    ToolMessage::DestroyDB(mut db) => {
+                        execute!(
+                            &stdout,
+                            LeaveAlternateScreen,
+                            cursor::Show,
+                            DisableMouseCapture
+                        )?;
+                        db.destroy(state.main_tx.clone(), 0)?;
+                        enable_raw_mode()?;
+                        execute!(
+                            &stdout,
+                            EnterAlternateScreen,
+                            EnableMouseCapture,
+                            cursor::Hide
+                        )?;
+                        terminal.clear()?;
+                    }
+                    ToolMessage::ConnectServer => {
+                        if let Some(server) = state.selected_server.clone() {
+                            let _ = server_tx.send(ServerBrokerCmd::ConnectServer(server.id));
+                        }
+                    }
+                    ToolMessage::SendServer(data) => {
+                        let _ = server_tx.send(ServerBrokerCmd::RegisterActionOutput(data));
+                    }
+                    ToolMessage::DisconnectAllServers => {
+                        println!("todo");
+                    }
+                    ToolMessage::DisconnectServer => {
+                        println!("todo");
+                    }
+                    ToolMessage::ServerBrokerExit => {
+                        state.server_broker_running = false;
+                    }
+                    ToolMessage::AppStateExit => {
+                        println!("todo");
+                    }
                     _ => {}
                 },
                 AppEvent::Key(key) => {
                     match key.code {
-                        KeyCode::Esc => break,
+                        KeyCode::Esc => {
+                            break;
+                        }
                         KeyCode::Enter => {
                             let trimmed = state.curent_intput.trim().to_string();
                             if !trimmed.is_empty() {
@@ -318,54 +522,81 @@ pub fn run_tui(
                                         state.execute_command(
                                             prompt.execute_command.as_str(),
                                             None,
+                                            0,
                                         )?;
+                                    } else {
+                                        state.execute_command("prompt", None, 0)?;
                                     }
                                 }
                                 let (command, args) =
                                     trimmed.split_once(' ').unwrap_or((&trimmed, ""));
-                                match command {
+                                if state.selected_client != 0 {
-                                    "exit" | "quit" => break,
+                                    if let Some(server) = state.selected_server.clone() {
-                                    "reload-modules" => {
+                                        let _ = server_tx.send(
-                                        state.initialize_modules();
+                                            ServerBrokerCmd::RegisterActionOutput((
-                                        state.output.push("Reloading module paths...".into());
+                                                server.clone(),
+                                                format!(
+                                                    "CMD|{}|{}|{} {}",
+                                                    server.client_id,
+                                                    state.selected_client,
+                                                    command,
+                                                    args
+                                                ),
+                                            )),
+                                        );
                                     }
-                                    "help" => {
+                                } else {
-                                        state.output.push("Available Modules:".into());
+                                    match command {
-                                        for name in state.module_loader.commands.keys() {
+                                        "exit" | "quit" => break,
-                                            state.output.push(format!("  - {}", name));
+                                        "reload-modules" => {
+                                            state.initialize_modules();
+                                            state.output.push("Reloading module paths...".into());
                                         }
-                                    }
+                                        "help" => {
-                                    "new_project" | "np" => {
+                                            let help_text = state.help.clone().join("\n");
-                                        if args.split_once(' ').is_some() {
+                                            for line in help_text.lines() {
-                                            let _ = state
+                                                state.output.push(line.to_string());
-                                                .execute_command(command, Some(args.to_string()));
-                                        } else {
-                                            state
-                                                .output
-                                                .push("Error: USAGE -> np <org> <name>".into());
-                                        }
-                                    }
-                                    command_name => {
-                                        if state.module_loader.commands.contains_key(command_name) {
-                                            state.output.push(format!(
-                                                "[Worker] Executing script '{}'...",
-                                                command_name
-                                            ));
-                                            if let Err(e) =
-                                                state.execute_command(command_name, None)
-                                            {
-                                                state
-                                                    .output
-                                                    .push(format!("[Error] Pipeline fail: {}", e));
                                             }
-                                        } else {
+                                        }
-                                            if args == "" {
+                                        "new_project" | "np" => {
-                                                let _ = state.execute_command(command, None);
+                                            if args.split_once(' ').is_some() {
-                                            } else {
                                                 let _ = state.execute_command(
                                                     command,
                                                     Some(args.to_string()),
+                                                    0,
                                                 );
+                                            } else {
+                                                state
+                                                    .output
+                                                    .push("Error: USAGE -> np <org> <name>".into());
+                                            }
+                                        }
+                                        command_name => {
+                                            if state.module_loader.commands.contains_key(command) {
+                                                state.output.push(format!(
+                                                    "[Worker] Executing script '{}'...",
+                                                    command_name
+                                                ));
+                                                if let Err(e) = state.execute_command(
+                                                    command,
+                                                    Some(args.to_string()),
+                                                    0,
+                                                ) {
+                                                    state.output.push(format!(
+                                                        "[Error] Pipeline fail: {}",
+                                                        e
+                                                    ));
+                                                }
+                                            } else {
+                                                if args == "" {
+                                                    let _ = state.execute_command(command, None, 0);
+                                                } else {
+                                                    let _ = state.execute_command(
+                                                        command,
+                                                        Some(args.to_string()),
+                                                        0,
+                                                    );
+                                                }
                                             }
                                         }
                                     }
@@ -457,6 +688,9 @@ pub fn run_tui(
                     }
                 }
             }
+            if !state.app_state_running && state.server_broker_running {
+                break;
+            }
         }
     }
     disable_raw_mode()?;

src/install.rs

@@ -9,6 +9,26 @@ use std::process::exit;
 use tetanus::funcs::get_user_input;
 
 pub fn install() -> Result<(), Box<dyn Error>> {
+    let client = get_user_input("would you like to install the Tetanus client on this computer?")?
+        .contains("y");
+    let server = get_user_input("would you like to install the Tetanus server on this computer?")?
+        .contains("y");
+    if client {
+        client_install()?;
+    }
+    if server {
+        println!("todo");
+    }
+    if !client && !server {
+        println!(
+            "no install selected, please answer y to one or both of the above prompts next time!"
+        );
+        exit(1);
+    }
+    return Ok(());
+}
+
+pub fn client_install() -> Result<(), Box<dyn Error>> {
     if let Some(homedir) = env::home_dir() {
         let mut config_path = homedir.clone();
         config_path.push(".config/tetanus");
@@ -40,6 +60,15 @@ pub fn install() -> Result<(), Box<dyn Error>> {
         let template_box = get_user_input(
             "enter the name of your template distrobox, NONE in all caps if you are not using distrobox",
         )?;
+        println!(
+            "enter the command used to launch your terminal while running a command. Put the place holder |||COMMAND||| where the command would go"
+        );
+        println!(
+            "if you use a profile to launch the distroboxes, and that profile needs the box name as an environment variable put ENV_NAME= then the name of the variable you use at the end of the command"
+        );
+        let term_cmd = get_user_input(
+            "for example: konsole -e |||COMMAND||| or konsole --profile profile_name ENV_NAME=CURRENT_PROJECT_BOX",
+        )?;
         config_path.push("client.conf");
         projects_path.push("default");
         create_dir_all(&projects_path)?;
@@ -70,6 +99,7 @@ pub fn install() -> Result<(), Box<dyn Error>> {
             .write(format!("upcoming_notes: {}\n", upcoming_notes_path.display()).as_bytes())?;
         client_config_file.write(format!("module_path: {}\n", module_path.display()).as_bytes())?;
         client_config_file.write(format!("template_box: {}\n", template_box).as_bytes())?;
+        client_config_file.write(format!("term_cmd: {}\n", term_cmd).as_bytes())?;
         println!("\ndefault config files written!");
         println!("downloading default notes and modules...");
         create_dir_all("./temp")?;

src/main.rs

@@ -1,7 +1,9 @@
 use clap::Parser;
+use std::sync::Arc;
 use std::{env, path::PathBuf, process::exit};
 use tetanus::AppState;
 use tetanus::funcs::*;
+use tokio::sync::Mutex;
 
 mod install;
 
@@ -20,45 +22,61 @@ struct Args {
 
     #[arg(short, long, help = "start gui client")]
     gui: bool,
+
+    #[arg(short, long, help = "for testing stuff...")]
+    test: bool,
 }
 
-fn main() {
+#[tokio::main]
-    println!("checking for server or client config files...");
+async fn main() {
-    let mut config_path = env::home_dir().unwrap();
+    let args = Args::parse();
-    config_path.push(".config/tetanus");
+    if args.server {
-    if !config_path.exists() {
+        let server = Arc::new(Mutex::new(tetanus::server::ServerState {
-        println!("no config directory found in home directory...");
+            id: 0,
-        config_path = PathBuf::from("/etc/tetanus");
+            running: true,
+            address: String::from("127.0.0.1:31337"),
+            clients: Vec::new(),
+            key: String::new(),
+        }));
+        tetanus::server::start_server(server).await.unwrap();
+    } else {
+        println!("checking for server or client config files...");
+        let mut config_path = env::home_dir().unwrap();
+        config_path.push(".config/tetanus");
         if !config_path.exists() {
-            println!("no config directory found in /etc/tetanus... Installing!");
+            println!("no config directory found in home directory...");
-            let install_result = install::install();
+            config_path = PathBuf::from("/etc/tetanus");
-            if install_result.is_ok() {
+            if !config_path.exists() {
-                install_result.unwrap();
+                println!("no config directory found in /etc/tetanus... Installing!");
-                println!("install succeeded!");
+                let install_result = install::install();
-                exit(0);
+                if install_result.is_ok() {
+                    install_result.unwrap();
+                    println!("install succeeded!");
+                    exit(0);
+                }
             }
         }
-    }
+        let args = Args::parse();
-    let args = Args::parse();
+        let mut client_config_path = config_path.clone();
-    let mut client_config_path = config_path.clone();
+        client_config_path.push("client.conf");
-    client_config_path.push("client.conf");
+        let mut server_config_path = config_path.clone();
-    let mut server_config_path = config_path.clone();
+        server_config_path.push("server.conf");
-    server_config_path.push("server.conf");
+        if args.client {
-    if args.client {
+            if !client_config_path.exists() {
-        if !client_config_path.exists() {
+                eprintln!(
-            eprintln!(
+                    "error: no client config path found at {}",
-                "error: no client config path found at {}",
+                    client_config_path.display()
-                client_config_path.display()
+                );
-            );
+                exit(1);
-            exit(1);
+            }
+            let (mut appstate, rx) = AppState::new();
+            let config_load = appstate.load_config(client_config_path.clone());
+            if config_load.is_err() {
+                eprintln!("error loading config!");
+                exit(1);
+            }
+            appstate.initialize_modules();
+            let _res = run_tui(appstate, rx).unwrap();
         }
-        let (mut appstate, rx) = AppState::new();
-        let config_load = appstate.load_config(client_config_path.clone());
-        if config_load.is_err() {
-            eprintln!("error loading config!");
-            exit(1);
-        }
-        appstate.initialize_modules();
-        let _res = run_tui(appstate, rx);
     }
 }