957 lines
34 KiB
Rust
957 lines
34 KiB
Rust
use fs_extra::dir::{CopyOptions, copy};
|
|
use ratatui::crossterm::event;
|
|
use rhai::{AST, Dynamic, Engine, Scope};
|
|
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::path::PathBuf;
|
|
use std::process::{Command, Stdio};
|
|
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: usize,
|
|
pub curent_intput: String,
|
|
pub module_loader: ModuleLoader,
|
|
pub output_scroll: u16,
|
|
prompt: Prompt,
|
|
}
|
|
|
|
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: 0,
|
|
curent_intput: String::new(),
|
|
module_loader: ModuleLoader::new(),
|
|
output_scroll: 0,
|
|
prompt: Prompt {
|
|
action: None,
|
|
responses: Vec::new(),
|
|
execute_command: String::new(),
|
|
num_responses: 0,
|
|
},
|
|
},
|
|
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());
|
|
}
|
|
_ => {
|
|
if parts[0].len() > 1 {
|
|
self.config
|
|
.insert(parts[0].trim().to_string(), parts[1].trim().to_string());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
let mut projet_folder_path = self.config_file.clone();
|
|
projet_folder_path.pop();
|
|
projet_folder_path.push("projects");
|
|
let project_dir_reses = read_dir(projet_folder_path)?;
|
|
for res in project_dir_reses {
|
|
if let Ok(project_folder) = res {
|
|
let project_folder_file_reses = read_dir(project_folder.path())?;
|
|
for res in project_folder_file_reses {
|
|
if let Ok(conf_file) = res {
|
|
if conf_file.file_name().to_string_lossy() == "project.conf".to_string() {
|
|
let mut new_project = Project::new();
|
|
new_project.config_folder(conf_file.path());
|
|
new_project.load_config()?;
|
|
let template_box = self.config.get("template_box").unwrap();
|
|
let tools = self.config.get("tools").unwrap();
|
|
let db = DistroBox {
|
|
name: format!(
|
|
"{}-{}-{}",
|
|
template_box,
|
|
new_project.org_name.clone(),
|
|
new_project.name.clone()
|
|
),
|
|
volumes: vec![
|
|
new_project.files.display().to_string(),
|
|
new_project.notes.display().to_string(),
|
|
tools.clone(),
|
|
],
|
|
created: false,
|
|
template: template_box.clone(),
|
|
};
|
|
new_project.db = Some(db);
|
|
self.projects.push(new_project);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
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!");
|
|
}
|
|
}
|
|
pub fn execute_command(
|
|
&mut self,
|
|
command_name: &str,
|
|
command_args: Option<String>,
|
|
) -> Result<(), Box<dyn Error>> {
|
|
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
|
|
)))?;
|
|
}
|
|
}
|
|
"current_project" | "cp" => {
|
|
let mut out_vec = Vec::new();
|
|
out_vec.push(format!(
|
|
"org: {}",
|
|
&self.projects[self.selected_project].org_name
|
|
));
|
|
out_vec.push(format!(
|
|
"name: {}",
|
|
&self.projects[self.selected_project].name
|
|
));
|
|
out_vec.push(format!(
|
|
"files: {}",
|
|
&self.projects[self.selected_project].files.display()
|
|
));
|
|
out_vec.push(format!(
|
|
"notes: {}",
|
|
&self.projects[self.selected_project].notes.display()
|
|
));
|
|
out_vec.push(format!(
|
|
"config_folder: {}",
|
|
&self.projects[self.selected_project].config_folder.display()
|
|
));
|
|
if let Some(db) = self.projects[self.selected_project].db.clone() {
|
|
out_vec.push(format!("distrobox: {}", db.name));
|
|
}
|
|
if self.projects[self.selected_project].current {
|
|
out_vec.push(format!("status: current"));
|
|
} else {
|
|
out_vec.push(format!("status: upcoming"));
|
|
}
|
|
for line in out_vec {
|
|
let _ = self.main_tx.send(ToolMessage::Output(line));
|
|
}
|
|
}
|
|
"promote_project" | "pp" => {
|
|
self.promote_project()?;
|
|
}
|
|
"remove_project_confirm" => {
|
|
self.prompt.action = None;
|
|
self.prompt.num_responses = 0;
|
|
if self.prompt.responses.len() > 0 {
|
|
if self.prompt.responses[0] == "y".to_string() {
|
|
let project = self.projects[self.selected_project].clone();
|
|
let mut config_file = project.config_folder.clone();
|
|
config_file.pop();
|
|
if let Err(e) = remove_dir_all(&project.config_folder) {
|
|
let _ = self.main_tx.send(ToolMessage::Output(format!(
|
|
"failed to delete config: {e} on {}",
|
|
&project.config_folder.display()
|
|
)));
|
|
}
|
|
if let Err(e) = remove_dir_all(&project.notes) {
|
|
let _ = self.main_tx.send(ToolMessage::Output(format!(
|
|
"failed to delete notes: {e} on {}",
|
|
&project.notes.display()
|
|
)));
|
|
} else if let Some(parent) = project.notes.parent() {
|
|
if let Ok(mut entries) = read_dir(parent) {
|
|
if entries.next().is_none() {
|
|
if let Err(e) = remove_dir(parent) {
|
|
let _ = self.main_tx.send(ToolMessage::Output(format!("failed to delete the empty client notes folder: {e} on {}", parent.display())));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if let Err(e) = remove_dir_all(&project.files) {
|
|
let _ = self.main_tx.send(ToolMessage::Output(format!(
|
|
"failed to delete files: {e} on {}",
|
|
&project.files.display()
|
|
)));
|
|
} else if let Some(parent) = project.files.parent() {
|
|
if let Ok(mut entries) = read_dir(parent) {
|
|
if entries.next().is_none() {
|
|
if let Err(e) = remove_dir(parent) {
|
|
let _ = self.main_tx.send(ToolMessage::Output(format!(
|
|
"failed to delete empty files parent: {e} on {}",
|
|
parent.display()
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
self.projects.remove(self.selected_project);
|
|
let _ = self.main_tx.send(ToolMessage::Output(format!(
|
|
"{}-{} was sucessfully removed!",
|
|
project.org_name, project.name
|
|
)));
|
|
}
|
|
}
|
|
self.prompt.responses.clear();
|
|
}
|
|
"remove_project" | "rp" => {
|
|
let project = self.projects[self.selected_project].clone();
|
|
self.prompt.action = Some(ToolMessage::RemoveProject);
|
|
self.prompt.num_responses = 1;
|
|
self.prompt.execute_command = String::from("remove_project_confirm");
|
|
let _ = self.main_tx.send(ToolMessage::Output(format!(
|
|
"{}, {} and all contents will be deleted. Continue? (y/N)",
|
|
project.files.display(),
|
|
project.notes.display()
|
|
)));
|
|
}
|
|
_ => {
|
|
let ast = self
|
|
.module_loader
|
|
.asts
|
|
.get(command_name)
|
|
.ok_or_else(|| format!("Command not found: {}", 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();
|
|
let mut scope = Scope::new();
|
|
for arg_requirement in &cmd_meta.args {
|
|
match arg_requirement.as_str() {
|
|
"project" => {
|
|
scope.push("project", self.projects[self.selected_project].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" => {
|
|
scope.push("host", "todo");
|
|
}
|
|
"hosts" => {
|
|
scope.push("hosts", self.projects[self.selected_project].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);
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
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
|
|
)));
|
|
}
|
|
}
|
|
});
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
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));
|
|
let mut template_path = self.config_file.clone();
|
|
let mut project_conf_folder = self.config_file.clone();
|
|
let tools = self.config.get("tools").unwrap();
|
|
template_path.pop();
|
|
template_path.push("note_templates");
|
|
project_conf_folder.pop();
|
|
project_conf_folder.push(format!("projects/{}-{}", org, name));
|
|
let mut options = CopyOptions::new();
|
|
options.overwrite = true;
|
|
match create_dir_all(&project_files) {
|
|
Ok(_) => {}
|
|
Err(e) => {
|
|
let _ = self.main_tx.send(ToolMessage::Output(format!(
|
|
"error making project files! {e} on {}",
|
|
project_files.display()
|
|
)));
|
|
}
|
|
}
|
|
match create_dir_all(&project_notes) {
|
|
Ok(_) => {}
|
|
Err(e) => {
|
|
let _ = self.main_tx.send(ToolMessage::Output(format!(
|
|
"error making project notes! {e} on {}",
|
|
project_notes.display()
|
|
)));
|
|
}
|
|
}
|
|
match create_dir_all(&project_conf_folder) {
|
|
Ok(_) => {}
|
|
Err(e) => {
|
|
let _ = self.main_tx.send(ToolMessage::Output(format!(
|
|
"error making config folder! {e} on {}",
|
|
project_conf_folder.display()
|
|
)));
|
|
}
|
|
}
|
|
for entry in read_dir(template_path)? {
|
|
let entry = entry?;
|
|
let template_name_os = entry.file_name();
|
|
let template_name = template_name_os.to_string_lossy();
|
|
if name.clone().contains(template_name.as_ref()) {
|
|
for entry in read_dir(entry.path())? {
|
|
let entry = entry?;
|
|
let file_name_os = entry.file_name();
|
|
let file_name = file_name_os.to_string_lossy();
|
|
if entry.file_type()?.is_dir() {
|
|
let dest = project_notes.join(file_name.as_ref());
|
|
match copy(entry.path(), dest, &options) {
|
|
Ok(_) => {}
|
|
Err(e) => {
|
|
let _ = self.main_tx.send(ToolMessage::Output(format!(
|
|
"error copying note template! {e}"
|
|
)));
|
|
}
|
|
}
|
|
} else {
|
|
match fs::copy(entry.path(), project_notes.join(file_name.as_ref())) {
|
|
Ok(_) => {}
|
|
Err(e) => {
|
|
let _ = self.main_tx.send(ToolMessage::Output(format!(
|
|
"error copying note template! {e}"
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
let db = DistroBox {
|
|
name: format!("{}-{}-{}", template_box, org, name),
|
|
volumes: vec![
|
|
project_files.display().to_string(),
|
|
project_notes.display().to_string(),
|
|
tools.to_string(),
|
|
],
|
|
created: false,
|
|
template: template_box.to_string(),
|
|
};
|
|
let mut new_project = Project {
|
|
org_name: org,
|
|
name: name,
|
|
notes: project_notes,
|
|
files: project_files,
|
|
hosts: Vec::new(),
|
|
config_folder: project_conf_folder,
|
|
current: false,
|
|
db: Some(db),
|
|
};
|
|
new_project.save_config()?;
|
|
self.projects.push(new_project);
|
|
return Ok(());
|
|
}
|
|
|
|
pub fn promote_project(&mut self) -> Result<String, Box<dyn Error>> {
|
|
let mut return_string = String::from("Project promoted successfully!");
|
|
let mut new_project = self.projects[self.selected_project].clone();
|
|
let new_files_path = PathBuf::from(self.config.get("current_files").unwrap())
|
|
.join(format!("{}/{}", new_project.org_name, new_project.name));
|
|
let new_notes_path = PathBuf::from(self.config.get("current_notes").unwrap())
|
|
.join(format!("{}/{}", new_project.org_name, new_project.name));
|
|
let mut options = CopyOptions::new();
|
|
create_dir_all(&new_files_path)?;
|
|
create_dir_all(&new_notes_path)?;
|
|
options.overwrite = true;
|
|
copy(new_project.files.clone(), new_files_path.clone(), &options)?;
|
|
copy(new_project.notes.clone(), new_notes_path.clone(), &options)?;
|
|
let cleanup_res = Command::new("rm")
|
|
.arg("-rf")
|
|
.arg(new_project.files.display().to_string())
|
|
.arg(new_project.notes.display().to_string())
|
|
.status();
|
|
if cleanup_res.is_err() {
|
|
return_string = format!(
|
|
"Error deleting {} and {}, please clean up manually. Otherwise Project promotion succeeded!",
|
|
new_files_path.display(),
|
|
new_notes_path.display()
|
|
);
|
|
}
|
|
new_project.files(new_files_path.clone());
|
|
new_project.notes(new_notes_path.clone());
|
|
if let Some(mut db) = new_project.db {
|
|
db.volumes = vec![
|
|
new_files_path.display().to_string(),
|
|
new_notes_path.display().to_string(),
|
|
self.config.get("tools").unwrap().to_string(),
|
|
];
|
|
new_project.db = Some(db);
|
|
let _ = self.main_tx.send(ToolMessage::RebuildDB);
|
|
}
|
|
self.projects[self.selected_project] = new_project;
|
|
return Ok(return_string);
|
|
}
|
|
}
|
|
|
|
#[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,
|
|
pub db: Option<DistroBox>,
|
|
}
|
|
|
|
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,
|
|
db: None,
|
|
}
|
|
}
|
|
|
|
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(());
|
|
}
|
|
|
|
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 {
|
|
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),
|
|
RebuildDB,
|
|
UpdateProject(usize, Project),
|
|
RemoveProject,
|
|
}
|
|
|
|
#[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>,
|
|
pub asts: HashMap<String, AST>,
|
|
}
|
|
|
|
impl ModuleLoader {
|
|
pub fn new() -> Self {
|
|
let mut engine = Engine::new();
|
|
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),
|
|
);
|
|
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);
|
|
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);
|
|
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> {
|
|
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) {
|
|
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;
|
|
}
|
|
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;
|
|
}
|
|
let help = read_to_string(&help_path).unwrap_or_default();
|
|
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;
|
|
}
|
|
};
|
|
let command = ToolCommand {
|
|
name,
|
|
path: script_path,
|
|
help,
|
|
args,
|
|
output_type,
|
|
finished: false,
|
|
result: false,
|
|
};
|
|
|
|
Some((command, ast))
|
|
}
|
|
}
|
|
|
|
enum AppEvent {
|
|
Key(event::KeyEvent),
|
|
Worker(ToolMessage),
|
|
Mouse(event::MouseEvent),
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct DistroBox {
|
|
pub name: String,
|
|
pub volumes: Vec<String>,
|
|
pub created: bool,
|
|
pub template: String,
|
|
}
|
|
|
|
impl DistroBox {
|
|
pub fn create(&mut self, tx: Sender<ToolMessage>) -> Result<(), Box<dyn Error>> {
|
|
let mut create_command = Command::new("distrobox");
|
|
create_command
|
|
.arg("create")
|
|
.arg("--root")
|
|
.arg("--clone")
|
|
.arg(self.template.clone())
|
|
.arg("--name")
|
|
.arg(self.name.clone());
|
|
for volume in &self.volumes {
|
|
let folder_name = volume.split("/").last().unwrap();
|
|
create_command
|
|
.arg("--volume")
|
|
.arg(format!("{}:/{}:rw", volume, folder_name));
|
|
}
|
|
create_command.stdout(Stdio::piped());
|
|
create_command.stderr(Stdio::piped());
|
|
let mut child = create_command.spawn()?;
|
|
if let Some(stdout) = child.stdout.take() {
|
|
let tx_out = tx.clone();
|
|
std::thread::spawn(move || {
|
|
let reader = BufReader::new(stdout);
|
|
for line in reader.lines().flatten() {
|
|
let _ = tx_out.send(ToolMessage::Output(line));
|
|
}
|
|
});
|
|
}
|
|
if let Some(stderr) = child.stderr.take() {
|
|
let tx_out = tx.clone();
|
|
std::thread::spawn(move || {
|
|
let reader = BufReader::new(stderr);
|
|
for line in reader.lines().flatten() {
|
|
let _ = tx_out.send(ToolMessage::Output(line));
|
|
}
|
|
});
|
|
}
|
|
let status = child.wait()?;
|
|
if !status.success() {
|
|
let _ = tx.send(ToolMessage::Output(format!("error creating distrobox!")));
|
|
}
|
|
|
|
self.created = true;
|
|
return Ok(());
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct Prompt {
|
|
action: Option<ToolMessage>,
|
|
responses: Vec<String>,
|
|
execute_command: String,
|
|
num_responses: usize,
|
|
}
|