wrote the tool to use rhai scripts as commands, and wrote a basic cli as
an example.
This commit is contained in:
Generated
+4444
File diff suppressed because it is too large
Load Diff
+12
@@ -0,0 +1,12 @@
|
||||
[package]
|
||||
name = "tetanus"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
clap = { version = "4.6.1", features = ["derive"] }
|
||||
iced = { version = "0.14.0", features = ["advanced", "tokio"] }
|
||||
rayon = "1.12.0"
|
||||
rhai = { version = "1.24.0", features = ["metadata", "sync"] }
|
||||
rustc-hash = "2.1.2"
|
||||
walkdir = "2.5.0"
|
||||
@@ -0,0 +1,85 @@
|
||||
use crate::*;
|
||||
use std::error::Error;
|
||||
use std::io::Write;
|
||||
use std::thread;
|
||||
|
||||
pub fn get_user_input(prompt: &str) -> Result<String, Box<dyn Error>> {
|
||||
println!("{}", prompt);
|
||||
let mut response = String::new();
|
||||
std::io::stdin().read_line(&mut response)?;
|
||||
return Ok(response.trim().to_string());
|
||||
}
|
||||
|
||||
pub fn cli(mut state: AppState, main_rx: Receiver<ToolMessage>) {
|
||||
println!("Starting tetanus CLI...");
|
||||
let (input_tx, input_rx) = channel::<String>();
|
||||
thread::spawn(move || {
|
||||
let stdin = io::stdin();
|
||||
let mut input_buffer = String::new();
|
||||
loop {
|
||||
input_buffer.clear();
|
||||
if stdin.read_line(&mut input_buffer).is_ok() {
|
||||
let trimmed = input_buffer.trim().to_string();
|
||||
if !trimmed.is_empty() {
|
||||
if input_tx.send(trimmed).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
print!("tetanus> ");
|
||||
let _ = io::stdout().flush();
|
||||
loop {
|
||||
while let Ok(msg) = main_rx.try_recv() {
|
||||
match msg {
|
||||
ToolMessage::Output(txt) => {
|
||||
println!("\n[result] {}", txt);
|
||||
state.log.push(txt.clone());
|
||||
state.output.push(txt.clone());
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
while let Ok(user_input) = input_rx.try_recv() {
|
||||
state.curent_intput = user_input.clone();
|
||||
state.history.push(user_input.clone());
|
||||
match user_input.as_str() {
|
||||
"exit" | "quit" => {
|
||||
println!("shutting down...");
|
||||
return;
|
||||
}
|
||||
"reload-modules" => {
|
||||
println!("reloadling modules...");
|
||||
state.initialize_modules();
|
||||
}
|
||||
"help" => {
|
||||
println!("\nAvailable Built-In & Custom Script Modules:");
|
||||
for (name, cmd) in &state.module_loader.commands {
|
||||
println!(" - {:<15} (Outputs: {})", name, cmd.output_type);
|
||||
if !cmd.help.is_empty() {
|
||||
println!(" Help: {}", cmd.help.trim());
|
||||
}
|
||||
}
|
||||
println!();
|
||||
}
|
||||
// Everything else gets dynamically routed to your Rhai engine execution matrix!
|
||||
command_name => {
|
||||
if state.module_loader.commands.contains_key(command_name) {
|
||||
println!(
|
||||
"[Worker] Dispatching script '{}' to Rayon engine layer...",
|
||||
command_name
|
||||
);
|
||||
if let Err(e) = state.execute_command(command_name, None) {
|
||||
eprintln!("[Error] Failed running script: {}", e);
|
||||
}
|
||||
} else {
|
||||
println!("[Error] Command '{}' unrecognized.", command_name);
|
||||
}
|
||||
}
|
||||
}
|
||||
print!("tetanus> ");
|
||||
let _ = io::stdout().flush();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
use std::env;
|
||||
use std::error::Error;
|
||||
use std::fs::{File, create_dir_all};
|
||||
use std::io::Write;
|
||||
use std::path::PathBuf;
|
||||
use std::process::exit;
|
||||
use tetanus::funcs::get_user_input;
|
||||
|
||||
pub fn install() -> Result<(), Box<dyn Error>> {
|
||||
if let Some(homedir) = env::home_dir() {
|
||||
let mut config_path = homedir.clone();
|
||||
config_path.push(".config/tetanus");
|
||||
let mut projects_path = config_path.clone();
|
||||
projects_path.push("projects");
|
||||
let mut module_path = config_path.clone();
|
||||
module_path.push("modules/default");
|
||||
create_dir_all(&module_path)?;
|
||||
module_path.pop();
|
||||
module_path.push("custom");
|
||||
create_dir_all(&module_path)?;
|
||||
module_path.pop();
|
||||
let current_files_path = PathBuf::from(get_user_input(
|
||||
"enter the path to store currently in progress project files (not notes), or none if you want to be propted everytime",
|
||||
)?);
|
||||
let current_notes_path = PathBuf::from(get_user_input(
|
||||
"enter the path to store currently in progress project notes, or none if you want to be prompted everytime",
|
||||
)?);
|
||||
let upcoming_files_path = PathBuf::from(get_user_input(
|
||||
"enter the path to store upcoming project files (not notes) or none if you want to be prompted everytime",
|
||||
)?);
|
||||
let upcoming_notes_path = PathBuf::from(get_user_input(
|
||||
"enter the path to store upcoming project notes, or none if you want to be prompted everytime",
|
||||
)?);
|
||||
config_path.push("client.conf");
|
||||
projects_path.push("default");
|
||||
create_dir_all(&projects_path)?;
|
||||
projects_path.push("project.conf");
|
||||
let mut client_config_file = File::create(&config_path)?;
|
||||
client_config_file.write(format!("projects: {}\n", projects_path.display()).as_bytes())?;
|
||||
client_config_file.write("servers: 127.0.0.1:31337\n".as_bytes())?;
|
||||
let mut default_project_file = File::create(projects_path)?;
|
||||
config_path.pop();
|
||||
config_path.push("server.conf");
|
||||
let mut server_config_file = File::create(config_path)?;
|
||||
server_config_file.write("address: 127.0.0.1:31337\n".as_bytes())?;
|
||||
server_config_file.write("running: false\n".as_bytes())?;
|
||||
default_project_file.write("org_name: default\n".as_bytes())?;
|
||||
default_project_file.write("name: default\n".as_bytes())?;
|
||||
default_project_file
|
||||
.write(format!("notes: {}\n", current_files_path.display()).as_bytes())?;
|
||||
default_project_file
|
||||
.write(format!("files: {}\n", current_notes_path.display()).as_bytes())?;
|
||||
default_project_file.write("stage: current".as_bytes())?;
|
||||
client_config_file
|
||||
.write(format!("current_files: {}\n", current_files_path.display()).as_bytes())?;
|
||||
client_config_file
|
||||
.write(format!("current_notes: {}\n", current_notes_path.display()).as_bytes())?;
|
||||
client_config_file
|
||||
.write(format!("upcoming_files: {}\n", upcoming_files_path.display()).as_bytes())?;
|
||||
client_config_file
|
||||
.write(format!("upcoming_notes: {}\n", upcoming_notes_path.display()).as_bytes())?;
|
||||
client_config_file.write(format!("module_path: {}\n", module_path.display()).as_bytes())?;
|
||||
} else {
|
||||
eprintln!("error finding home directory!");
|
||||
exit(1);
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
+605
@@ -0,0 +1,605 @@
|
||||
use rhai::{AST, Dynamic, Engine, Scope};
|
||||
use std::collections::HashMap;
|
||||
use std::error::Error;
|
||||
use std::fs::{read_dir, read_to_string};
|
||||
use std::io;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::mpsc::Receiver;
|
||||
use std::sync::{Arc, mpsc::Sender, mpsc::channel};
|
||||
|
||||
pub mod funcs;
|
||||
|
||||
pub struct AppState {
|
||||
pub projects: Vec<Project>,
|
||||
pub servers: Vec<Server>,
|
||||
pub config: HashMap<String, String>,
|
||||
pub config_file: PathBuf,
|
||||
pub workers: rayon::ThreadPool,
|
||||
pub worker_txes: Vec<Sender<ToolMessage>>,
|
||||
pub main_tx: Sender<ToolMessage>,
|
||||
pub history: Vec<String>,
|
||||
pub log: Vec<String>,
|
||||
pub output: Vec<String>,
|
||||
pub selected_server: Option<Server>,
|
||||
pub selected_project: Option<Project>,
|
||||
pub curent_intput: String,
|
||||
pub module_loader: ModuleLoader,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
pub fn new() -> (Self, Receiver<ToolMessage>) {
|
||||
let (main_tx, main_rx) = channel();
|
||||
(
|
||||
Self {
|
||||
projects: Vec::new(),
|
||||
servers: Vec::new(),
|
||||
config: HashMap::new(),
|
||||
config_file: PathBuf::new(),
|
||||
workers: rayon::ThreadPoolBuilder::new()
|
||||
.num_threads(4)
|
||||
.build()
|
||||
.unwrap(),
|
||||
worker_txes: Vec::new(),
|
||||
main_tx,
|
||||
history: Vec::new(),
|
||||
log: Vec::new(),
|
||||
output: Vec::new(),
|
||||
selected_server: None,
|
||||
selected_project: None,
|
||||
curent_intput: String::new(),
|
||||
module_loader: ModuleLoader::new(),
|
||||
},
|
||||
main_rx,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn load_config(&mut self, file: PathBuf) -> Result<(), Box<dyn Error>> {
|
||||
self.config_file = file.clone();
|
||||
let config_contents = read_to_string(file)?;
|
||||
for line in config_contents.lines() {
|
||||
let parts: Vec<&str> = line.split(": ").collect();
|
||||
if parts.len() == 2 {
|
||||
match parts[0].trim() {
|
||||
"servers" => {
|
||||
let addresses: Vec<&str> = parts[1].trim().split(", ").collect();
|
||||
for address in addresses {
|
||||
let new_server = Server {
|
||||
address: address.trim().to_string(),
|
||||
connected: false,
|
||||
};
|
||||
self.servers.push(new_server);
|
||||
}
|
||||
self.config
|
||||
.insert("servers".to_string(), line.trim().to_string());
|
||||
}
|
||||
"projects" => {
|
||||
let strs: Vec<&str> = parts[1].split(", ").collect();
|
||||
let mut paths = Vec::new();
|
||||
for path in strs {
|
||||
paths.push(PathBuf::from(path));
|
||||
}
|
||||
for path in paths {
|
||||
let mut new_project = Project::new();
|
||||
new_project.config_folder(path);
|
||||
new_project.load_config()?;
|
||||
self.projects.push(new_project);
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
if parts[0].len() > 1 {
|
||||
self.config
|
||||
.insert(parts[0].trim().to_string(), parts[1].trim().to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
pub fn initialize_modules(&mut self) {
|
||||
if let Some(base_path_str) = self.config.get("module_path") {
|
||||
let base_path = PathBuf::from(base_path_str);
|
||||
|
||||
println!("Loading Tetanus modules from: {:?}", base_path);
|
||||
if let Err(e) = self.module_loader.load_all(&base_path) {
|
||||
eprintln!("Failed to load modules: {}", e);
|
||||
} else {
|
||||
println!(
|
||||
"Successfully loaded {} modules.",
|
||||
self.module_loader.commands.len()
|
||||
);
|
||||
}
|
||||
} else {
|
||||
eprintln!("Warning: 'module_path' is missing from AppState config map!");
|
||||
}
|
||||
}
|
||||
|
||||
/// Dispatches a command execution payload to the Rayon thread pool
|
||||
pub fn execute_command(
|
||||
&self,
|
||||
command_name: &str,
|
||||
raw_input_string: Option<String>,
|
||||
) -> Result<(), Box<dyn Error>> {
|
||||
// 1. Retrieve the compiled AST and Engine reference
|
||||
let ast = self
|
||||
.module_loader
|
||||
.asts
|
||||
.get(command_name)
|
||||
.ok_or_else(|| format!("AST not found for command: {}", command_name))?
|
||||
.clone();
|
||||
|
||||
let engine = Arc::clone(&self.module_loader.engine);
|
||||
let cmd_meta = self
|
||||
.module_loader
|
||||
.commands
|
||||
.get(command_name)
|
||||
.ok_or_else(|| format!("Metadata not found for command: {}", command_name))?
|
||||
.clone();
|
||||
|
||||
// 2. Build out the script Scope using values cloned from AppState
|
||||
let mut scope = Scope::new();
|
||||
|
||||
// Dynamically add requested args into script context based on config.conf parsing rules
|
||||
for arg_requirement in &cmd_meta.args {
|
||||
match arg_requirement.as_str() {
|
||||
"project" => {
|
||||
if let Some(ref proj) = self.selected_project {
|
||||
scope.push("project", proj.clone());
|
||||
}
|
||||
}
|
||||
"projects" => {
|
||||
let mut arr = rhai::Array::new();
|
||||
for proj in &self.projects {
|
||||
arr.push(rhai::Dynamic::from(proj.clone()));
|
||||
}
|
||||
scope.push("projects", arr);
|
||||
}
|
||||
"host" => {
|
||||
if let Some(host) = self.selected_project.as_ref().and_then(|p| p.hosts.first())
|
||||
{
|
||||
scope.push("host", host.clone());
|
||||
}
|
||||
}
|
||||
"hosts" => {
|
||||
if let Some(ref proj) = self.selected_project {
|
||||
scope.push("hosts", proj.hosts.clone());
|
||||
}
|
||||
}
|
||||
"config" => {
|
||||
// Turn Hashmap into rhai::Map
|
||||
let mut map = rhai::Map::new();
|
||||
for (k, v) in &self.config {
|
||||
map.insert(k.clone().into(), v.clone().into());
|
||||
}
|
||||
scope.push("config", map);
|
||||
}
|
||||
_ if arg_requirement.starts_with("string") => {
|
||||
// If it asks for an input string, supply what the user provided or the current interactive line input
|
||||
let input_val = raw_input_string
|
||||
.clone()
|
||||
.unwrap_or_else(|| self.curent_intput.clone());
|
||||
scope.push("input_string", input_val);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Clone our main communication MPSC tx channel to ship results back out of the Rayon worker threads
|
||||
let tx = self.main_tx.clone();
|
||||
|
||||
// 4. Delegate to Rayon
|
||||
self.workers.spawn(move || {
|
||||
// This runs in a background thread inside Rayon!
|
||||
match engine.eval_ast_with_scope::<Dynamic>(&mut scope, &ast) {
|
||||
Ok(result) => {
|
||||
// Handle output based on the output type requested
|
||||
let result_str = if result.is_array() {
|
||||
if let Ok(arr) = result.into_array() {
|
||||
let lines: Vec<String> = arr
|
||||
.into_iter()
|
||||
.map(|item| item.into_string().unwrap_or_else(|_| "".into()))
|
||||
.collect();
|
||||
lines.join("\n")
|
||||
} else {
|
||||
"failed to process array output".into()
|
||||
}
|
||||
} else if result.is_string() {
|
||||
result
|
||||
.into_string()
|
||||
.unwrap_or_else(|_| "failed to parse string".into())
|
||||
} else {
|
||||
format!("{:?}", result)
|
||||
};
|
||||
let _ = tx.send(ToolMessage::Output(result_str));
|
||||
}
|
||||
Err(err) => {
|
||||
let _ = tx.send(ToolMessage::Output(format!(
|
||||
"Script Error in execution: {}",
|
||||
err
|
||||
)));
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn new_project(&mut self) {}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Server {
|
||||
pub address: String,
|
||||
pub connected: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Project {
|
||||
pub org_name: String,
|
||||
pub config_folder: PathBuf,
|
||||
pub name: String,
|
||||
pub notes: PathBuf,
|
||||
pub files: PathBuf,
|
||||
pub hosts: Vec<Host>,
|
||||
pub current: bool,
|
||||
}
|
||||
|
||||
impl Project {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
org_name: String::new(),
|
||||
config_folder: PathBuf::new(),
|
||||
name: String::new(),
|
||||
notes: PathBuf::new(),
|
||||
files: PathBuf::new(),
|
||||
hosts: Vec::new(),
|
||||
current: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn org_name(&mut self, name: String) {
|
||||
self.org_name = name;
|
||||
}
|
||||
|
||||
pub fn notes(&mut self, path: PathBuf) {
|
||||
self.notes = path;
|
||||
}
|
||||
|
||||
pub fn files(&mut self, path: PathBuf) {
|
||||
self.files = path;
|
||||
}
|
||||
|
||||
pub fn hosts(&mut self, hosts: Vec<Host>) {
|
||||
self.hosts = hosts;
|
||||
}
|
||||
|
||||
pub fn add_host(&mut self, host: Host) {
|
||||
self.hosts.push(host);
|
||||
}
|
||||
|
||||
pub fn config_folder(&mut self, path: PathBuf) {
|
||||
self.config_folder = path;
|
||||
}
|
||||
|
||||
pub fn load_config(&mut self) -> Result<(), Box<dyn Error>> {
|
||||
let config_contents = read_to_string(&self.config_folder)?;
|
||||
for line in config_contents.lines() {
|
||||
let parts: Vec<&str> = line.split(": ").collect();
|
||||
if parts.len() == 2 {
|
||||
match parts[0].trim() {
|
||||
"org_name" => self.org_name(parts[1].trim().to_string()),
|
||||
"name" => self.name = parts[1].trim().to_string(),
|
||||
"notes" => self.notes(PathBuf::from(parts[1].trim())),
|
||||
"files" => self.files(PathBuf::from(parts[1].trim())),
|
||||
"stage" => {
|
||||
if parts[1].trim() == "current" {
|
||||
self.current = true;
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
println!("{} | {} loaded!", self.org_name, self.name);
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for Project {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{} | {}", self.org_name, self.name)
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for Project {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
format!("{} | {}", self.org_name, self.name)
|
||||
== format!("{} | {}", other.org_name, other.name)
|
||||
}
|
||||
}
|
||||
|
||||
impl Eq for Project {}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Host {
|
||||
pub ip: String,
|
||||
pub hostname: String,
|
||||
pub open_ports: Vec<usize>,
|
||||
pub control_port: usize,
|
||||
pub users: Vec<User>,
|
||||
pub pwned: bool,
|
||||
pub id: usize,
|
||||
pub findings: Vec<String>,
|
||||
}
|
||||
|
||||
impl Host {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
ip: String::new(),
|
||||
hostname: String::new(),
|
||||
open_ports: Vec::new(),
|
||||
control_port: 0,
|
||||
users: Vec::new(),
|
||||
pwned: false,
|
||||
id: 0,
|
||||
findings: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn ip(&mut self, ip: String) {
|
||||
self.ip = ip;
|
||||
}
|
||||
|
||||
pub fn hostname(&mut self, name: String) {
|
||||
self.hostname = name;
|
||||
}
|
||||
|
||||
pub fn open_ports(&mut self, ports: Vec<usize>) {
|
||||
self.open_ports = ports;
|
||||
}
|
||||
|
||||
pub fn control_port(&mut self, port: usize) {
|
||||
self.control_port = port;
|
||||
}
|
||||
|
||||
pub fn users(&mut self, users: Vec<User>) {
|
||||
self.users = users;
|
||||
}
|
||||
|
||||
pub fn pwnd(&mut self) {
|
||||
self.pwned = true;
|
||||
}
|
||||
|
||||
pub fn id(&mut self, id: usize) {
|
||||
self.id = id;
|
||||
}
|
||||
|
||||
pub fn findings(&mut self, findings: Vec<String>) {
|
||||
self.findings = findings;
|
||||
}
|
||||
|
||||
pub fn add_port(&mut self, port: usize) {
|
||||
self.open_ports.push(port);
|
||||
}
|
||||
|
||||
pub fn add_user(&mut self, user: User) {
|
||||
self.users.push(user);
|
||||
}
|
||||
|
||||
pub fn add_finding(&mut self, finding: String) {
|
||||
self.findings.push(finding);
|
||||
}
|
||||
}
|
||||
|
||||
pub enum Destination {
|
||||
Server,
|
||||
Victim,
|
||||
Attacker,
|
||||
Control,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct User {
|
||||
pub name: String,
|
||||
pub password: Option<String>,
|
||||
pub hash: Option<String>,
|
||||
pub ticket: Option<PathBuf>,
|
||||
pub compromised: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum ToolMessage {
|
||||
Input(String),
|
||||
Output(String),
|
||||
UpdateHost(Host),
|
||||
UpdateProject(Project),
|
||||
}
|
||||
|
||||
#[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,
|
||||
pub path: PathBuf,
|
||||
pub help: String,
|
||||
pub args: Vec<String>,
|
||||
pub output_type: String,
|
||||
pub finished: bool,
|
||||
pub result: bool,
|
||||
}
|
||||
|
||||
pub struct ModuleLoader {
|
||||
pub engine: Arc<Engine>,
|
||||
pub commands: HashMap<String, ToolCommand>,
|
||||
// Stored separately because AST doesn't implement Debug/Clone easily
|
||||
pub asts: HashMap<String, AST>,
|
||||
}
|
||||
|
||||
impl ModuleLoader {
|
||||
pub fn new() -> Self {
|
||||
let mut engine = Engine::new();
|
||||
// 1. Register PathBuf
|
||||
engine
|
||||
.register_type_with_name::<PathBuf>("PathBuf")
|
||||
.register_get_set(
|
||||
"display",
|
||||
|p: &mut PathBuf| p.to_string_lossy().into_owned(),
|
||||
|p: &mut PathBuf, s: String| *p = PathBuf::from(s),
|
||||
);
|
||||
|
||||
// 2. Register Host
|
||||
engine
|
||||
.register_type_with_name::<Host>("Host")
|
||||
.register_get_set(
|
||||
"ip",
|
||||
|h: &mut Host| h.ip.clone(),
|
||||
|h: &mut Host, val: String| h.ip = val,
|
||||
)
|
||||
.register_get_set(
|
||||
"hostname",
|
||||
|h: &mut Host| h.hostname.clone(),
|
||||
|h: &mut Host, val: String| h.hostname = val,
|
||||
)
|
||||
.register_get_set(
|
||||
"pwned",
|
||||
|h: &mut Host| h.pwned,
|
||||
|h: &mut Host, val: bool| h.pwned = val,
|
||||
)
|
||||
.register_fn("add_port", Host::add_port)
|
||||
.register_fn("add_finding", Host::add_finding);
|
||||
|
||||
// 3. Register Project
|
||||
engine
|
||||
.register_type_with_name::<Project>("Project")
|
||||
.register_get_set(
|
||||
"name",
|
||||
|p: &mut Project| p.name.clone(),
|
||||
|p: &mut Project, val: String| p.name = val,
|
||||
)
|
||||
.register_get_set(
|
||||
"org_name",
|
||||
|p: &mut Project| p.org_name.clone(),
|
||||
|p: &mut Project, val: String| p.org_name = val,
|
||||
)
|
||||
.register_get_set(
|
||||
"current",
|
||||
|p: &mut Project| p.current.clone(),
|
||||
|p: &mut Project, val: bool| p.current = val,
|
||||
)
|
||||
.register_fn("add_host", Project::add_host);
|
||||
|
||||
// 4. Register custom Vector indexers / custom array conversion helpers
|
||||
// Rhai works naturally with its own rhai::Array (Vec<Dynamic>).
|
||||
// To bridge Vec<Host> to Rhai scripts seamlessly, we can expose specialized utility methods.
|
||||
engine.register_fn("get_host", |hosts: &mut Vec<Host>, index: i64| {
|
||||
hosts.get(index as usize).cloned().unwrap_or_else(Host::new)
|
||||
});
|
||||
Self {
|
||||
engine: Arc::new(engine),
|
||||
commands: HashMap::new(),
|
||||
asts: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn load_all(&mut self, base_path: &PathBuf) -> Result<(), std::io::Error> {
|
||||
// Clear out existing maps in case this is called during a hot-reload
|
||||
self.commands.clear();
|
||||
self.asts.clear();
|
||||
|
||||
self.load_from_dir(&base_path.join("default"))?;
|
||||
self.load_from_dir(&base_path.join("custom"))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn load_from_dir(&mut self, dir: &PathBuf) -> Result<(), std::io::Error> {
|
||||
if !dir.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
for entry in read_dir(dir)? {
|
||||
let entry = entry?;
|
||||
let path = entry.path();
|
||||
|
||||
if path.is_dir() {
|
||||
if let Some((command, ast)) = self.parse_module_directory(&path) {
|
||||
// Cache both side-by-side using the parsed command name as the key
|
||||
self.asts.insert(command.name.clone(), ast);
|
||||
self.commands.insert(command.name.clone(), command);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_module_directory(&self, dir: &PathBuf) -> Option<(ToolCommand, AST)> {
|
||||
let config_path = dir.join("config.conf");
|
||||
let help_path = dir.join("help.txt");
|
||||
let script_path = dir.join("script.rhai");
|
||||
|
||||
if !config_path.exists() || !help_path.exists() || !script_path.exists() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// 1. Parse the config.conf file line by line
|
||||
let config_content = read_to_string(&config_path).ok()?;
|
||||
let mut name = String::new();
|
||||
let mut output_type = String::new();
|
||||
let mut args = Vec::new();
|
||||
|
||||
for line in config_content.lines() {
|
||||
if let Some((key, val)) = line.split_once(':') {
|
||||
match key.trim() {
|
||||
"name" => name = val.trim().to_string(),
|
||||
"output_type" => output_type = val.trim().to_string(),
|
||||
"args" => {
|
||||
args = val
|
||||
.split(',')
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.collect();
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if name.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// 2. Read the help text
|
||||
let help = read_to_string(&help_path).unwrap_or_default();
|
||||
|
||||
// 3. Compile the Rhai script into an AST
|
||||
let ast = match self.engine.compile_file(script_path.clone()) {
|
||||
Ok(compiled_ast) => compiled_ast,
|
||||
Err(e) => {
|
||||
eprintln!("Error compiling script in {:?}: {}", script_path, e);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
// 4. Assemble your exact ToolCommand struct
|
||||
let command = ToolCommand {
|
||||
name,
|
||||
path: script_path,
|
||||
help,
|
||||
args,
|
||||
output_type,
|
||||
finished: false,
|
||||
result: false,
|
||||
};
|
||||
|
||||
Some((command, ast))
|
||||
}
|
||||
}
|
||||
+64
@@ -0,0 +1,64 @@
|
||||
use clap::Parser;
|
||||
use std::{env, path::PathBuf, process::exit};
|
||||
use tetanus::AppState;
|
||||
use tetanus::funcs::cli;
|
||||
|
||||
mod install;
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(
|
||||
author,
|
||||
version,
|
||||
about = "The Tetanus Redteaming tool. This will start the client, server, or both"
|
||||
)]
|
||||
struct Args {
|
||||
#[arg(short, long, help = "start in client mode")]
|
||||
client: bool,
|
||||
|
||||
#[arg(short, long, help = "start in server mode")]
|
||||
server: bool,
|
||||
|
||||
#[arg(short, long, help = "start gui client")]
|
||||
gui: bool,
|
||||
}
|
||||
|
||||
fn main() {
|
||||
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!("no config directory found in home directory...");
|
||||
config_path = PathBuf::from("/etc/tetanus");
|
||||
if !config_path.exists() {
|
||||
println!("no config directory found in /etc/tetanus... Installing!");
|
||||
let install_result = install::install();
|
||||
if install_result.is_ok() {
|
||||
install_result.unwrap();
|
||||
println!("install succeeded!");
|
||||
exit(0);
|
||||
}
|
||||
}
|
||||
}
|
||||
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();
|
||||
cli(appstate, rx);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user