did some refinement on the client/server protocol

This commit is contained in:
2026-06-23 17:12:29 -05:00
parent 283a336643
commit ae91cc0ce4
6 changed files with 1341 additions and 352 deletions
Generated
+186 -149
View File
File diff suppressed because it is too large Load Diff
+2
View File
@@ -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"
+321 -87
View File
@@ -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");
let template_box = state.config.get("template_box").unwrap();
status
.arg("create")
.arg("--root")
.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 db.created {
if let Err(_) = db.destroy(state.main_tx.clone(), 0) {
state.main_tx.send(ToolMessage::Output((
0,
String::from("failed to destroy the project box!"),
)))?;
}
if volume.contains("/notes") {
dir_name = "notes";
}
status
.arg("--volume")
.arg(format!("{}:{}:rw", volume, dir_name));
}
status
.stdin(Stdio::inherit())
.stdout(Stdio::inherit())
.stderr(Stdio::inherit());
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);
}
_ => {}
if let Err(_) = db.create(state.main_tx.clone(), 0) {
state.main_tx.send(ToolMessage::Output((
0,
String::from("failed to create the project box!"),
)))?;
}
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 {
"exit" | "quit" => break,
"reload-modules" => {
state.initialize_modules();
state.output.push("Reloading module paths...".into());
if state.selected_client != 0 {
if let Some(server) = state.selected_server.clone() {
let _ = server_tx.send(
ServerBrokerCmd::RegisterActionOutput((
server.clone(),
format!(
"CMD|{}|{}|{} {}",
server.client_id,
state.selected_client,
command,
args
),
)),
);
}
"help" => {
state.output.push("Available Modules:".into());
for name in state.module_loader.commands.keys() {
state.output.push(format!(" - {}", name));
} else {
match command {
"exit" | "quit" => break,
"reload-modules" => {
state.initialize_modules();
state.output.push("Reloading module paths...".into());
}
}
"new_project" | "np" => {
if args.split_once(' ').is_some() {
let _ = state
.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));
"help" => {
let help_text = state.help.clone().join("\n");
for line in help_text.lines() {
state.output.push(line.to_string());
}
} else {
if args == "" {
let _ = state.execute_command(command, None);
} else {
}
"new_project" | "np" => {
if args.split_once(' ').is_some() {
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()?;
+30
View File
@@ -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")?;
+750 -82
View File
File diff suppressed because it is too large Load Diff
+52 -34
View File
@@ -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() {
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 home directory...");
config_path = PathBuf::from("/etc/tetanus");
#[tokio::main]
async fn main() {
let args = Args::parse();
if args.server {
let server = Arc::new(Mutex::new(tetanus::server::ServerState {
id: 0,
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!");
let install_result = install::install();
if install_result.is_ok() {
install_result.unwrap();
println!("install succeeded!");
exit(0);
println!("no config directory found in home directory...");
config_path = PathBuf::from("/etc/tetanus");
if !config_path.exists() {
println!("no config directory found in /etc/tetanus... Installing!");
let install_result = install::install();
if install_result.is_ok() {
install_result.unwrap();
println!("install succeeded!");
exit(0);
}
}
}
}
let args = Args::parse();
let mut client_config_path = config_path.clone();
client_config_path.push("client.conf");
let mut server_config_path = config_path.clone();
server_config_path.push("server.conf");
if args.client {
if !client_config_path.exists() {
eprintln!(
"error: no client config path found at {}",
client_config_path.display()
);
exit(1);
let args = Args::parse();
let mut client_config_path = config_path.clone();
client_config_path.push("client.conf");
let mut server_config_path = config_path.clone();
server_config_path.push("server.conf");
if args.client {
if !client_config_path.exists() {
eprintln!(
"error: no client config path found at {}",
client_config_path.display()
);
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);
}
}