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

made the server and the client/server communication functionality work.

Browse Source
This commit is contained in:
pyro
2026-08-04 13:35:52 -05:00
parent ae91cc0ce4
commit 4ce461025c
7 changed files with 1162 additions and 438 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lib.rs
+267 -157
View File
@@ -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,