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

did some refinement on the client/server protocol

Browse Source
This commit is contained in:
pyro
2026-06-23 17:12:29 -05:00
parent 283a336643
commit ae91cc0ce4
6 changed files with 1341 additions and 352 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+186 -149
View File
File diff suppressed because it is too large Load Diff
Cargo.toml
+2
View File
@@ -8,8 +8,10 @@ clap = { version = "4.6.1", features = ["derive"] }
crossterm = "0.29.0" crossterm = "0.29.0"
fs_extra = "1.3.0" fs_extra = "1.3.0"
iced = { version = "0.14.0", features = ["advanced", "tokio"] } iced = { version = "0.14.0", features = ["advanced", "tokio"] }
ipnet = "2.12.0"
ratatui = "0.30.0" ratatui = "0.30.0"
rayon = "1.12.0" rayon = "1.12.0"
rhai = { version = "1.24.0", features = ["metadata", "sync"] } rhai = { version = "1.24.0", features = ["metadata", "sync"] }
rustc-hash = "2.1.2" rustc-hash = "2.1.2"
tokio = { version = "1.52.3", features = ["full"] }
walkdir = "2.5.0" walkdir = "2.5.0"
src/funcs.rs
+298 -64
View File
@@ -1,17 +1,31 @@
use crate::*; use crate::{ToolMessage::Input, *};
use crossterm::{ use crossterm::{
cursor, cursor,
event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode, KeyEventKind}, event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode, KeyEventKind},
execute, execute,
terminal::{EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode}, terminal::{EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode},
}; };
use iced::alignment::Vertical::Bottom;
use ratatui::{ use ratatui::{
prelude::*, prelude::*,
widgets::{Block, Borders, List, ListItem, ListState, Paragraph}, 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::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>> { pub fn get_user_input(prompt: &str) -> Result<String, Box<dyn Error>> {
println!("{}", prompt); println!("{}", prompt);
@@ -20,6 +34,158 @@ pub fn get_user_input(prompt: &str) -> Result<String, Box<dyn Error>> {
return Ok(response.trim().to_string()); 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( pub fn run_tui(
mut state: AppState, mut state: AppState,
main_rx: Receiver<ToolMessage>, main_rx: Receiver<ToolMessage>,
@@ -32,6 +198,12 @@ pub fn run_tui(
cursor::Hide, cursor::Hide,
EnableMouseCapture 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 backend = CrosstermBackend::new(&stdout);
let mut terminal = Terminal::new(backend)?; let mut terminal = Terminal::new(backend)?;
let (event_tx, event_rx) = channel::<AppEvent>(); let (event_tx, event_rx) = channel::<AppEvent>();
@@ -43,6 +215,9 @@ pub fn run_tui(
} }
std::thread::spawn(move || { std::thread::spawn(move || {
loop { loop {
if !state.server_broker_running && !state.app_state_running {
break;
}
if event::poll(Duration::from_millis(100)).unwrap_or(false) { if event::poll(Duration::from_millis(100)).unwrap_or(false) {
match event::read() { match event::read() {
Ok(Event::Key(key)) => { Ok(Event::Key(key)) => {
@@ -222,10 +397,24 @@ pub fn run_tui(
if let Ok(event) = event_rx.recv() { if let Ok(event) = event_rx.recv() {
match event { match event {
AppEvent::Worker(msg) => match msg { 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.log.push(txt.clone());
state.output.push(txt); state.output.push(txt.clone());
state.output_scroll = u16::MAX; 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) => { ToolMessage::UpdateProject(index, project) => {
if index < state.projects.len() { if index < state.projects.len() {
@@ -233,8 +422,16 @@ pub fn run_tui(
} }
} }
ToolMessage::RebuildDB => { ToolMessage::RebuildDB => {
state.main_tx.send(ToolMessage::Output((
0,
String::from("rebuilddb message recieved..."),
)))?;
let project = state.projects[state.selected_project].clone(); 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()?; disable_raw_mode()?;
execute!( execute!(
&stdout, &stdout,
@@ -243,67 +440,74 @@ pub fn run_tui(
DisableMouseCapture DisableMouseCapture
)?; )?;
std::io::stdout().flush()?; 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 if let Err(_) = db.create(state.main_tx.clone(), 0) {
.arg("--volume") state.main_tx.send(ToolMessage::Output((
.arg(format!("{}:{}:rw", volume, dir_name)); 0,
} String::from("failed to create the project box!"),
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);
}
_ => {}
} }
enable_raw_mode()?; enable_raw_mode()?;
execute!( execute!(
&stdout, &stdout,
EnterAlternateScreen, EnterAlternateScreen,
DisableMouseCapture, EnableMouseCapture,
cursor::Hide cursor::Hide
)?; )?;
terminal.clear()?; 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) => { AppEvent::Key(key) => {
match key.code { match key.code {
KeyCode::Esc => break, KeyCode::Esc => {
break;
}
KeyCode::Enter => { KeyCode::Enter => {
let trimmed = state.curent_intput.trim().to_string(); let trimmed = state.curent_intput.trim().to_string();
if !trimmed.is_empty() { if !trimmed.is_empty() {
@@ -318,11 +522,30 @@ pub fn run_tui(
state.execute_command( state.execute_command(
prompt.execute_command.as_str(), prompt.execute_command.as_str(),
None, None,
0,
)?; )?;
} else {
state.execute_command("prompt", None, 0)?;
} }
} }
let (command, args) = let (command, args) =
trimmed.split_once(' ').unwrap_or((&trimmed, "")); trimmed.split_once(' ').unwrap_or((&trimmed, ""));
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
),
)),
);
}
} else {
match command { match command {
"exit" | "quit" => break, "exit" | "quit" => break,
"reload-modules" => { "reload-modules" => {
@@ -330,15 +553,18 @@ pub fn run_tui(
state.output.push("Reloading module paths...".into()); state.output.push("Reloading module paths...".into());
} }
"help" => { "help" => {
state.output.push("Available Modules:".into()); let help_text = state.help.clone().join("\n");
for name in state.module_loader.commands.keys() { for line in help_text.lines() {
state.output.push(format!(" - {}", name)); state.output.push(line.to_string());
} }
} }
"new_project" | "np" => { "new_project" | "np" => {
if args.split_once(' ').is_some() { if args.split_once(' ').is_some() {
let _ = state let _ = state.execute_command(
.execute_command(command, Some(args.to_string())); command,
Some(args.to_string()),
0,
);
} else { } else {
state state
.output .output
@@ -346,30 +572,35 @@ pub fn run_tui(
} }
} }
command_name => { command_name => {
if state.module_loader.commands.contains_key(command_name) { if state.module_loader.commands.contains_key(command) {
state.output.push(format!( state.output.push(format!(
"[Worker] Executing script '{}'...", "[Worker] Executing script '{}'...",
command_name command_name
)); ));
if let Err(e) = if let Err(e) = state.execute_command(
state.execute_command(command_name, None) command,
{ Some(args.to_string()),
state 0,
.output ) {
.push(format!("[Error] Pipeline fail: {}", e)); state.output.push(format!(
"[Error] Pipeline fail: {}",
e
));
} }
} else { } else {
if args == "" { if args == "" {
let _ = state.execute_command(command, None); let _ = state.execute_command(command, None, 0);
} else { } else {
let _ = state.execute_command( let _ = state.execute_command(
command, command,
Some(args.to_string()), Some(args.to_string()),
0,
); );
} }
} }
} }
} }
}
state.curent_intput.clear(); state.curent_intput.clear();
} }
history_index = state.history.len(); history_index = state.history.len();
@@ -457,6 +688,9 @@ pub fn run_tui(
} }
} }
} }
if !state.app_state_running && state.server_broker_running {
break;
}
} }
} }
disable_raw_mode()?; disable_raw_mode()?;
src/install.rs
+30
View File
@@ -9,6 +9,26 @@ use std::process::exit;
use tetanus::funcs::get_user_input; use tetanus::funcs::get_user_input;
pub fn install() -> Result<(), Box<dyn Error>> { 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() { if let Some(homedir) = env::home_dir() {
let mut config_path = homedir.clone(); let mut config_path = homedir.clone();
config_path.push(".config/tetanus"); config_path.push(".config/tetanus");
@@ -40,6 +60,15 @@ pub fn install() -> Result<(), Box<dyn Error>> {
let template_box = get_user_input( let template_box = get_user_input(
"enter the name of your template distrobox, NONE in all caps if you are not using distrobox", "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"); config_path.push("client.conf");
projects_path.push("default"); projects_path.push("default");
create_dir_all(&projects_path)?; 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())?; .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!("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!("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!("\ndefault config files written!");
println!("downloading default notes and modules..."); println!("downloading default notes and modules...");
create_dir_all("./temp")?; create_dir_all("./temp")?;
src/lib.rs
+728 -60
View File
File diff suppressed because it is too large Load Diff
src/main.rs
+20 -2
View File
@@ -1,7 +1,9 @@
use clap::Parser; use clap::Parser;
use std::sync::Arc;
use std::{env, path::PathBuf, process::exit}; use std::{env, path::PathBuf, process::exit};
use tetanus::AppState; use tetanus::AppState;
use tetanus::funcs::*; use tetanus::funcs::*;
use tokio::sync::Mutex;
mod install; mod install;
@@ -20,9 +22,24 @@ struct Args {
#[arg(short, long, help = "start gui client")] #[arg(short, long, help = "start gui client")]
gui: bool, gui: bool,
#[arg(short, long, help = "for testing stuff...")]
test: bool,
} }
fn main() { #[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..."); println!("checking for server or client config files...");
let mut config_path = env::home_dir().unwrap(); let mut config_path = env::home_dir().unwrap();
config_path.push(".config/tetanus"); config_path.push(".config/tetanus");
@@ -59,6 +76,7 @@ fn main() {
exit(1); exit(1);
} }
appstate.initialize_modules(); appstate.initialize_modules();
let _res = run_tui(appstate, rx); let _res = run_tui(appstate, rx).unwrap();
}
} }
} }