made the server and the client/server communication functionality work.
This commit is contained in:
+178
-241
@@ -1,31 +1,20 @@
|
||||
use crate::{ToolMessage::Input, *};
|
||||
use crate::*;
|
||||
use crossterm::{
|
||||
cursor,
|
||||
event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode, KeyEventKind},
|
||||
execute,
|
||||
terminal::{EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode},
|
||||
};
|
||||
use iced::keyboard::key::Code::Sleep;
|
||||
use ratatui::{
|
||||
prelude::*,
|
||||
widgets::{Block, Borders, List, ListItem, ListState, Paragraph},
|
||||
};
|
||||
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,
|
||||
}
|
||||
use std::{io::Write, usize};
|
||||
use std::{sync::Arc, thread::sleep};
|
||||
use std::{sync::Mutex, thread::spawn};
|
||||
|
||||
pub fn get_user_input(prompt: &str) -> Result<String, Box<dyn Error>> {
|
||||
println!("{}", prompt);
|
||||
@@ -34,158 +23,6 @@ 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>,
|
||||
@@ -198,12 +35,6 @@ 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>();
|
||||
@@ -250,6 +81,7 @@ pub fn run_tui(
|
||||
project_list_state.select(Some(0));
|
||||
}
|
||||
let mut history_index = state.history.len();
|
||||
let mut handles = Vec::new();
|
||||
loop {
|
||||
terminal.draw(|f| {
|
||||
let main_chunks = Layout::default()
|
||||
@@ -402,17 +234,22 @@ pub fn run_tui(
|
||||
}
|
||||
ToolMessage::Input(cmd) => {
|
||||
state.log.push(cmd.1.clone());
|
||||
if cmd.1.contains(" ") {
|
||||
let (gcmd, args) = cmd.1.split_once(" ").unwrap();
|
||||
state.execute_command(gcmd, Some(args.to_string()), cmd.0)?;
|
||||
} else {
|
||||
state.execute_command(cmd.1.as_str(), None, cmd.0)?;
|
||||
}
|
||||
}
|
||||
ToolMessage::Output((rid, txt)) => {
|
||||
state.log.push(txt.clone());
|
||||
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),
|
||||
)));
|
||||
if let Some(server) = state.servers.get(state.selected_server) {
|
||||
if let Ok(mut server) = server.lock() {
|
||||
server.message_que.push(format!("OUTPUT|{}", txt));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -482,12 +319,121 @@ pub fn run_tui(
|
||||
terminal.clear()?;
|
||||
}
|
||||
ToolMessage::ConnectServer => {
|
||||
if let Some(server) = state.selected_server.clone() {
|
||||
let _ = server_tx.send(ServerBrokerCmd::ConnectServer(server.id));
|
||||
if let Some(server) = state.servers.get_mut(state.selected_server) {
|
||||
let mut address = String::new();
|
||||
let lock_res = server.try_lock();
|
||||
let mut connected = false;
|
||||
match lock_res {
|
||||
Ok(mut locked_server) => match locked_server.connect() {
|
||||
Ok(id) => {
|
||||
state.output.push(format!(
|
||||
"Server Connected! ID:{} ADDRESS:{}",
|
||||
id,
|
||||
locked_server.address.clone()
|
||||
));
|
||||
address = locked_server.address.clone();
|
||||
connected = true;
|
||||
}
|
||||
Err(e) => {
|
||||
state
|
||||
.output
|
||||
.push(format!("Error connecting to server {e}"));
|
||||
}
|
||||
},
|
||||
Err(e) => {
|
||||
state
|
||||
.output
|
||||
.push(format!("Error locking server state! {e}"));
|
||||
}
|
||||
}
|
||||
if connected {
|
||||
let server_clone = server.clone();
|
||||
let tx_clone = state.main_tx.clone();
|
||||
let mut timer = Duration::from_secs(5);
|
||||
let mut last_check = Instant::now();
|
||||
let server_handle = spawn(move || {
|
||||
let server = server_clone;
|
||||
loop {
|
||||
let mut stop = false;
|
||||
if Instant::now().duration_since(last_check) >= timer {
|
||||
let lock_res = server.try_lock();
|
||||
match lock_res {
|
||||
Ok(mut locked_server) => {
|
||||
locked_server.checkin();
|
||||
last_check = locked_server.last_check.clone();
|
||||
for action in locked_server.action_que.clone() {
|
||||
if action.contains("|") {
|
||||
let (action, data) =
|
||||
action.split_once("|").unwrap();
|
||||
match action.trim() {
|
||||
"OUTPUT" => {
|
||||
let _ = tx_clone.send(
|
||||
ToolMessage::Output((
|
||||
0,
|
||||
data.trim().to_string(),
|
||||
)),
|
||||
);
|
||||
}
|
||||
"CMD" => {
|
||||
let _ = tx_clone.send(
|
||||
ToolMessage::Input((
|
||||
locked_server
|
||||
.client_id
|
||||
.clone(),
|
||||
data.trim().to_string(),
|
||||
)),
|
||||
);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
} else {
|
||||
match action.as_str() {
|
||||
"disconnect" => {
|
||||
locked_server.connected = false;
|
||||
stop = true;
|
||||
}
|
||||
_ => {
|
||||
let _ = tx_clone.send(
|
||||
ToolMessage::Output((
|
||||
0,
|
||||
action.clone(),
|
||||
)),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
timer = locked_server.timer.clone();
|
||||
locked_server.action_que.clear();
|
||||
if locked_server.message_que.is_empty() {
|
||||
let id = locked_server.client_id.clone();
|
||||
locked_server
|
||||
.message_que
|
||||
.push(format!("{}|||PING\n", id));
|
||||
}
|
||||
}
|
||||
Err(_) => {}
|
||||
}
|
||||
}
|
||||
if stop {
|
||||
println!("{} disconnected!", address);
|
||||
break;
|
||||
}
|
||||
sleep(Duration::from_secs(1));
|
||||
}
|
||||
});
|
||||
handles.push(server_handle);
|
||||
}
|
||||
}
|
||||
}
|
||||
ToolMessage::SendServer(data) => {
|
||||
let _ = server_tx.send(ServerBrokerCmd::RegisterActionOutput(data));
|
||||
if let Some(server) = state.servers.get(state.selected_server) {
|
||||
if let Ok(mut server) = server.lock() {
|
||||
server
|
||||
.message_que
|
||||
.push(format!("{}|||{}\n", data.0, data.1));
|
||||
}
|
||||
}
|
||||
}
|
||||
ToolMessage::DisconnectAllServers => {
|
||||
println!("todo");
|
||||
@@ -530,73 +476,55 @@ pub fn run_tui(
|
||||
}
|
||||
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
|
||||
),
|
||||
)),
|
||||
);
|
||||
match command {
|
||||
"exit" | "quit" => break,
|
||||
"reload-modules" => {
|
||||
state.initialize_modules();
|
||||
state.output.push("Reloading module paths...".into());
|
||||
}
|
||||
} else {
|
||||
match command {
|
||||
"exit" | "quit" => break,
|
||||
"reload-modules" => {
|
||||
state.initialize_modules();
|
||||
state.output.push("Reloading module paths...".into());
|
||||
"help" => {
|
||||
let help_text = state.help.clone().join("\n");
|
||||
for line in help_text.lines() {
|
||||
state.output.push(line.to_string());
|
||||
}
|
||||
"help" => {
|
||||
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()),
|
||||
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));
|
||||
}
|
||||
}
|
||||
"new_project" | "np" => {
|
||||
if args.split_once(' ').is_some() {
|
||||
} else {
|
||||
if args == "" {
|
||||
let _ = state.execute_command(command, None, 0);
|
||||
} 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,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -700,5 +628,14 @@ pub fn run_tui(
|
||||
cursor::Show,
|
||||
DisableMouseCapture
|
||||
)?;
|
||||
println!("disconnecting from servers...");
|
||||
state.servers.iter().for_each(|server| {
|
||||
if let Ok(mut server_lock) = server.lock() {
|
||||
server_lock.action_que.push("disconnect".to_string());
|
||||
}
|
||||
});
|
||||
for handle in handles {
|
||||
handle.join().unwrap();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -69,6 +69,8 @@ pub fn client_install() -> Result<(), Box<dyn Error>> {
|
||||
let term_cmd = get_user_input(
|
||||
"for example: konsole -e |||COMMAND||| or konsole --profile profile_name ENV_NAME=CURRENT_PROJECT_BOX",
|
||||
)?;
|
||||
let tools_folder =
|
||||
get_user_input("path where your tools are installed ex /opt or /home/user/tools")?;
|
||||
config_path.push("client.conf");
|
||||
projects_path.push("default");
|
||||
create_dir_all(&projects_path)?;
|
||||
@@ -100,6 +102,7 @@ pub fn client_install() -> Result<(), Box<dyn Error>> {
|
||||
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())?;
|
||||
client_config_file.write(format!("tools: {}\n", tools_folder).as_bytes())?;
|
||||
println!("\ndefault config files written!");
|
||||
println!("downloading default notes and modules...");
|
||||
create_dir_all("./temp")?;
|
||||
|
||||
+267
-157
@@ -2,27 +2,72 @@ use fs_extra::dir::{CopyOptions, copy};
|
||||
use ipnet::IpNet;
|
||||
use ratatui::crossterm::event;
|
||||
use rhai::{AST, Dynamic, Engine, Scope};
|
||||
use rustls::{ClientConfig, ClientConnection, RootCertStore, StreamOwned, pki_types::ServerName};
|
||||
use std::collections::HashMap;
|
||||
use std::error::Error;
|
||||
use std::fs::{self, File, create_dir_all, read_dir, read_to_string, remove_dir, remove_dir_all};
|
||||
use std::io::{self, BufRead, BufReader, Write};
|
||||
use std::io::{self, BufRead, BufReader, Read, Write};
|
||||
use std::net::IpAddr;
|
||||
use std::net::TcpStream;
|
||||
use std::path::PathBuf;
|
||||
use std::process::{Command, Stdio};
|
||||
use std::sync::mpsc::Receiver;
|
||||
use std::sync::{Arc, Mutex, mpsc::Sender, mpsc::channel};
|
||||
|
||||
use std::time;
|
||||
use std::time::Duration;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpStream;
|
||||
use std::time::{self, Instant};
|
||||
|
||||
pub mod funcs;
|
||||
pub mod server;
|
||||
|
||||
enum AppEvent {
|
||||
Key(event::KeyEvent),
|
||||
Worker(ToolMessage),
|
||||
Mouse(event::MouseEvent),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum Destination {
|
||||
Server,
|
||||
Victim,
|
||||
Attacker,
|
||||
Control,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum ToolMessage {
|
||||
Input((usize, String)),
|
||||
Output((usize, String)),
|
||||
UpdateHost(Host),
|
||||
RebuildDB,
|
||||
DestroyDB(DistroBox),
|
||||
StopDB(DistroBox),
|
||||
StopTemplate(DistroBox),
|
||||
UpdateProject(usize, Project),
|
||||
RemoveProject,
|
||||
AddServer,
|
||||
EndPrompt,
|
||||
ConnectServer,
|
||||
SelectServer,
|
||||
SendServer((usize, String)),
|
||||
ServerBrokerExit,
|
||||
AppStateExit,
|
||||
DisconnectServer,
|
||||
DisconnectAllServers,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum ToolArg {
|
||||
Project(Project),
|
||||
Projects(Vec<Project>),
|
||||
Host(Host),
|
||||
Hosts(Vec<Host>),
|
||||
Config(HashMap<String, String>),
|
||||
Path(PathBuf),
|
||||
}
|
||||
|
||||
pub struct AppState {
|
||||
pub projects: Vec<Project>,
|
||||
pub servers: Arc<Mutex<Vec<Server>>>,
|
||||
pub servers: Vec<Arc<Mutex<Server>>>,
|
||||
pub config: HashMap<String, String>,
|
||||
pub config_file: PathBuf,
|
||||
pub workers: rayon::ThreadPool,
|
||||
@@ -31,9 +76,8 @@ pub struct AppState {
|
||||
pub history: Vec<String>,
|
||||
pub log: Vec<String>,
|
||||
pub output: Vec<String>,
|
||||
pub selected_server: Option<Server>,
|
||||
pub selected_server: usize,
|
||||
pub selected_project: usize,
|
||||
pub selected_client: usize,
|
||||
pub curent_intput: String,
|
||||
pub module_loader: ModuleLoader,
|
||||
pub output_scroll: u16,
|
||||
@@ -51,7 +95,7 @@ impl AppState {
|
||||
(
|
||||
Self {
|
||||
projects: Vec::new(),
|
||||
servers: Arc::new(Mutex::new(Vec::new())),
|
||||
servers: Vec::new(),
|
||||
config: HashMap::new(),
|
||||
config_file: PathBuf::new(),
|
||||
workers: rayon::ThreadPoolBuilder::new()
|
||||
@@ -63,9 +107,8 @@ impl AppState {
|
||||
history: Vec::new(),
|
||||
log: Vec::new(),
|
||||
output: Vec::new(),
|
||||
selected_server: None,
|
||||
selected_server: 0,
|
||||
selected_project: 0,
|
||||
selected_client: 0,
|
||||
curent_intput: String::new(),
|
||||
module_loader: ModuleLoader::new(),
|
||||
output_scroll: 0,
|
||||
@@ -100,14 +143,14 @@ impl AppState {
|
||||
address: address.trim().to_string(),
|
||||
connected: false,
|
||||
timer: Duration::from_secs(60),
|
||||
id: 0,
|
||||
last_check: time::Instant::now(),
|
||||
message_que: Vec::new(),
|
||||
action_que: Vec::new(),
|
||||
client_id: 0,
|
||||
selected_client: 0,
|
||||
config: None,
|
||||
};
|
||||
if let Ok(mut lock) = self.servers.lock() {
|
||||
lock.push(new_server);
|
||||
}
|
||||
self.servers.push(Arc::new(Mutex::new(new_server)));
|
||||
}
|
||||
self.config
|
||||
.insert("servers".to_string(), line.trim().to_string());
|
||||
@@ -223,6 +266,27 @@ impl AppState {
|
||||
) -> Result<(), Box<dyn Error>> {
|
||||
let tx = self.main_tx.clone();
|
||||
match command_name {
|
||||
"remote" => {
|
||||
if let Some(server) = self.servers.get(self.selected_server) {
|
||||
if let Ok(mut server) = server.lock() {
|
||||
if server.selected_client != 0 {
|
||||
if let Some(args) = command_args.clone() {
|
||||
let args = args.trim();
|
||||
if args.contains(" ") {
|
||||
let (cmd, iargs) = args.split_once(" ").unwrap();
|
||||
server.message_que.push(format!("CMD|{} {}", cmd, iargs));
|
||||
} else {
|
||||
server.message_que.push(format!("CMD|{}", args));
|
||||
}
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
0,
|
||||
format!("Tasked {} to run {}", server.selected_client, args),
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"new_project" | "np" => {
|
||||
let args = command_args.unwrap();
|
||||
let (org, name) = args.split_once(" ").unwrap();
|
||||
@@ -407,17 +471,14 @@ impl AppState {
|
||||
address: format!("{}:{}", self.prompt.responses[0], self.prompt.responses[1]),
|
||||
connected: false,
|
||||
timer: Duration::from_secs(5),
|
||||
id: 0,
|
||||
last_check: time::Instant::now(),
|
||||
message_que: Vec::new(),
|
||||
action_que: Vec::new(),
|
||||
client_id: 0,
|
||||
selected_client: 0,
|
||||
config: None,
|
||||
};
|
||||
if let Ok(mut lock) = self.servers.lock() {
|
||||
if lock.len() == 0 {
|
||||
self.selected_server = Some(new_server.clone());
|
||||
}
|
||||
lock.push(new_server);
|
||||
}
|
||||
self.servers.push(Arc::new(Mutex::new(new_server)));
|
||||
self.prompt.reset();
|
||||
}
|
||||
"connect_server" => {
|
||||
@@ -440,71 +501,64 @@ impl AppState {
|
||||
}
|
||||
}
|
||||
"list_servers" => {
|
||||
let mut selected_addr = String::new();
|
||||
if let Some(selected) = self.selected_server.clone() {
|
||||
selected_addr = selected.address.clone();
|
||||
}
|
||||
if let Ok(lock) = self.servers.lock() {
|
||||
lock.iter().for_each(|s| {
|
||||
if selected_addr == s.address {
|
||||
self.servers.iter().enumerate().for_each(|(id, m)| {
|
||||
if let Ok(s) = m.lock() {
|
||||
if id == self.selected_server {
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
rid,
|
||||
format!("{} - Currently Selected", s.address),
|
||||
format!(
|
||||
"Server ID:{} Address:{} ClientID:{} Sleep:{} Last:{} - Currently Selected",
|
||||
id, s.address, s.client_id, s.timer.as_secs(), s.last_check.elapsed().as_secs()
|
||||
),
|
||||
)));
|
||||
} else {
|
||||
let _ = self
|
||||
.main_tx
|
||||
.send(ToolMessage::Output((rid, s.address.clone())));
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
rid,
|
||||
format!(
|
||||
"Server ID:{} Address:{} ClientID:{} Sleep:{} Last:{}",
|
||||
id, s.address, s.client_id, s.timer.as_secs(), s.last_check.elapsed().as_secs()
|
||||
),
|
||||
)));
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
"select_server" => {
|
||||
self.prompt.action = Some(ToolMessage::SelectServer);
|
||||
self.prompt.num_responses = 1;
|
||||
self.prompt.prompts.push(String::from("Selection?"));
|
||||
self.prompt.execute_command = String::from("server_selected");
|
||||
let mut counter: usize = 0;
|
||||
let mut selected_addr = String::new();
|
||||
if let Some(selected) = self.selected_server.clone() {
|
||||
selected_addr = selected.address.clone();
|
||||
}
|
||||
if let Ok(lock) = self.servers.lock() {
|
||||
lock.iter().for_each(|server| {
|
||||
if server.address == selected_addr {
|
||||
self.servers.iter().enumerate().for_each(|(id, servermut)| {
|
||||
if let Ok(server) = servermut.lock() {
|
||||
if id == self.selected_server {
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
rid,
|
||||
format!(
|
||||
"{}:{} - Currently Selected",
|
||||
counter,
|
||||
server.address.clone()
|
||||
),
|
||||
format!("{}:{} - Currently Selected", id, server.address.clone()),
|
||||
)));
|
||||
} else {
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
rid,
|
||||
format!("{}: {}", counter, server.address.clone()),
|
||||
format!("{}: {}", id, server.address.clone()),
|
||||
)));
|
||||
}
|
||||
counter += 1;
|
||||
});
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
"server_selected" => {
|
||||
if self.prompt.responses.len() == 1 {
|
||||
if let Ok(lock) = self.servers.lock() {
|
||||
if let Ok(id) = self.prompt.responses[0].parse::<usize>() {
|
||||
self.selected_server = Some(lock[id].clone());
|
||||
if let Ok(id) = self.prompt.responses[0].parse::<usize>() {
|
||||
self.selected_server = id;
|
||||
if let Ok(server) = self.servers[id].lock() {
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
rid,
|
||||
format!("{} Selected!", lock[id].address),
|
||||
)));
|
||||
} else {
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
rid,
|
||||
String::from("Invalid selection."),
|
||||
format!("{} Selected!", server.address),
|
||||
)));
|
||||
}
|
||||
} else {
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
rid,
|
||||
String::from("Invalid selection."),
|
||||
)));
|
||||
}
|
||||
} else {
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
@@ -514,46 +568,81 @@ impl AppState {
|
||||
}
|
||||
let _ = self.main_tx.send(ToolMessage::EndPrompt);
|
||||
}
|
||||
"test_server" => {
|
||||
if let Some(server) = self.servers.get(self.selected_server) {
|
||||
if let Ok(mut server) = server.lock() {
|
||||
server.message_que.push("TEST|NONE".to_string());
|
||||
} else {
|
||||
self.output.push(String::from("error getting server lock!"));
|
||||
}
|
||||
} else {
|
||||
self.output
|
||||
.push(String::from("Error getting selected server!"));
|
||||
}
|
||||
}
|
||||
"server_sleep" => {
|
||||
if let Some(args) = command_args {
|
||||
if let Ok(seconds) = args.trim().parse::<u64>() {
|
||||
if let Some(server) = self.servers.get(self.selected_server) {
|
||||
if let Ok(mut server) = server.lock() {
|
||||
server.timer = Duration::from_secs(seconds);
|
||||
self.output
|
||||
.push(format!("Server timer set to {} seconds", seconds));
|
||||
} else {
|
||||
self.output
|
||||
.push(String::from("error locking server object!"));
|
||||
}
|
||||
} else {
|
||||
self.output
|
||||
.push(String::from("error not a valid server selected!"));
|
||||
}
|
||||
} else {
|
||||
self.output
|
||||
.push(String::from("invalid second count provided!"));
|
||||
}
|
||||
} else {
|
||||
self.output.push(String::from(
|
||||
"Error no seconds provided! please use server_sleep number_of_seconds",
|
||||
));
|
||||
}
|
||||
}
|
||||
"save_all" => {
|
||||
self.save_all()?;
|
||||
}
|
||||
"save_projects" => {
|
||||
self.save_projects()?;
|
||||
}
|
||||
"test_server" => {
|
||||
if let Some(server) = self.selected_server.clone() {
|
||||
let _ = self
|
||||
.main_tx
|
||||
.send(ToolMessage::SendServer((server, "TEST".to_string())));
|
||||
} else {
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
rid,
|
||||
"You need to select a server first!".to_string(),
|
||||
)));
|
||||
}
|
||||
}
|
||||
"list_clients" => {
|
||||
if let Some(server) = self.selected_server.clone() {
|
||||
let _ = self.main_tx.send(ToolMessage::SendServer((
|
||||
server,
|
||||
"LIST_CLIENTS".to_string(),
|
||||
)));
|
||||
if let Some(server) = self.servers.get(self.selected_server) {
|
||||
if let Ok(mut server) = server.lock() {
|
||||
server.message_que.push("LIST_CLIENTS|NONE".to_string());
|
||||
} else {
|
||||
self.output.push(String::from("error getting server lock!"));
|
||||
}
|
||||
} else {
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
rid,
|
||||
"You need to select a server first!".to_string(),
|
||||
)));
|
||||
self.output
|
||||
.push(String::from("Error getting selected server!"));
|
||||
}
|
||||
}
|
||||
"select_client" => {
|
||||
if let Some(arg) = command_args {
|
||||
if let Ok(id) = arg.parse::<usize>() {
|
||||
self.selected_client = id;
|
||||
} else {
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
rid,
|
||||
format!("error parsing client id from arg: {}", arg),
|
||||
)));
|
||||
if let Some(server) = self.servers.get(self.selected_server) {
|
||||
if let Ok(mut server) = server.lock() {
|
||||
if let Ok(id) = arg.parse::<usize>() {
|
||||
server.message_que.push(format!("STOP_CONTROL"));
|
||||
server.message_que.push(format!("CONTROL|{}", id.clone()));
|
||||
server.selected_client = id;
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
0,
|
||||
format!("client {} selected!", id),
|
||||
)));
|
||||
} else {
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
rid,
|
||||
format!("error parsing client id from arg: {}", arg),
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let _ = self.main_tx.send(ToolMessage::Output((
|
||||
@@ -562,6 +651,17 @@ impl AppState {
|
||||
)));
|
||||
}
|
||||
}
|
||||
"deselect_client" => {
|
||||
if let Some(server) = self.servers.get(self.selected_server) {
|
||||
if let Ok(mut server) = server.lock() {
|
||||
server.message_que.push(format!("STOP_CONTROL"));
|
||||
server.selected_client = 0;
|
||||
let _ = self
|
||||
.main_tx
|
||||
.send(ToolMessage::Output((0, format!("client deselected!"))));
|
||||
}
|
||||
}
|
||||
}
|
||||
"add_scope" => {
|
||||
if let Some(args) = command_args {
|
||||
args.split_whitespace().into_iter().for_each(|host| {
|
||||
@@ -902,17 +1002,17 @@ impl AppState {
|
||||
),
|
||||
connected: false,
|
||||
timer: Duration::from_secs(60),
|
||||
id: 0,
|
||||
last_check: time::Instant::now(),
|
||||
message_que: Vec::new(),
|
||||
action_que: Vec::new(),
|
||||
client_id: 0,
|
||||
selected_client: 0,
|
||||
config: None,
|
||||
};
|
||||
if let Ok(mut lock) = self.servers.lock() {
|
||||
lock.push(new_server);
|
||||
self.prompt.action = None;
|
||||
self.prompt.responses.clear();
|
||||
self.prompt.num_responses = 0;
|
||||
}
|
||||
self.servers.push(Arc::new(Mutex::new(new_server)));
|
||||
self.prompt.action = None;
|
||||
self.prompt.responses.clear();
|
||||
self.prompt.num_responses = 0;
|
||||
return Ok(());
|
||||
} else {
|
||||
match prompt.responses.len() {
|
||||
@@ -940,34 +1040,90 @@ pub struct Server {
|
||||
pub address: String,
|
||||
pub connected: bool,
|
||||
pub timer: Duration,
|
||||
pub id: usize,
|
||||
pub config: Option<Arc<ClientConfig>>,
|
||||
pub last_check: time::Instant,
|
||||
pub message_que: Vec<String>,
|
||||
pub action_que: Vec<String>,
|
||||
pub client_id: usize,
|
||||
pub selected_client: usize,
|
||||
}
|
||||
|
||||
impl Server {
|
||||
pub async fn checkin(&mut self, outputs: Vec<String>) -> (String, String) {
|
||||
let mut out_string = (String::new(), String::new());
|
||||
if let Ok(mut stream) = TcpStream::connect(self.address.clone()).await {
|
||||
let payload = format!("{}:{}\n", self.id.clone(), outputs.join(":"));
|
||||
if let Err(e) = stream.write_all(payload.as_bytes()).await {
|
||||
out_string.0 = format!("Error sending response output to server! {e}:");
|
||||
} else {
|
||||
out_string.0 = String::from("Success");
|
||||
pub fn donwload_certificate(&self) -> Result<String, Box<dyn Error>> {
|
||||
let mut stream = TcpStream::connect(self.address.clone())?;
|
||||
stream.write_all("CERT_REQ".as_bytes())?;
|
||||
let mut buf = [0u8; 8192];
|
||||
let bytes_read = stream.read(&mut buf)?;
|
||||
let response = String::from_utf8_lossy(&buf[..bytes_read]);
|
||||
if let Some((_, cert)) = response.split_once("|") {
|
||||
return Ok(cert.to_string());
|
||||
}
|
||||
Err("Server did not send certificate".into())
|
||||
}
|
||||
|
||||
pub fn connect(&mut self) -> Result<usize, Box<dyn Error>> {
|
||||
if let Ok(cert) = self.donwload_certificate() {
|
||||
self.config = Some(self.build_tls_config(cert.as_str())?);
|
||||
}
|
||||
let mut stream = self.tls_connect()?;
|
||||
stream.write("HELLO".as_bytes())?;
|
||||
let mut buf = [0; 8192];
|
||||
let bytes_read = stream.read(&mut buf)?;
|
||||
let response = String::from_utf8_lossy(&buf[..bytes_read]);
|
||||
if let Some((_, id)) = response.split_once("|") {
|
||||
self.client_id = id.trim().parse()?;
|
||||
}
|
||||
self.connected = true;
|
||||
self.last_check = Instant::now();
|
||||
return Ok(self.client_id);
|
||||
}
|
||||
|
||||
pub fn checkin(&mut self) {
|
||||
if let Ok(mut stream) = self.tls_connect() {
|
||||
self.last_check = Instant::now();
|
||||
let payload = format!(
|
||||
"{}|||{}\n",
|
||||
self.client_id.clone(),
|
||||
self.message_que.join("||")
|
||||
);
|
||||
if let Err(e) = stream.write_all(payload.as_bytes()) {
|
||||
self.action_que
|
||||
.push(format!("Error|sending reponse output to server! {e}"));
|
||||
}
|
||||
self.message_que.clear();
|
||||
let mut buffer = [0; 4096];
|
||||
if let Ok(bytes_read) = stream.read(&mut buffer).await {
|
||||
if let Ok(bytes_read) = stream.read(&mut buffer) {
|
||||
if bytes_read > 0 {
|
||||
let response = String::from_utf8_lossy(&buffer[..bytes_read]);
|
||||
if response.starts_with("ACTIONS:") {
|
||||
let actions_str = response.trim_start_matches("ACTIONS:").trim();
|
||||
out_string.1.push_str(actions_str);
|
||||
}
|
||||
response.split("||").into_iter().for_each(|action| {
|
||||
self.action_que.push(action.trim().to_string());
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
return out_string;
|
||||
}
|
||||
|
||||
pub fn build_tls_config(&self, pem: &str) -> Result<Arc<ClientConfig>, Box<dyn Error>> {
|
||||
let mut roots = RootCertStore::empty();
|
||||
let mut reader = BufReader::new(pem.as_bytes());
|
||||
let certs = rustls_pemfile::certs(&mut reader).collect::<Result<Vec<_>, _>>()?;
|
||||
for cert in certs {
|
||||
roots.add(cert)?;
|
||||
}
|
||||
let config = ClientConfig::builder()
|
||||
.with_root_certificates(roots)
|
||||
.with_no_client_auth();
|
||||
Ok(Arc::new(config))
|
||||
}
|
||||
|
||||
fn tls_connect(&self) -> Result<StreamOwned<ClientConnection, TcpStream>, Box<dyn Error>> {
|
||||
let config = self.config.clone().ok_or("TLS config not loaded")?;
|
||||
let tcp = TcpStream::connect(&self.address)?;
|
||||
let (host, _) = self.address.split_once(':').ok_or("Invalid address")?;
|
||||
let server_name = ServerName::try_from(host.to_string())?;
|
||||
let connection = ClientConnection::new(config, server_name)?;
|
||||
let tls_stream = StreamOwned::new(connection, tcp);
|
||||
Ok(tls_stream)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1314,14 +1470,6 @@ impl Host {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum Destination {
|
||||
Server,
|
||||
Victim,
|
||||
Attacker,
|
||||
Control,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct User {
|
||||
pub name: String,
|
||||
@@ -1343,38 +1491,6 @@ impl User {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum ToolMessage {
|
||||
Input((usize, String)),
|
||||
Output((usize, String)),
|
||||
UpdateHost(Host),
|
||||
RebuildDB,
|
||||
DestroyDB(DistroBox),
|
||||
StopDB(DistroBox),
|
||||
StopTemplate(DistroBox),
|
||||
UpdateProject(usize, Project),
|
||||
RemoveProject,
|
||||
AddServer,
|
||||
EndPrompt,
|
||||
ConnectServer,
|
||||
SelectServer,
|
||||
SendServer((Server, String)),
|
||||
ServerBrokerExit,
|
||||
AppStateExit,
|
||||
DisconnectServer,
|
||||
DisconnectAllServers,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum ToolArg {
|
||||
Project(Project),
|
||||
Projects(Vec<Project>),
|
||||
Host(Host),
|
||||
Hosts(Vec<Host>),
|
||||
Config(HashMap<String, String>),
|
||||
Path(PathBuf),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ToolCommand {
|
||||
pub name: String,
|
||||
@@ -1549,12 +1665,6 @@ impl ModuleLoader {
|
||||
}
|
||||
}
|
||||
|
||||
enum AppEvent {
|
||||
Key(event::KeyEvent),
|
||||
Worker(ToolMessage),
|
||||
Mouse(event::MouseEvent),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DistroBox {
|
||||
pub name: String,
|
||||
|
||||
+58
-38
@@ -1,9 +1,10 @@
|
||||
use clap::Parser;
|
||||
use std::sync::Arc;
|
||||
use rcgen::generate_simple_self_signed;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::{env, path::PathBuf, process::exit};
|
||||
use tetanus::AppState;
|
||||
use tetanus::funcs::*;
|
||||
use tokio::sync::Mutex;
|
||||
use tetanus::server::start_server;
|
||||
use tetanus::{AppState, server};
|
||||
|
||||
mod install;
|
||||
|
||||
@@ -29,21 +30,13 @@ struct Args {
|
||||
|
||||
#[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!("checking for server or client config files...");
|
||||
let mut config_path = PathBuf::new();
|
||||
let mut config_path_opt = env::home_dir();
|
||||
if let Some(home_config_path) = config_path_opt.as_mut() {
|
||||
println!("made it to home_config_path");
|
||||
home_config_path.push(".config/tetanus");
|
||||
if !home_config_path.exists() {
|
||||
println!("no config directory found in home directory...");
|
||||
config_path = PathBuf::from("/etc/tetanus");
|
||||
if !config_path.exists() {
|
||||
@@ -55,28 +48,55 @@ async fn main() {
|
||||
exit(0);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
println!("found home config_path.");
|
||||
config_path = home_config_path.clone();
|
||||
}
|
||||
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();
|
||||
println!("entering tui...");
|
||||
let _res = run_tui(appstate, rx).unwrap();
|
||||
} else if args.server {
|
||||
let address = "127.0.0.1:31337".to_string();
|
||||
let names = vec!["127.0.0.1".to_string(), "localhost".to_string()];
|
||||
let mut certificate_path = config_path.clone();
|
||||
certificate_path.push("server");
|
||||
let mut key_path = certificate_path.clone();
|
||||
certificate_path.push("cert.pem");
|
||||
key_path.push("key.pem");
|
||||
println!("key paht: {}", key_path.display());
|
||||
println!("cert path: {}", certificate_path.display());
|
||||
let new_server = server::Server {
|
||||
address,
|
||||
clients: Vec::new(),
|
||||
certificate_path,
|
||||
key_path,
|
||||
};
|
||||
let res = start_server(Arc::new(Mutex::new(new_server))).await;
|
||||
match res {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
println!("error running server {e}");
|
||||
}
|
||||
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();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+317
@@ -0,0 +1,317 @@
|
||||
use rcgen::generate_simple_self_signed;
|
||||
use std::error::Error;
|
||||
use std::fs::File;
|
||||
use std::fs::write;
|
||||
use std::io::BufReader;
|
||||
use std::net::SocketAddr;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::usize;
|
||||
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
use tokio_rustls::TlsAcceptor;
|
||||
use tokio_rustls::rustls::ServerConfig;
|
||||
use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer};
|
||||
|
||||
#[derive(Clone)]
|
||||
enum ClientAction {
|
||||
Output(String),
|
||||
Cmd(String),
|
||||
Ping,
|
||||
}
|
||||
|
||||
pub struct Client {
|
||||
address: String,
|
||||
id: usize,
|
||||
hostname: Option<String>,
|
||||
actions: Vec<ClientAction>,
|
||||
controlling: usize,
|
||||
controlled: usize,
|
||||
output_que: Vec<String>,
|
||||
}
|
||||
|
||||
pub struct Server {
|
||||
pub address: String,
|
||||
pub clients: Vec<Client>,
|
||||
pub certificate_path: PathBuf,
|
||||
pub key_path: PathBuf,
|
||||
}
|
||||
|
||||
pub async fn start_server(server: Arc<Mutex<Server>>) -> Result<(), Box<dyn Error>> {
|
||||
let lock = server.lock().unwrap();
|
||||
if !lock.certificate_path.exists() || !lock.key_path.exists() {
|
||||
let (server_ip, _) = lock.address.split_once(':').unwrap();
|
||||
|
||||
let cert = generate_simple_self_signed(vec![
|
||||
server_ip.to_string(),
|
||||
"127.0.0.1".to_string(),
|
||||
"localhost".to_string(),
|
||||
])?;
|
||||
|
||||
write(&lock.certificate_path, cert.cert.pem())?;
|
||||
|
||||
write(&lock.key_path, cert.signing_key.serialize_pem())?;
|
||||
}
|
||||
|
||||
let certs = load_certs(&lock.certificate_path)?;
|
||||
let key = load_key(&lock.key_path)?;
|
||||
|
||||
let tls_config = ServerConfig::builder()
|
||||
.with_no_client_auth()
|
||||
.with_single_cert(certs, key)?;
|
||||
|
||||
let acceptor = TlsAcceptor::from(Arc::new(tls_config));
|
||||
|
||||
let listener = TcpListener::bind(lock.address.clone()).await?;
|
||||
|
||||
println!("Listening on {}", lock.address);
|
||||
|
||||
drop(lock);
|
||||
|
||||
loop {
|
||||
let (stream, addr) = listener.accept().await?;
|
||||
|
||||
let acceptor = acceptor.clone();
|
||||
let server = server.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut peek = [0u8; 8];
|
||||
if let Ok(n) = stream.peek(&mut peek).await {
|
||||
if n >= 8 && &peek[..8] == b"CERT_REQ" {
|
||||
handle_bootstrap(server, stream).await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
match acceptor.accept(stream).await {
|
||||
Ok(tls_stream) => {
|
||||
handle_connection_stat9ic(server, tls_stream, addr).await;
|
||||
}
|
||||
|
||||
Err(e) => {
|
||||
eprintln!("TLS handshake failed: {}", e);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn load_certs(path: &std::path::Path) -> Result<Vec<CertificateDer<'static>>, Box<dyn Error>> {
|
||||
let mut reader = BufReader::new(File::open(path)?);
|
||||
|
||||
Ok(rustls_pemfile::certs(&mut reader).collect::<Result<Vec<_>, _>>()?)
|
||||
}
|
||||
|
||||
fn load_key(path: &std::path::Path) -> Result<PrivateKeyDer<'static>, Box<dyn Error>> {
|
||||
let mut reader = BufReader::new(File::open(path)?);
|
||||
|
||||
let key = rustls_pemfile::private_key(&mut reader)?.ok_or("No private key found")?;
|
||||
|
||||
Ok(key)
|
||||
}
|
||||
|
||||
async fn handle_bootstrap(server: Arc<Mutex<Server>>, mut stream: TcpStream) {
|
||||
let mut buf = [0u8; 1024];
|
||||
let n = match stream.read(&mut buf).await {
|
||||
Ok(n) => n,
|
||||
Err(_) => return,
|
||||
};
|
||||
if n == 0 {
|
||||
return;
|
||||
}
|
||||
let msg = String::from_utf8_lossy(&buf[..n]);
|
||||
if msg.trim() != "CERT_REQ" {
|
||||
return;
|
||||
}
|
||||
let cert_path = {
|
||||
let lock = server.lock().unwrap();
|
||||
lock.certificate_path.clone()
|
||||
};
|
||||
let cert_pem = match std::fs::read_to_string(cert_path) {
|
||||
Ok(c) => c,
|
||||
Err(_) => return,
|
||||
};
|
||||
let response = format!("CERT|{}\n", cert_pem);
|
||||
let _ = stream.write_all(response.as_bytes()).await;
|
||||
}
|
||||
|
||||
async fn handle_connection_stat9ic<S>(server: Arc<Mutex<Server>>, mut stream: S, addr: SocketAddr)
|
||||
where
|
||||
S: AsyncRead + AsyncWrite + Unpin,
|
||||
{
|
||||
let mut buf = vec![0u8; 4096];
|
||||
let n = match stream.read(&mut buf).await {
|
||||
Ok(n) => n,
|
||||
Err(e) => {
|
||||
eprintln!("Read Error: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
if n == 0 {
|
||||
return;
|
||||
}
|
||||
let msg = String::from_utf8_lossy(&buf[..n]);
|
||||
if msg.contains("HELLO") {
|
||||
let mut connected = false;
|
||||
let mut id = 0;
|
||||
if let Ok(mut lock) = server.lock() {
|
||||
let new_client = Client {
|
||||
address: addr.to_string(),
|
||||
id: lock.clients.len() + 1,
|
||||
hostname: None,
|
||||
actions: Vec::new(),
|
||||
controlled: 0,
|
||||
controlling: 0,
|
||||
output_que: Vec::new(),
|
||||
};
|
||||
connected = true;
|
||||
id = new_client.id.clone();
|
||||
lock.clients.push(new_client);
|
||||
}
|
||||
if connected {
|
||||
stream
|
||||
.write_all(format!("HELLO|{}\n", id).as_bytes())
|
||||
.await
|
||||
.unwrap();
|
||||
println!("client connected! ID:{}", id,);
|
||||
}
|
||||
} else {
|
||||
if let Some((source, data)) = msg.split_once("|||") {
|
||||
if let Ok(source_id) = source.trim().parse::<usize>() {
|
||||
let mut responses = Vec::new();
|
||||
if let Ok(mut lock) = server.lock() {
|
||||
if let Some(source_client) = lock.clients.iter_mut().find(|c| c.id == source_id)
|
||||
{
|
||||
responses = source_client.actions.clone();
|
||||
}
|
||||
data.split("||").into_iter().for_each(|action| {
|
||||
if let Some((cmd, data)) = action.split_once("|") {
|
||||
match cmd.trim() {
|
||||
"OUTPUT" => {
|
||||
if let Some(dest_client) =
|
||||
lock.clients.iter_mut().find(|c| c.controlling == source_id)
|
||||
{
|
||||
dest_client.actions.push(ClientAction::Output(format!(
|
||||
"from {}: {}",
|
||||
dest_client.id,
|
||||
data.trim()
|
||||
)));
|
||||
} else {
|
||||
if let Some(source_client) =
|
||||
lock.clients.iter_mut().find(|c| c.id == source_id)
|
||||
{
|
||||
source_client.output_que.push(data.trim().to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
"CMD" => {
|
||||
if let Some(dest_client) =
|
||||
lock.clients.iter_mut().find(|c| c.controlled == source_id)
|
||||
{
|
||||
dest_client
|
||||
.actions
|
||||
.push(ClientAction::Cmd(data.to_string()));
|
||||
responses.push(ClientAction::Output(format!(
|
||||
"tasked client {} to run {}",
|
||||
dest_client.id, cmd
|
||||
)));
|
||||
}
|
||||
}
|
||||
"CONTROL" => {
|
||||
let new_control_id = data.trim().parse::<usize>().unwrap();
|
||||
if let Some(controlling_client) =
|
||||
lock.clients.iter_mut().find(|c| c.id == source_id)
|
||||
{
|
||||
controlling_client.controlling = new_control_id.clone();
|
||||
}
|
||||
if let Some(contlled_client) =
|
||||
lock.clients.iter_mut().find(|c| c.id == new_control_id)
|
||||
{
|
||||
contlled_client.controlled = source_id;
|
||||
} else {
|
||||
if let Some(controlling_client) =
|
||||
lock.clients.iter_mut().find(|c| c.id == source_id)
|
||||
{
|
||||
controlling_client.controlling = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
"STOP_CONTROL" => {
|
||||
if let Some(source_client) =
|
||||
lock.clients.iter_mut().find(|c| c.id == source_id)
|
||||
{
|
||||
if source_client.controlling != 0 {
|
||||
let dest_client_id = source_client.controlling.clone();
|
||||
if let Some(dest_client) = lock
|
||||
.clients
|
||||
.iter_mut()
|
||||
.find(|c| c.id == dest_client_id)
|
||||
{
|
||||
dest_client.controlled = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"BREAK_CONTROL" => {
|
||||
if let Some(source_client) =
|
||||
lock.clients.iter_mut().find(|c| c.id == source_id)
|
||||
{
|
||||
if source_client.controlled != 0 {
|
||||
let dest_id = source_client.controlled.clone();
|
||||
source_client.controlled = 0;
|
||||
if let Some(dest_client) =
|
||||
lock.clients.iter_mut().find(|c| c.id == dest_id)
|
||||
{
|
||||
dest_client.controlling = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"LIST_CLIENTS" => {
|
||||
lock.clients.iter().for_each(|c| {
|
||||
let out = format!("{}: {}", c.id, c.address);
|
||||
println!("client list requested!");
|
||||
println!("adding {} to response...", out);
|
||||
responses.push(ClientAction::Output(out));
|
||||
});
|
||||
}
|
||||
"TEST" => {
|
||||
responses.push(ClientAction::Output("TEST BACK".to_string()));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
if responses.len() > 0 {
|
||||
let mut messages = Vec::new();
|
||||
for r in responses {
|
||||
match r {
|
||||
ClientAction::Cmd(cmd) => {
|
||||
messages.push(format!("CMD|{}", cmd));
|
||||
}
|
||||
ClientAction::Output(text) => {
|
||||
println!("adding {} to output", text);
|
||||
messages.push(format!("OUTPUT|{}", text));
|
||||
}
|
||||
ClientAction::Ping => {
|
||||
messages.push(format!("PONG"));
|
||||
}
|
||||
}
|
||||
}
|
||||
let full_message = messages.join("||");
|
||||
println!("attempting to send {}", full_message);
|
||||
let buf = full_message.as_bytes();
|
||||
stream.write(buf).await.unwrap();
|
||||
println!("buffer written!");
|
||||
if let Ok(mut lock) = server.lock() {
|
||||
if let Some(source_client) =
|
||||
lock.clients.iter_mut().find(|c| c.id == source_id)
|
||||
{
|
||||
source_client.actions.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user