did some refinement on the client/server protocol
This commit is contained in:
Generated
+186
-149
File diff suppressed because it is too large
Load Diff
@@ -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"
|
||||
|
||||
+298
-64
@@ -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,11 +522,30 @@ 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, ""));
|
||||
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 {
|
||||
"exit" | "quit" => break,
|
||||
"reload-modules" => {
|
||||
@@ -330,15 +553,18 @@ pub fn run_tui(
|
||||
state.output.push("Reloading module paths...".into());
|
||||
}
|
||||
"help" => {
|
||||
state.output.push("Available Modules:".into());
|
||||
for name in state.module_loader.commands.keys() {
|
||||
state.output.push(format!(" - {}", name));
|
||||
let help_text = state.help.clone().join("\n");
|
||||
for line in help_text.lines() {
|
||||
state.output.push(line.to_string());
|
||||
}
|
||||
}
|
||||
"new_project" | "np" => {
|
||||
if args.split_once(' ').is_some() {
|
||||
let _ = state
|
||||
.execute_command(command, Some(args.to_string()));
|
||||
let _ = state.execute_command(
|
||||
command,
|
||||
Some(args.to_string()),
|
||||
0,
|
||||
);
|
||||
} else {
|
||||
state
|
||||
.output
|
||||
@@ -346,30 +572,35 @@ pub fn run_tui(
|
||||
}
|
||||
}
|
||||
command_name => {
|
||||
if state.module_loader.commands.contains_key(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_name, None)
|
||||
{
|
||||
state
|
||||
.output
|
||||
.push(format!("[Error] Pipeline fail: {}", e));
|
||||
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);
|
||||
let _ = state.execute_command(command, None, 0);
|
||||
} else {
|
||||
let _ = state.execute_command(
|
||||
command,
|
||||
Some(args.to_string()),
|
||||
0,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
state.curent_intput.clear();
|
||||
}
|
||||
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()?;
|
||||
|
||||
@@ -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")?;
|
||||
|
||||
+728
-60
File diff suppressed because it is too large
Load Diff
+20
-2
@@ -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,9 +22,24 @@ struct Args {
|
||||
|
||||
#[arg(short, long, help = "start gui client")]
|
||||
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...");
|
||||
let mut config_path = env::home_dir().unwrap();
|
||||
config_path.push(".config/tetanus");
|
||||
@@ -59,6 +76,7 @@ fn main() {
|
||||
exit(1);
|
||||
}
|
||||
appstate.initialize_modules();
|
||||
let _res = run_tui(appstate, rx);
|
||||
let _res = run_tui(appstate, rx).unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user