did some refinement on the client/server protocol
This commit is contained in:
+321
-87
@@ -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()?;
|
||||
|
||||
Reference in New Issue
Block a user