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

changed the cli to a tui environment.

Browse Source
This commit is contained in:
pyro
2026-05-20 13:33:53 -05:00
parent 3eeecb0010
commit a907de8ff3
6 changed files with 1376 additions and 217 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lib.rs
+171 -124
View File
@@ -1,8 +1,9 @@
use ratatui::crossterm::event;
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::fs::{File, create_dir_all, read_dir, read_to_string};
use std::io::{self, Write};
use std::path::PathBuf;
use std::sync::mpsc::Receiver;
use std::sync::{Arc, mpsc::Sender, mpsc::channel};
@@ -24,6 +25,7 @@ pub struct AppState {
pub selected_project: Option<Project>,
pub curent_intput: String,
pub module_loader: ModuleLoader,
pub output_scroll: u16,
}
impl AppState {
@@ -48,6 +50,7 @@ impl AppState {
selected_project: None,
curent_intput: String::new(),
module_loader: ModuleLoader::new(),
output_scroll: 0,
},
main_rx,
)
@@ -114,118 +117,151 @@ impl AppState {
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,
&mut self,
command_name: &str,
raw_input_string: Option<String>,
command_args: 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 tx = self.main_tx.clone();
match command_name {
"new_project" | "np" => {
let args = command_args.unwrap();
let (org, name) = args.split_once(" ").unwrap();
if let Err(e) = self.new_project(org.to_string(), name.to_string()) {
tx.send(ToolMessage::Output(format!(
"Error making {}-{}: {e}",
org, name
)))?;
} else {
tx.send(ToolMessage::Output(format!(
"{}-{} created successfully!",
org, name
)))?;
}
}
_ => {
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());
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();
let mut scope = Scope::new();
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" => {
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") => {
let input_val = command_args
.clone()
.unwrap_or_else(|| self.curent_intput.clone());
scope.push("input_string", input_val);
}
_ => {}
}
}
"projects" => {
let mut arr = rhai::Array::new();
for proj in &self.projects {
arr.push(rhai::Dynamic::from(proj.clone()));
self.workers.spawn(move || {
match engine.eval_ast_with_scope::<Dynamic>(&mut scope, &ast) {
Ok(result) => {
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)
};
if result_str.contains("\n") {
result_str.lines().into_iter().for_each(|line| {
let _ = tx.send(ToolMessage::Output(line.to_string()));
});
} else {
let _ = tx.send(ToolMessage::Output(result_str));
}
}
Err(err) => {
let _ = tx.send(ToolMessage::Output(format!(
"Script Error in execution: {}",
err
)));
}
}
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) {}
pub fn new_project(&mut self, org: String, name: String) -> Result<(), Box<dyn Error>> {
let template_box = self.config.get("template_box").unwrap();
let project_files = PathBuf::from(self.config.get("upcoming_files").unwrap())
.join(format!("{}/{}", org, name));
let project_notes = PathBuf::from(self.config.get("upcoming_notes").unwrap())
.join(format!("{}/{}", org, name));
create_dir_all(&project_files)?;
create_dir_all(&project_notes)?;
let mut new_project = Project::new();
let mut project_conf_folder = self.config_file.clone();
project_conf_folder.pop();
project_conf_folder.push(format!("projects/{}-{}", org, name));
new_project.config_folder(project_conf_folder);
new_project.box_name(format!("{}-{}-{}", template_box, org, name));
new_project.name = name;
new_project.org_name(org);
new_project.files(project_files);
new_project.notes(project_notes);
new_project.save_config()?;
self.projects.push(new_project);
return Ok(());
}
}
#[derive(Clone)]
@@ -243,6 +279,7 @@ pub struct Project {
pub files: PathBuf,
pub hosts: Vec<Host>,
pub current: bool,
pub boxname: String,
}
impl Project {
@@ -255,6 +292,7 @@ impl Project {
files: PathBuf::new(),
hosts: Vec::new(),
current: false,
boxname: String::new(),
}
}
@@ -282,6 +320,10 @@ impl Project {
self.config_folder = path;
}
pub fn box_name(&mut self, boxname: String) {
self.boxname = boxname;
}
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() {
@@ -304,6 +346,27 @@ impl Project {
println!("{} | {} loaded!", self.org_name, self.name);
return Ok(());
}
pub fn save_config(&mut self) -> Result<(), Box<dyn Error>> {
create_dir_all(&self.config_folder)?;
let mut conf_file = self.config_folder.clone();
conf_file.push("project.conf");
let mut file = File::create(conf_file)?;
let mut out_string = format!(
"org_name: {}\nname: {}\nnotes: {}\nfiles: {}\n",
self.org_name,
self.name,
self.notes.display(),
self.files.display()
);
if self.current {
out_string.push_str("stage: current");
} else {
out_string.push_str("stage: upcoming");
}
file.write_all(out_string.as_bytes())?;
return Ok(());
}
}
impl std::fmt::Display for Project {
@@ -440,14 +503,12 @@ pub struct ToolCommand {
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(
@@ -455,8 +516,6 @@ impl ModuleLoader {
|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(
@@ -476,8 +535,6 @@ impl ModuleLoader {
)
.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(
@@ -496,10 +553,6 @@ impl ModuleLoader {
|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)
});
@@ -511,7 +564,6 @@ impl ModuleLoader {
}
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();
@@ -532,7 +584,6 @@ impl ModuleLoader {
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);
}
@@ -549,13 +600,10 @@ impl ModuleLoader {
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() {
@@ -572,15 +620,10 @@ impl ModuleLoader {
}
}
}
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) => {
@@ -588,8 +631,6 @@ impl ModuleLoader {
return None;
}
};
// 4. Assemble your exact ToolCommand struct
let command = ToolCommand {
name,
path: script_path,
@@ -603,3 +644,9 @@ impl ModuleLoader {
Some((command, ast))
}
}
enum AppEvent {
Key(event::KeyEvent),
Worker(ToolMessage),
Mouse(event::MouseEvent),
}