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

added a new module and the ability to specify a config file if you want

Browse Source 
to.
This commit is contained in:
pyro
2026-08-12 17:05:35 -05:00
parent 4ce461025c
commit 7b1f9f9d4e
10 changed files with 208 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/funcs.rs
+2 -3
View File
@@ -5,16 +5,15 @@ use crossterm::{
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::thread::sleep;
use std::thread::spawn;
use std::time::Duration;
use std::{error::Error, time::Instant};
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);
src/lib.rs
+111 -16
View File
@@ -689,6 +689,62 @@ impl AppState {
}
}
}
"show_config" | "sc" => {
let _ = self
.main_tx
.send(ToolMessage::Output((0, "Current Config:".to_string())));
self.config.iter().for_each(|setting| {
let _ = self.main_tx.send(ToolMessage::Output((
0,
format!("{}:{}", setting.0, setting.1),
)));
});
}
"set" => {
let mut ready = false;
let mut key = String::new();
let mut value = String::new();
if let Some(args) = command_args.clone() {
let args_vec: Vec<&str> = args.split(" ").collect();
if args_vec.len() == 2 {
key = args_vec[0].to_string();
value = args_vec[1].to_string();
if self.config.contains_key(&key) {
ready = true;
} else {
let _ = self.main_tx.send(ToolMessage::Output((
0,
"Error Unknown setting! {key}, {value}".to_string(),
)));
return Ok(());
}
}
}
if ready {
self.config.insert(key, value);
self.save_config()?;
let _ = self.main_tx.send(ToolMessage::Output((
0,
"Setting changed successfully!".to_string(),
)));
} else {
let _ = self.main_tx.send(ToolMessage::Output((
0,
"Error malformed command!".to_string(),
)));
let _ = self.main_tx.send(ToolMessage::Output((
0,
"Usag: set setting value".to_string(),
)));
}
}
"list_modules" => {
self.module_loader.commands.iter().for_each(|module| {
let _ = self
.main_tx
.send(ToolMessage::Output((0, format!("{}", module.0))));
});
}
"exit" => {
self.save_all()?;
let _ = self.main_tx.send(ToolMessage::AppStateExit);
@@ -752,6 +808,7 @@ impl AppState {
_ => {}
}
}
let worker_tx = self.main_tx.clone();
self.workers.spawn(move || {
match engine.eval_ast_with_scope::<Dynamic>(&mut scope, &ast) {
Ok(result) => {
@@ -768,9 +825,14 @@ impl AppState {
"failed to process array output".into()
}
} else if result.is_string() {
result
let out_string = result
.into_string()
.unwrap_or_else(|_| "failed to parse string".into())
.unwrap_or_else(|_| "failed to parse string".into());
if cmd_meta.calls != "none" {
let _ =
worker_tx.send(ToolMessage::Input((0, out_string.clone())));
}
out_string
} else {
format!("{:?}", result)
};
@@ -795,6 +857,16 @@ impl AppState {
Ok(())
}
pub fn edit_config(&mut self, key: String, value: String) {
if self.config.contains_key(&key) {
self.config.insert(key.clone(), value.clone());
}
let _ = self.main_tx.send(ToolMessage::Output((
0,
format!("Setting saved! {key}:{value}"),
)));
}
pub fn new_project(
&mut self,
org: String,
@@ -1086,6 +1158,7 @@ impl Server {
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}"));
@@ -1500,6 +1573,7 @@ pub struct ToolCommand {
pub output_type: String,
pub finished: bool,
pub result: bool,
pub calls: String,
}
pub struct ModuleLoader {
@@ -1575,6 +1649,33 @@ impl ModuleLoader {
engine.register_fn("add_port", |host: &mut Host, port: i64| {
host.add_port(port as usize);
});
engine.register_fn("expand_target", |s: &str| -> rhai::Array {
let mut results = rhai::Array::new();
if let Ok(net) = s.trim().parse::<IpNet>() {
for host in net.hosts() {
results.push(host.to_string().into());
}
} else if let Ok(ip) = s.trim().parse::<IpAddr>() {
results.push(ip.to_string().into());
} else if s.contains(".") {
results.push(s.to_string().into());
}
results
});
engine.register_fn("write_file", |path: &str, contents: &str| {
std::fs::write(path, contents).is_ok()
});
engine.register_fn("find_scope_file", |path: PathBuf| -> String {
for entry in walkdir::WalkDir::new(path).max_depth(10) {
if let Ok(entry) = entry {
if entry.file_name().to_string_lossy().contains("scope") {
return entry.path().display().to_string();
}
}
}
String::new()
});
Self {
engine: Arc::new(engine),
commands: HashMap::new(),
@@ -1623,6 +1724,7 @@ impl ModuleLoader {
let mut name = String::new();
let mut output_type = String::new();
let mut args = Vec::new();
let mut calls = String::new();
for line in config_content.lines() {
if let Some((key, val)) = line.split_once(':') {
match key.trim() {
@@ -1635,6 +1737,9 @@ impl ModuleLoader {
.filter(|s| !s.is_empty())
.collect();
}
"calls" => {
calls = val.trim().to_string();
}
_ => {}
}
}
@@ -1659,6 +1764,7 @@ impl ModuleLoader {
output_type,
finished: false,
result: false,
calls,
};
Some((command, ast))
@@ -1685,7 +1791,6 @@ impl DistroBox {
let mut create_command = Command::new("distrobox");
create_command
.arg("create")
.arg("--root")
.arg("--clone")
.arg(self.template.clone())
.arg("--name")
@@ -1743,7 +1848,6 @@ impl DistroBox {
let mut stop_command = Command::new("distrobox");
stop_command
.arg("stop")
.arg("--root")
.arg(self.template.clone())
.arg("--yes");
stop_command.stdin(Stdio::piped());
@@ -1782,11 +1886,7 @@ impl DistroBox {
pub fn stop(&mut self, tx: Sender<ToolMessage>, rid: usize) -> Result<(), Box<dyn Error>> {
let mut stop_command = Command::new("distrobox");
stop_command
.arg("stop")
.arg("--root")
.arg(self.name.clone())
.arg("--yes");
stop_command.arg("stop").arg(self.name.clone()).arg("--yes");
stop_command.stdin(Stdio::piped());
stop_command.stdout(Stdio::piped());
stop_command.stderr(Stdio::piped());
@@ -1823,11 +1923,7 @@ impl DistroBox {
pub fn destroy(&mut self, tx: Sender<ToolMessage>, rid: usize) -> Result<(), Box<dyn Error>> {
let mut destroycmd = Command::new("distrobox");
destroycmd
.arg("rm")
.arg("--root")
.arg(self.name.clone())
.arg("-f");
destroycmd.arg("rm").arg(self.name.clone()).arg("-f");
let res = destroycmd.status()?;
if !res.success() {
let _ = tx.send(ToolMessage::Output((
@@ -1840,7 +1936,7 @@ impl DistroBox {
)));
let _ = tx.send(ToolMessage::Output((
rid,
format!("distrobox rm --root {} -f", self.name),
format!("distrobox rm {} -f", self.name),
)));
}
self.created = false;
@@ -1855,7 +1951,6 @@ impl DistroBox {
term_cmd
.arg("distrobox")
.arg("enter")
.arg("--root")
.arg(self.name.clone());
} else if arg.contains("ENV_NAME=") {
let (_, env_name) = arg.split_once("=").unwrap();
src/main.rs
+31 -7
View File
@@ -1,5 +1,4 @@
use clap::Parser;
use rcgen::generate_simple_self_signed;
use std::sync::{Arc, Mutex};
use std::{env, path::PathBuf, process::exit};
use tetanus::funcs::*;
@@ -26,14 +25,38 @@ struct Args {
#[arg(short, long, help = "for testing stuff...")]
test: bool,
#[arg(
short,
long,
help = "The IP address for the server, in server mode it will listen on this IP, in client mode it will add this ip as its first server. Default is 127.0.0.1"
)]
ip: Option<String>,
#[arg(
short,
long,
help = "The port for the server, in server mode this will be the port the server listens on, in client mode this will be the default server port it connects to."
)]
port: Option<String>,
#[arg(
short = 'C',
long,
help = "Custom Config file. Give the path to a config file to use."
)]
config: Option<PathBuf>,
}
#[tokio::main]
async fn main() {
println!("checking for server or client config files...");
let args = Args::parse();
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() {
if let Some(given_config_path) = args.config {
config_path = given_config_path.clone();
} else 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() {
@@ -53,11 +76,8 @@ async fn main() {
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!(
@@ -76,8 +96,7 @@ async fn main() {
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 address = "127.0.0.1:31337".to_string();
let mut certificate_path = config_path.clone();
certificate_path.push("server");
let mut key_path = certificate_path.clone();
@@ -85,6 +104,11 @@ async fn main() {
key_path.push("key.pem");
println!("key paht: {}", key_path.display());
println!("cert path: {}", certificate_path.display());
if let Some(ip) = args.ip {
if let Some(port) = args.port {
address = format!("{}:{}", ip.trim(), port.trim());
}
}
let new_server = server::Server {
address,
clients: Vec::new(),
src/server.rs
+19 -5
View File
@@ -17,7 +17,6 @@ use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer};
enum ClientAction {
Output(String),
Cmd(String),
Ping,
}
pub struct Client {
@@ -240,6 +239,10 @@ where
lock.clients.iter_mut().find(|c| c.id == source_id)
{
if source_client.controlling != 0 {
println!(
"deslelcting client: {}",
source_client.controlling
);
let dest_client_id = source_client.controlling.clone();
if let Some(dest_client) = lock
.clients
@@ -247,6 +250,7 @@ where
.find(|c| c.id == dest_client_id)
{
dest_client.controlled = 0;
println!("client {} deselected", dest_client.id);
}
}
}
@@ -268,7 +272,12 @@ where
}
"LIST_CLIENTS" => {
lock.clients.iter().for_each(|c| {
let out = format!("{}: {}", c.id, c.address);
let out;
if let Some(hostname) = c.hostname.clone() {
out = format!("{}: {}({})", c.id, hostname, c.address);
} else {
out = format!("{}: {}", c.id, c.address);
}
println!("client list requested!");
println!("adding {} to response...", out);
responses.push(ClientAction::Output(out));
@@ -277,6 +286,14 @@ where
"TEST" => {
responses.push(ClientAction::Output("TEST BACK".to_string()));
}
"SET_HOSTNAME" => {
lock.clients
.iter_mut()
.filter(|c| c.id == source_id)
.for_each(|c| {
c.hostname = Some(data.trim().to_string());
});
}
_ => {}
}
}
@@ -293,9 +310,6 @@ where
println!("adding {} to output", text);
messages.push(format!("OUTPUT|{}", text));
}
ClientAction::Ping => {
messages.push(format!("PONG"));
}
}
}
let full_message = messages.join("||");