wrote the tool to use rhai scripts as commands, and wrote a basic cli as
an example.
This commit is contained in:
+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))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user