pyro
0bb3a125aa
being passed correctly. Also started the groundwork for adding more of my rust based tooling to tetanus.
2967 lines
126 KiB
Rust
2967 lines
126 KiB
Rust
use clipboard::ClipboardContext;
|
|
use clipboard::ClipboardProvider;
|
|
use fs_extra::dir::{CopyOptions, copy};
|
|
use ipnet::IpNet;
|
|
use keyring::Entry;
|
|
use ratatui::crossterm::event;
|
|
use rhai::{AST, Dynamic, Engine, Scope};
|
|
use rustls::{ClientConfig, ClientConnection, RootCertStore, StreamOwned, pki_types::ServerName};
|
|
use std::collections::HashMap;
|
|
use std::env;
|
|
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, Read, Write};
|
|
use std::net::IpAddr;
|
|
use std::net::TcpStream;
|
|
use std::path::PathBuf;
|
|
use std::process::{Command, Stdio};
|
|
use std::sync::mpsc::Receiver;
|
|
use std::sync::{Arc, Mutex, mpsc::Sender, mpsc::channel};
|
|
use std::time::Duration;
|
|
use std::time::{self, Instant};
|
|
use walkdir::WalkDir;
|
|
|
|
use crate::rust_tools::rustwitness;
|
|
|
|
pub mod funcs;
|
|
pub mod rust_tools;
|
|
pub mod server;
|
|
pub mod victim;
|
|
|
|
enum AppEvent {
|
|
Key(event::KeyEvent),
|
|
Worker(ToolMessage),
|
|
Mouse(event::MouseEvent),
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub enum Destination {
|
|
Server,
|
|
Victim,
|
|
Attacker,
|
|
Control,
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub enum ToolMessage {
|
|
Input((usize, String)),
|
|
Output((usize, String)),
|
|
UpdateHost(Host),
|
|
RebuildDB,
|
|
DestroyDB(DistroBox),
|
|
StopDB(DistroBox),
|
|
StopTemplate(DistroBox),
|
|
UpdateProject(usize, Project),
|
|
RemoveProject,
|
|
AddServer,
|
|
EndPrompt,
|
|
ConnectServer,
|
|
SelectServer,
|
|
SendServer((usize, String)),
|
|
ServerBrokerExit,
|
|
AppStateExit,
|
|
DisconnectServer,
|
|
DisconnectAllServers,
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub enum ToolArg {
|
|
Project(Project),
|
|
Projects(Vec<Project>),
|
|
Host(Host),
|
|
Hosts(Vec<Host>),
|
|
Config(HashMap<String, String>),
|
|
Path(PathBuf),
|
|
}
|
|
|
|
pub struct AppState {
|
|
pub projects: Vec<Project>,
|
|
pub servers: Vec<Arc<Mutex<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: usize,
|
|
pub selected_project: usize,
|
|
pub curent_intput: String,
|
|
pub module_loader: ModuleLoader,
|
|
pub output_scroll: u16,
|
|
pub output_follow: bool,
|
|
prompt: Prompt,
|
|
pub info_scroll: u16,
|
|
pub server_broker_running: bool,
|
|
pub app_state_running: bool,
|
|
pub remoting: bool,
|
|
pub help: Vec<String>,
|
|
pub name: String,
|
|
}
|
|
|
|
impl AppState {
|
|
pub fn new() -> Self {
|
|
let (main_tx, _) = 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: 0,
|
|
selected_project: 0,
|
|
curent_intput: String::new(),
|
|
module_loader: ModuleLoader::new(),
|
|
output_scroll: 0,
|
|
output_follow: true,
|
|
prompt: Prompt {
|
|
action: None,
|
|
responses: Vec::new(),
|
|
execute_command: String::new(),
|
|
num_responses: 0,
|
|
prompts: Vec::new(),
|
|
last_prompted: None,
|
|
},
|
|
info_scroll: 0,
|
|
server_broker_running: true,
|
|
app_state_running: true,
|
|
remoting: false,
|
|
help: Vec::new(),
|
|
name: String::new(),
|
|
}
|
|
}
|
|
|
|
pub fn load_config(&mut self, file: PathBuf, display: bool) -> 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" => {
|
|
for server_data in parts[1].trim().split(", ").collect::<Vec<&str>>() {
|
|
if let Some((address, name)) = server_data.split_once("|") {
|
|
let new_server = Server {
|
|
address: address.trim().to_string(),
|
|
connected: false,
|
|
timer: Duration::from_secs(60),
|
|
last_check: time::Instant::now(),
|
|
message_que: Vec::new(),
|
|
action_que: Vec::new(),
|
|
client_id: 0,
|
|
selected_client: 0,
|
|
config: None,
|
|
logged_in: false,
|
|
password: String::new(),
|
|
name: name.to_string(),
|
|
cert_text: String::new(),
|
|
};
|
|
self.servers.push(Arc::new(Mutex::new(new_server)));
|
|
}
|
|
}
|
|
self.config
|
|
.insert("servers".to_string(), parts[1].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(display)?;
|
|
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);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
self.help.push("help\nThat's MEEEEEEE\n".to_string());
|
|
self.help.push("new_project\nadd a project to the upcoming project pool and create default notes based on templates\nusage:\nnew_project organization_name, project_name\n".to_string());
|
|
self.help.push(
|
|
"current_project\nshow information on the currently selected project\n".to_string(),
|
|
);
|
|
self.help
|
|
.push("promote_project\npromote a project from upcoming to current\n".to_string());
|
|
self.help.push(
|
|
"remove_project\nremove a project, delete its folders, and destroy its distrobox\n"
|
|
.to_string(),
|
|
);
|
|
self.help.push("new_terminal\nopen a new terminal in the distrobox of the currently selected project\n".to_string());
|
|
self.help
|
|
.push("new_server\ncreate a new server object\n".to_string());
|
|
self.help
|
|
.push("list_servers\nlist all servers currently tracked by the tool\n".to_string());
|
|
self.help
|
|
.push("connect_server\nconnect to a tracked server\n".to_string());
|
|
self.help
|
|
.push("save_all\nsave all projects and settings\n".to_string());
|
|
self.help
|
|
.push("save_projects\nsave just the projects\n".to_string());
|
|
self.help
|
|
.push("test_server\ntest the server connection\n".to_string());
|
|
self.help.push(
|
|
"add_scope\nadd a host, cidr range, or ip address to the current project's scope\n"
|
|
.to_string(),
|
|
);
|
|
self.help.push("parse_scope\nparse the projects general.md notes file for the scope copied from the workbook\n".to_string());
|
|
self.help.push("parse_portscan | pps\nParse a services.tsv file exported from cobalt stirke in your current project's files directory".to_string());
|
|
self.help
|
|
.push("stop_db\nrestart this project's distrobox".to_string());
|
|
self.help.push("rustwitness\nRun a scan on the hosts in the urls.txt file in yoru project files directory that captures screenshots of those urls, optionally through a proxy.".to_string());
|
|
self.help.push("exit\nquit the tool\n".to_string());
|
|
self.initialize_modules();
|
|
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 {
|
|
for module in &self.module_loader.commands {
|
|
self.help.push(format!(
|
|
"{}\n{}\n",
|
|
module.1.name.clone(),
|
|
module.1.help.clone()
|
|
));
|
|
}
|
|
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>,
|
|
rid: usize,
|
|
) -> Result<(), Box<dyn Error>> {
|
|
let tx = self.main_tx.clone();
|
|
match command_name {
|
|
"remote" => {
|
|
if let Some(server) = self.servers.get(self.selected_server) {
|
|
if let Ok(mut server) = server.lock() {
|
|
if server.selected_client != 0 {
|
|
if let Some(args) = command_args.clone() {
|
|
let args = args.trim();
|
|
if args.contains(" ") {
|
|
let (cmd, iargs) = args.split_once(" ").unwrap();
|
|
server.message_que.push(format!("CMD|{} {}", cmd, iargs));
|
|
} else {
|
|
server.message_que.push(format!("CMD|{}", args));
|
|
}
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
0,
|
|
format!("Tasked {} to run {}", server.selected_client, args),
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
"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(), rid) {
|
|
tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("[error] making {}-{}: {e}", org, name),
|
|
)))?;
|
|
} else {
|
|
tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("[success] {}-{} created!", org, name),
|
|
)))?;
|
|
}
|
|
}
|
|
"generate_victim" => {
|
|
let mut success = false;
|
|
if let Some(args) = command_args {
|
|
if args.split(" ").into_iter().count() == 1 {
|
|
let mut address = String::new();
|
|
let mut password = String::new();
|
|
let config_file = self.config_file.clone();
|
|
let files = self.projects[self.selected_project].files.clone();
|
|
let mut cert_text = String::new();
|
|
if let Ok(lock) = self.servers[self.selected_server].lock() {
|
|
address = lock.address.clone();
|
|
password = lock.password.clone();
|
|
cert_text = lock.cert_text.clone();
|
|
}
|
|
match victim::generate_code(
|
|
address,
|
|
password,
|
|
config_file,
|
|
files,
|
|
args,
|
|
self.main_tx.clone(),
|
|
cert_text,
|
|
) {
|
|
Ok(_) => {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
"[Success] Victim code generated! Compilation thread spawned!"
|
|
.to_string(),
|
|
)));
|
|
success = true;
|
|
}
|
|
Err(e) => {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
0,
|
|
format!("[error] generating victim code: {e}"),
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if !success {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
0,
|
|
format!("Usage: generate_victim (windows/linux)"),
|
|
)));
|
|
}
|
|
}
|
|
"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"));
|
|
}
|
|
out_vec.push(format!(
|
|
"Num hosts: {}",
|
|
&self.projects[self.selected_project].hosts.iter().count()
|
|
));
|
|
for line in out_vec {
|
|
let _ = self.main_tx.send(ToolMessage::Output((rid, line)));
|
|
}
|
|
}
|
|
"rebuild_distro_box" | "rdb" => {
|
|
if let Some(db) = self.projects[self.selected_project].db.clone() {
|
|
if db.created {
|
|
let _ = self.main_tx.send(ToolMessage::DestroyDB(db.clone()));
|
|
}
|
|
let _ = self.main_tx.send(ToolMessage::RebuildDB);
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
"[success] Distorbox made!".to_string(),
|
|
)));
|
|
}
|
|
}
|
|
"promote_project" | "pp" => {
|
|
self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
String::from("promote_project command recieved!"),
|
|
)))?;
|
|
match self.promote_project(rid) {
|
|
Ok(_) => {
|
|
self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
String::from("[success] project promoted!"),
|
|
)))?;
|
|
}
|
|
Err(e) => {
|
|
self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("[error] promoting project: {}", e),
|
|
)))?;
|
|
}
|
|
}
|
|
}
|
|
"remove_project_confirm" => {
|
|
if self.prompt.responses.len() > 0 {
|
|
if self.prompt.responses[0]
|
|
.response
|
|
.to_lowercase()
|
|
.contains("y")
|
|
{
|
|
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(config_file) {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
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((
|
|
rid,
|
|
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((rid, 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((
|
|
rid,
|
|
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((
|
|
rid,
|
|
format!(
|
|
"failed to delete empty files parent: {e} on {}",
|
|
parent.display()
|
|
),
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
let remove_id = self.selected_project.clone();
|
|
let project = self.projects[remove_id].clone();
|
|
if let Some(db) = project.db {
|
|
self.main_tx.send(ToolMessage::DestroyDB(db))?;
|
|
}
|
|
self.selected_project -= 1;
|
|
self.projects.remove(remove_id);
|
|
self.prompt.responses.clear();
|
|
self.prompt.action = None;
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!(
|
|
"{}-{} was sucessfully removed!",
|
|
project.org_name, project.name
|
|
),
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
"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");
|
|
self.prompt.prompts.push(PromptQuery {
|
|
id: 0,
|
|
query: format!(
|
|
"{}, {} and all contentes wil lbe deleted. Contineue? (y/N",
|
|
project.files.display(),
|
|
project.notes.display()
|
|
),
|
|
});
|
|
let _ = self
|
|
.main_tx
|
|
.send(ToolMessage::Input((rid, "prompt".to_string())));
|
|
}
|
|
"new_terminal" | "nt" => {
|
|
let project = self.projects[self.selected_project].clone();
|
|
if let Some(db) = project.db {
|
|
if let Some(cmd) = self.config.get("term_cmd") {
|
|
db.launch_terminal(cmd.to_string());
|
|
}
|
|
}
|
|
}
|
|
"new_server" | "ns" => {
|
|
self.prompt.num_responses = 2;
|
|
self.prompt.action = Some(ToolMessage::AddServer);
|
|
self.prompt.execute_command = String::from("add_server");
|
|
self.prompt.prompts.push(PromptQuery {
|
|
id: 0,
|
|
query: "What is the IP of the server?".to_string(),
|
|
});
|
|
self.prompt.prompts.push(PromptQuery {
|
|
id: 1,
|
|
query: "What is the port for the server?".to_string(),
|
|
});
|
|
let _ = self
|
|
.main_tx
|
|
.send(ToolMessage::Input((0, "prompt".to_string())));
|
|
}
|
|
"add_server" => {
|
|
let mut ip = String::new();
|
|
let mut port = String::new();
|
|
self.prompt.responses.iter().for_each(|res| {
|
|
if res.query.id == 0 {
|
|
ip = res.response.clone();
|
|
} else if res.query.id == 1 {
|
|
port = res.response.clone();
|
|
}
|
|
});
|
|
let new_server = Server {
|
|
address: format!("{}:{}", ip, port),
|
|
connected: false,
|
|
timer: Duration::from_secs(5),
|
|
last_check: time::Instant::now(),
|
|
message_que: Vec::new(),
|
|
action_que: Vec::new(),
|
|
client_id: 0,
|
|
selected_client: 0,
|
|
config: None,
|
|
logged_in: false,
|
|
password: String::new(),
|
|
name: String::new(),
|
|
cert_text: String::new(),
|
|
};
|
|
self.servers.push(Arc::new(Mutex::new(new_server)));
|
|
self.prompt.reset();
|
|
}
|
|
"connect_server" => {
|
|
let _ = self.main_tx.send(ToolMessage::ConnectServer);
|
|
}
|
|
"prompt" => {
|
|
if self.prompt.responses.len() == self.prompt.num_responses {
|
|
let _ = self
|
|
.main_tx
|
|
.send(ToolMessage::Input((0, self.prompt.execute_command.clone())));
|
|
let _ = self.main_tx.send(ToolMessage::EndPrompt);
|
|
} else {
|
|
for query in &self.prompt.prompts {
|
|
if !self
|
|
.prompt
|
|
.responses
|
|
.iter()
|
|
.any(|res| res.query.id == query.id)
|
|
{
|
|
self.prompt.last_prompted = Some(query.clone());
|
|
let _ = self
|
|
.main_tx
|
|
.send(ToolMessage::Output((rid, query.query.clone())));
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
"list_servers" => {
|
|
self.servers.iter().enumerate().for_each(|(id, m)| {
|
|
if let Ok(s) = m.lock() {
|
|
if id == self.selected_server {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!(
|
|
"Server ID:{} Name:{} Address:{} ClientID:{} Sleep:{} Last:{} Loggedin: {} - Currently Selected",
|
|
id, s.name, s.address, s.client_id, s.timer.as_secs(), s.last_check.elapsed().as_secs(), s.logged_in
|
|
),
|
|
)));
|
|
} else {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!(
|
|
"Server ID:{} Name:{} Address:{} ClientID:{} Sleep:{} Last:{}, loggedin: {}",
|
|
id, s.name, s.address, s.client_id, s.timer.as_secs(), s.last_check.elapsed().as_secs(), s.logged_in
|
|
),
|
|
)));
|
|
}
|
|
}
|
|
});
|
|
}
|
|
"remove_server" => {
|
|
let _ = self.main_tx.send(ToolMessage::DisconnectServer);
|
|
}
|
|
"server_remove_confim" => {
|
|
let args = command_args.unwrap();
|
|
let id = args.trim().parse::<usize>().unwrap();
|
|
self.servers.remove(id);
|
|
if self.selected_server == id && self.selected_server >= 1 {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
0,
|
|
format!("[success] server {} removed!", self.selected_server),
|
|
)));
|
|
self.selected_server = self.selected_server - 1;
|
|
}
|
|
}
|
|
"select_server" => {
|
|
self.prompt.action = Some(ToolMessage::SelectServer);
|
|
self.prompt.num_responses = 1;
|
|
self.prompt.prompts.push(PromptQuery {
|
|
id: 0,
|
|
query: String::from("Selection?"),
|
|
});
|
|
self.prompt.execute_command = String::from("server_selected");
|
|
self.servers.iter().enumerate().for_each(|(id, servermut)| {
|
|
if let Ok(server) = servermut.lock() {
|
|
if id == self.selected_server {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("{}:{} - Currently Selected", id, server.address.clone()),
|
|
)));
|
|
} else {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("{}: {}", id, server.address.clone()),
|
|
)));
|
|
}
|
|
}
|
|
});
|
|
}
|
|
"server_selected" => {
|
|
if self.prompt.responses.len() == 1 {
|
|
if let Ok(id) = self.prompt.responses[0].response.parse::<usize>() {
|
|
self.selected_server = id;
|
|
if let Ok(server) = self.servers[id].lock() {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("{} Selected!", server.address),
|
|
)));
|
|
}
|
|
} else {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
String::from("Invalid selection."),
|
|
)));
|
|
}
|
|
} else {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
String::from("[error] selecting server: no prompt responses found!"),
|
|
)));
|
|
}
|
|
let _ = self.main_tx.send(ToolMessage::EndPrompt);
|
|
}
|
|
"test_server" => {
|
|
if let Some(server) = self.servers.get(self.selected_server) {
|
|
if let Ok(mut server) = server.lock() {
|
|
server.message_que.push("TEST|NONE".to_string());
|
|
} else {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
"[error] getting server lock!".to_string(),
|
|
)));
|
|
}
|
|
} else {
|
|
self.output
|
|
.push(String::from("[error] getting selected server!"));
|
|
}
|
|
}
|
|
"server_sleep" => {
|
|
if let Some(args) = command_args {
|
|
if let Ok(seconds) = args.trim().parse::<u64>() {
|
|
if let Some(server) = self.servers.get(self.selected_server) {
|
|
if let Ok(mut server) = server.lock() {
|
|
server.timer = Duration::from_secs(seconds);
|
|
self.output
|
|
.push(format!("Server timer set to {} seconds", seconds));
|
|
} else {
|
|
self.output
|
|
.push(String::from("[error] locking server object!"));
|
|
}
|
|
} else {
|
|
self.output
|
|
.push(String::from("[error] not a valid server selected!"));
|
|
}
|
|
} else {
|
|
self.output
|
|
.push(String::from("invalid second count provided!"));
|
|
}
|
|
} else {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
"[error] no seconds provided! please use server_sleep number_of_seconds"
|
|
.to_string(),
|
|
)));
|
|
}
|
|
}
|
|
"save_all" => {
|
|
self.save_all()?;
|
|
}
|
|
"save_projects" => {
|
|
self.save_projects()?;
|
|
}
|
|
"list_clients" => {
|
|
if let Some(server) = self.servers.get(self.selected_server) {
|
|
if let Ok(mut server) = server.lock() {
|
|
server.message_que.push("LIST_CLIENTS|NONE".to_string());
|
|
} else {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
"[error] getting server lock!".to_string(),
|
|
)));
|
|
}
|
|
} else {
|
|
self.output
|
|
.push(String::from("[error] getting selected server!"));
|
|
}
|
|
}
|
|
"select_client" => {
|
|
if let Some(arg) = command_args {
|
|
if let Some(server) = self.servers.get(self.selected_server) {
|
|
if let Ok(mut server) = server.lock() {
|
|
if let Ok(id) = arg.parse::<usize>() {
|
|
server.message_que.push(format!("STOP_CONTROL"));
|
|
server.message_que.push(format!("CONTROL|{}", id.clone()));
|
|
server.selected_client = id;
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
0,
|
|
format!("client {} selected!", id),
|
|
)));
|
|
} else {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("[error] parsing client id from arg: {}", arg),
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("No client identified. Usage: select_client client_id"),
|
|
)));
|
|
}
|
|
}
|
|
"deselect_client" => {
|
|
if let Some(server) = self.servers.get(self.selected_server) {
|
|
if let Ok(mut server) = server.lock() {
|
|
server.message_que.push(format!("STOP_CONTROL"));
|
|
server.selected_client = 0;
|
|
let _ = self
|
|
.main_tx
|
|
.send(ToolMessage::Output((rid, format!("client deselected!"))));
|
|
}
|
|
}
|
|
}
|
|
"add_scope" => {
|
|
if let Some(args) = command_args {
|
|
args.split_whitespace().into_iter().for_each(|host| {
|
|
self.projects[self.selected_project].add_scope(host.to_string());
|
|
});
|
|
}
|
|
}
|
|
"parse_scope" => {
|
|
let mut general_path = self.projects[self.selected_project].notes.clone();
|
|
general_path.push("general.md");
|
|
let general_string = read_to_string(general_path)?;
|
|
if let Some((_, scope_section)) = general_string.split_once("# Scope") {
|
|
if let Some((scope_section, _)) = scope_section.split_once("#") {
|
|
scope_section.lines().into_iter().for_each(|line| {
|
|
if line.contains("|") {
|
|
if !line.contains("Subnet") {
|
|
if !line.contains("--") {
|
|
let host =
|
|
line.split("|").collect::<Vec<&str>>()[1].to_string();
|
|
self.projects[self.selected_project].add_scope(host);
|
|
}
|
|
}
|
|
}
|
|
});
|
|
}
|
|
}
|
|
}
|
|
"show_config" | "sc" => {
|
|
let _ = self
|
|
.main_tx
|
|
.send(ToolMessage::Output((rid, "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,
|
|
"[success]Setting changed!".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((rid, format!("{}", module.0))));
|
|
});
|
|
}
|
|
"server_login" => {
|
|
if let Some(pass) = command_args {
|
|
let pass = pass.trim().to_string();
|
|
if let Some(server) = self.servers.get(self.selected_server) {
|
|
if let Ok(mut lock) = server.lock() {
|
|
lock.login(pass)?;
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
"[success] Logged into Server!".to_string(),
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
"new_config" => {
|
|
self.prompt.num_responses = 8;
|
|
self.prompt.action = Some(ToolMessage::Input((0, "add_config".to_string())));
|
|
self.prompt.execute_command = String::from("add_config");
|
|
self.prompt.prompts.push(PromptQuery {
|
|
id: 0,
|
|
query: "name for the new config?".to_string(),
|
|
});
|
|
self.prompt.prompts.push(PromptQuery {
|
|
id: 1,
|
|
query: "Path to save the config folder?".to_string(),
|
|
});
|
|
self.prompt.prompts.push(PromptQuery {
|
|
id: 2,
|
|
query: "Path for upcoming project files (not notes)?".to_string(),
|
|
});
|
|
self.prompt.prompts.push(PromptQuery {
|
|
id: 3,
|
|
query: "Path for current project files (not notes)?".to_string(),
|
|
});
|
|
self.prompt.prompts.push(PromptQuery {
|
|
id: 4,
|
|
query: "Path for upcoming notes?".to_string(),
|
|
});
|
|
self.prompt.prompts.push(PromptQuery {
|
|
id: 5,
|
|
query: "Path for current notes?".to_string(),
|
|
});
|
|
self.prompt.prompts.push(PromptQuery {
|
|
id: 6,
|
|
query: "Path to tool directory?".to_string(),
|
|
});
|
|
self.prompt.prompts.push(PromptQuery {
|
|
id: 7,
|
|
query: "Template box name?".to_string(),
|
|
});
|
|
let _ = self
|
|
.main_tx
|
|
.send(ToolMessage::Input((0, "prompt".to_string())));
|
|
}
|
|
"add_config" => {
|
|
self.prompt.responses.iter().for_each(|res| {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("{} | {}", res.query.query, res.response),
|
|
)));
|
|
});
|
|
self.new_config()?;
|
|
self.save_config()?;
|
|
}
|
|
"list_configs" => {
|
|
if let Some(config_line) = self.config.get("alternate_configs") {
|
|
config_line.split(",").into_iter().for_each(|entry| {
|
|
if let Some((name, path)) = entry.split_once("|") {
|
|
let _ = self
|
|
.main_tx
|
|
.send(ToolMessage::Output((rid, format!("{}: {}", name, path))));
|
|
}
|
|
});
|
|
}
|
|
}
|
|
"collab_with" => {
|
|
if let Some(args) = command_args {
|
|
if let Ok(id) = args.trim().parse::<usize>() {
|
|
if let Some(server) = self.servers.get(self.selected_server) {
|
|
if let Ok(mut lock) = server.lock() {
|
|
lock.action_que.push(format!("COLLAB|{}", id));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
"collab" => {
|
|
if let Some(args) = command_args {
|
|
if let Some(server) = self.servers.get(self.selected_server) {
|
|
if let Ok(mut lock) = server.lock() {
|
|
lock.action_que.push(format!("CMD|{}", args));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
"parse_port_scan" | "pps" => {
|
|
self.parse_cs_portscan()?;
|
|
}
|
|
"stop_db" => {
|
|
if let Some(project) = self.projects.get_mut(self.selected_project) {
|
|
if let Some(db) = project.db.as_mut() {
|
|
db.stop(self.main_tx.clone(), rid)?;
|
|
}
|
|
}
|
|
}
|
|
"rustwitness" => {
|
|
if let Some(project) = self.projects.get(self.selected_project) {
|
|
if let Some(args) = command_args {
|
|
if let Some((outfile, proxy)) = args.split_once(" ") {
|
|
let mut proxy_string = None;
|
|
if let Some((ip, port)) = proxy.split_once(" ") {
|
|
proxy_string = Some(format!("socks5://{}:{}", ip, port));
|
|
}
|
|
WalkDir::new(&project.files)
|
|
.into_iter()
|
|
.filter(|res| res.is_ok())
|
|
.for_each(|res| {
|
|
if let Ok(entry) = res {
|
|
if entry.file_name().to_string_lossy().contains("urls.txt")
|
|
{
|
|
let input = PathBuf::from(entry.path());
|
|
let mut output = project.files.clone();
|
|
output.push(outfile.trim());
|
|
let proxy_intput = proxy_string.clone();
|
|
let tx = self.main_tx.clone();
|
|
self.workers.spawn(move || {
|
|
rustwitness(input, output, proxy_intput, tx)
|
|
});
|
|
}
|
|
}
|
|
});
|
|
}
|
|
} else {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
0,
|
|
"[error] incorrect command usage!".to_string(),
|
|
)));
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
0,
|
|
"usage: rustwitness output_file_name optional_proxy_ip optional_proxy_port".to_string(),
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
"exit" => {
|
|
self.save_all()?;
|
|
let _ = self.main_tx.send(ToolMessage::AppStateExit);
|
|
let _ = self
|
|
.main_tx
|
|
.send(ToolMessage::Output((rid, String::from("Good bye!"))));
|
|
}
|
|
_ => {
|
|
if self.module_loader.asts.get(command_name).is_none() {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("[error] command not found: {}", command_name),
|
|
)));
|
|
return Ok(());
|
|
}
|
|
let ast = self
|
|
.module_loader
|
|
.asts
|
|
.get(command_name)
|
|
.ok_or_else(|| format!("[error] 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.trim() {
|
|
"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);
|
|
}
|
|
"distrobox" => {
|
|
if let Some(db) = self.projects[self.selected_project].db.clone() {
|
|
scope.push("distrobox", db);
|
|
}
|
|
}
|
|
_ 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);
|
|
}
|
|
_ if arg_requirement.starts_with("bool") => {
|
|
let mut input_val = false;
|
|
if command_args
|
|
.clone()
|
|
.unwrap_or_else(|| self.curent_intput.clone())
|
|
.contains("true")
|
|
{
|
|
input_val = true;
|
|
}
|
|
scope.push("input_bool", input_val);
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
let worker_tx = self.main_tx.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() {
|
|
let out_string = result
|
|
.into_string()
|
|
.unwrap_or_else(|_| "failed to parse string".into());
|
|
if cmd_meta.calls == "clipboard" {
|
|
let mut ctx: ClipboardContext =
|
|
ClipboardProvider::new().unwrap();
|
|
ctx.set_contents(out_string.clone()).unwrap();
|
|
} else if cmd_meta.calls != "none" {
|
|
let _ =
|
|
worker_tx.send(ToolMessage::Input((0, out_string.clone())));
|
|
}
|
|
out_string
|
|
} else {
|
|
format!("{:?}", result)
|
|
};
|
|
if result_str.contains("\n") {
|
|
result_str.lines().into_iter().for_each(|line| {
|
|
let _ = tx.send(ToolMessage::Output((rid, line.to_string())));
|
|
});
|
|
} else {
|
|
let _ = tx.send(ToolMessage::Output((rid, result_str)));
|
|
}
|
|
}
|
|
Err(err) => {
|
|
let _ = tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("[error] in script execution: {}", err),
|
|
)));
|
|
}
|
|
}
|
|
});
|
|
}
|
|
}
|
|
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,
|
|
name: String,
|
|
rid: usize,
|
|
) -> Result<(), Box<dyn Error>> {
|
|
let mut main_config_file = self.config_file.clone();
|
|
main_config_file.pop();
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("{} is the main config folder.", main_config_file.display()),
|
|
)));
|
|
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 _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("{} is the project_files folder", project_files.display()),
|
|
)));
|
|
let project_notes = PathBuf::from(self.config.get("upcoming_notes").unwrap())
|
|
.join(format!("{}/{}", org, name));
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("{} is the project_notes folder", project_notes.display()),
|
|
)));
|
|
let mut template_path = main_config_file.clone();
|
|
let mut project_conf_folder = main_config_file.clone();
|
|
let tools = self.config.get("tools").unwrap();
|
|
template_path.push("note_templates");
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("{} is the note template folder", template_path.display()),
|
|
)));
|
|
project_conf_folder.push(format!("projects/{}-{}", org, name));
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!(
|
|
"{} is the project config folder",
|
|
project_conf_folder.display()
|
|
),
|
|
)));
|
|
let mut options = CopyOptions::new();
|
|
options.overwrite = true;
|
|
if !project_files.exists() {
|
|
match create_dir_all(&project_files) {
|
|
Ok(_) => {}
|
|
Err(e) => {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!(
|
|
"[error] making project files! {e} on {}",
|
|
project_files.display()
|
|
),
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
if !project_notes.exists() {}
|
|
match create_dir_all(&project_notes) {
|
|
Ok(_) => {}
|
|
Err(e) => {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!(
|
|
"[error] making project notes! {e} on {}",
|
|
project_notes.display()
|
|
),
|
|
)));
|
|
}
|
|
}
|
|
if !project_conf_folder.exists() {
|
|
match create_dir_all(&project_conf_folder) {
|
|
Ok(_) => {}
|
|
Err(e) => {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!(
|
|
"[error] making config folder! {e} on {}",
|
|
project_conf_folder.display()
|
|
),
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
let mut project_config_file = project_conf_folder.clone();
|
|
project_config_file.push("project.conf");
|
|
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((
|
|
rid,
|
|
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((
|
|
rid,
|
|
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_config_file,
|
|
current: false,
|
|
db: Some(db),
|
|
scope: Vec::new(),
|
|
};
|
|
match new_project.save_config() {
|
|
Ok(_) => {}
|
|
Err(e) => {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
0,
|
|
format!("[error] saving project! {e}"),
|
|
)));
|
|
}
|
|
}
|
|
new_project.save_config()?;
|
|
self.projects.push(new_project);
|
|
return Ok(());
|
|
}
|
|
|
|
pub fn promote_project(&mut self, rid: usize) -> Result<String, Box<dyn Error>> {
|
|
self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
String::from("promoting project..."),
|
|
)))?;
|
|
let return_string = String::from("[success] Project promoted!");
|
|
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)?;
|
|
self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
String::from("[success] folders copied!"),
|
|
)))?;
|
|
remove_dir_all(&new_project.files)?;
|
|
remove_dir_all(&new_project.notes)?;
|
|
new_project.files(new_files_path.clone());
|
|
new_project.notes(new_notes_path.clone());
|
|
new_project.current = true;
|
|
self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
String::from("upcoming folders cleaned up!"),
|
|
)))?;
|
|
self.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
String::from("creating distrobox..."),
|
|
)))?;
|
|
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.main_tx.send(ToolMessage::Output((
|
|
rid,
|
|
String::from("distrobox sent to rebuild function..."),
|
|
)))?;
|
|
}
|
|
new_project.save_config()?;
|
|
self.projects[self.selected_project] = new_project;
|
|
return Ok(return_string);
|
|
}
|
|
|
|
pub fn save_projects(&mut self) -> Result<(), Box<dyn Error>> {
|
|
for mut project in self.projects.clone() {
|
|
project.save_config()?;
|
|
}
|
|
return Ok(());
|
|
}
|
|
|
|
pub fn save_config(&mut self) -> Result<(), Box<dyn Error>> {
|
|
let mut config_vec = Vec::new();
|
|
for key in self.config.keys() {
|
|
if let Some(value) = self.config.get(key) {
|
|
config_vec.push(format!("{}: {}", key, value));
|
|
}
|
|
}
|
|
if config_vec.len() > 0 {
|
|
fs::write(self.config_file.clone(), config_vec.join("\n"))?;
|
|
} else {
|
|
return Err("No config loaded!".into());
|
|
}
|
|
return Ok(());
|
|
}
|
|
|
|
pub fn save_all(&mut self) -> Result<(), Box<dyn Error>> {
|
|
self.save_config()?;
|
|
self.save_projects()?;
|
|
return Ok(());
|
|
}
|
|
|
|
pub fn new_server(&mut self) -> Result<(), Box<dyn Error>> {
|
|
let mut ip = String::new();
|
|
let mut port = String::new();
|
|
self.prompt.responses.iter().for_each(|res| {
|
|
if res.query.id == 0 {
|
|
ip = res.response.clone();
|
|
} else if res.query.id == 1 {
|
|
port = res.response.clone();
|
|
}
|
|
});
|
|
let new_server = Server {
|
|
address: format!("{}:{}", ip, port),
|
|
connected: false,
|
|
timer: Duration::from_secs(60),
|
|
last_check: time::Instant::now(),
|
|
message_que: Vec::new(),
|
|
action_que: Vec::new(),
|
|
client_id: 0,
|
|
selected_client: 0,
|
|
config: None,
|
|
logged_in: false,
|
|
password: String::new(),
|
|
name: String::new(),
|
|
cert_text: String::new(),
|
|
};
|
|
self.servers.push(Arc::new(Mutex::new(new_server)));
|
|
self.prompt.reset();
|
|
Ok(())
|
|
}
|
|
|
|
pub fn new_config(&mut self) -> Result<(), Box<dyn Error>> {
|
|
let mut name = String::new();
|
|
let mut config_path = PathBuf::new();
|
|
let mut upcoming_files_path = PathBuf::new();
|
|
let mut current_files_path = PathBuf::new();
|
|
let mut upcoming_notes_path = PathBuf::new();
|
|
let mut current_notes_path = PathBuf::new();
|
|
let mut tools_folder = PathBuf::new();
|
|
let mut template_box = String::new();
|
|
let term_cmd = self.config.get("term_cmd").unwrap().clone();
|
|
self.prompt
|
|
.responses
|
|
.iter()
|
|
.for_each(|res| match res.query.id {
|
|
0 => name = res.response.clone(),
|
|
1 => config_path = PathBuf::from(res.response.clone()),
|
|
2 => upcoming_files_path = PathBuf::from(res.response.clone()),
|
|
3 => current_files_path = PathBuf::from(res.response.clone()),
|
|
4 => upcoming_notes_path = PathBuf::from(res.response.clone()),
|
|
5 => current_notes_path = PathBuf::from(res.response.clone()),
|
|
6 => tools_folder = PathBuf::from(res.response.clone()),
|
|
7 => template_box = res.response.clone(),
|
|
_ => {}
|
|
});
|
|
let mut existing = false;
|
|
if let Some(line) = self.config.get("alternate_configs") {
|
|
line.split(",").into_iter().for_each(|line| {
|
|
if let Some((ename, _)) = line.split_once("|") {
|
|
if ename.trim().to_string() == name {
|
|
existing = true;
|
|
}
|
|
}
|
|
});
|
|
}
|
|
if existing {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
0,
|
|
format!("{} already exists!", name),
|
|
)));
|
|
return Ok(());
|
|
}
|
|
let root_config_folder_path = config_path.clone();
|
|
let mut projects_path = config_path.clone();
|
|
projects_path.push("projects");
|
|
let mut module_path = config_path.clone();
|
|
module_path.push("modules");
|
|
let mut note_templates_path = config_path.clone();
|
|
note_templates_path.push("note_templates");
|
|
create_dir_all(&module_path)?;
|
|
config_path.push("client.conf");
|
|
let client_conf_path = config_path.clone();
|
|
projects_path.push("default");
|
|
match create_dir_all(&projects_path) {
|
|
Ok(_) => {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
0,
|
|
format!(
|
|
"[success] create projects path at {}",
|
|
projects_path.display()
|
|
),
|
|
)));
|
|
}
|
|
Err(e) => {
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
0,
|
|
format!(
|
|
"[error] couldn't create projects path at {}: {e}",
|
|
projects_path.display()
|
|
),
|
|
)));
|
|
return Ok(());
|
|
}
|
|
}
|
|
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();
|
|
let mut server_path = config_path.clone();
|
|
server_path.push("server");
|
|
create_dir_all(&server_path)?;
|
|
server_path.push("server.conf");
|
|
let mut server_config_file = File::create(&server_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())?;
|
|
client_config_file.write(format!("template_box: {}\n", template_box).as_bytes())?;
|
|
client_config_file.write(format!("term_cmd: {}\n", term_cmd).as_bytes())?;
|
|
client_config_file.write(format!("tools: {}\n", tools_folder.display()).as_bytes())?;
|
|
create_dir_all("./temp")?;
|
|
env::set_current_dir("./temp")?;
|
|
let git_clone_output = Command::new("git")
|
|
.arg("clone")
|
|
.arg("https://git.pyro.monster/pyro/tetanus.git")
|
|
.output()?;
|
|
let git_clone_status = git_clone_output.status.success();
|
|
if git_clone_status {
|
|
let mut options = CopyOptions::new();
|
|
options.overwrite = true;
|
|
options.copy_inside = true;
|
|
copy(
|
|
"./tetanus/note_templates",
|
|
root_config_folder_path.clone(),
|
|
&options,
|
|
)?;
|
|
copy(
|
|
"./tetanus/default-modules",
|
|
module_path.join("default"),
|
|
&options,
|
|
)?;
|
|
copy(
|
|
"./tetanus/victim_templates",
|
|
root_config_folder_path.join("victim"),
|
|
&options,
|
|
)?;
|
|
env::set_current_dir("../")?;
|
|
remove_dir_all("./temp")?;
|
|
let _ = self.main_tx.send(ToolMessage::Output((
|
|
0,
|
|
"New config structure created!".to_string(),
|
|
)));
|
|
}
|
|
if let Some(original_value) = self.config.get("alternate_configs") {
|
|
let new_value = format!(
|
|
"{}|{},{}",
|
|
name,
|
|
client_conf_path.display().to_string(),
|
|
original_value,
|
|
);
|
|
self.config
|
|
.insert("alternate_configs".to_string(), new_value);
|
|
} else {
|
|
self.config.insert(
|
|
"alternate_configs".to_string(),
|
|
format!("{}|{},", name, client_conf_path.display().to_string()),
|
|
);
|
|
}
|
|
self.prompt.responses.clear();
|
|
self.prompt.prompts.clear();
|
|
self.prompt.action = None;
|
|
self.prompt.execute_command.clear();
|
|
Ok(())
|
|
}
|
|
|
|
pub fn parse_cs_portscan(&mut self) -> Result<(), Box<dyn Error>> {
|
|
if let Some(project) = self.projects.get_mut(self.selected_project) {
|
|
let walk = WalkDir::new(project.files.clone());
|
|
walk.into_iter().for_each(|res| {
|
|
if let Ok(entry) = res {
|
|
if entry.file_name().to_string_lossy().contains("services.tsv") {
|
|
if let Ok(scan) = read_to_string(entry.path()) {
|
|
let mut url_lines = Vec::new();
|
|
let mut web_host_lines = Vec::new();
|
|
let mut rdp_host_lines = Vec::new();
|
|
let mut ftp_host_lines = Vec::new();
|
|
let mut telnet_host_lines = Vec::new();
|
|
let mut winrm_host_lines = Vec::new();
|
|
let mut mysql_host_lines = Vec::new();
|
|
let mut mssql_host_lines = Vec::new();
|
|
let mut snmp_host_lines = Vec::new();
|
|
let mut kerberos_host_lines = Vec::new();
|
|
let mut ldap_host_lines = Vec::new();
|
|
let mut smb_host_lines = Vec::new();
|
|
let mut unknonw_host_lines = Vec::new();
|
|
let mut ssh_host_lines = Vec::new();
|
|
scan.lines().into_iter().for_each(|line| {
|
|
if line.len() > 0 {
|
|
let data: Vec<&str> = line.split_whitespace().collect();
|
|
match data[1].trim() {
|
|
"80" | "443" | "8080" | "8443" | "8000" | "4433" => {
|
|
let http_line =
|
|
format!("http://{}:{}", data[0], data[1]);
|
|
let https_line =
|
|
format!("https://{}:{}", data[0], data[1]);
|
|
url_lines.push(http_line);
|
|
url_lines.push(https_line);
|
|
web_host_lines.push(data[0].to_string());
|
|
let mut existing = false;
|
|
project
|
|
.hosts
|
|
.iter_mut()
|
|
.filter(|h| h.ip == data[0].to_string())
|
|
.for_each(|h| {
|
|
h.add_port(data[1]);
|
|
existing = true;
|
|
});
|
|
if !existing {
|
|
let mut new_host = Host::new();
|
|
new_host.ip = data[0].trim().to_string();
|
|
new_host.add_port(data[1]);
|
|
project.add_host(new_host);
|
|
}
|
|
}
|
|
"3389" => {
|
|
rdp_host_lines.push(data[0].to_string());
|
|
let mut existing = false;
|
|
project
|
|
.hosts
|
|
.iter_mut()
|
|
.filter(|h| h.ip == data[0].to_string())
|
|
.for_each(|h| {
|
|
h.add_port(data[1]);
|
|
existing = true;
|
|
});
|
|
if !existing {
|
|
let mut new_host = Host::new();
|
|
new_host.ip = data[0].trim().to_string();
|
|
new_host.add_port(data[1]);
|
|
project.add_host(new_host);
|
|
}
|
|
}
|
|
"21" => {
|
|
ftp_host_lines.push(data[0].to_string());
|
|
let mut existing = false;
|
|
project
|
|
.hosts
|
|
.iter_mut()
|
|
.filter(|h| h.ip == data[0].to_string())
|
|
.for_each(|h| {
|
|
h.add_port(data[1]);
|
|
existing = true;
|
|
});
|
|
if !existing {
|
|
let mut new_host = Host::new();
|
|
new_host.ip = data[0].trim().to_string();
|
|
new_host.add_port(data[1]);
|
|
project.add_host(new_host);
|
|
}
|
|
}
|
|
"23" => {
|
|
telnet_host_lines.push(data[0].to_string());
|
|
let mut existing = false;
|
|
project
|
|
.hosts
|
|
.iter_mut()
|
|
.filter(|h| h.ip == data[0].to_string())
|
|
.for_each(|h| {
|
|
h.add_port(data[1]);
|
|
existing = true;
|
|
});
|
|
if !existing {
|
|
let mut new_host = Host::new();
|
|
new_host.ip = data[0].trim().to_string();
|
|
new_host.add_port(data[1]);
|
|
project.add_host(new_host);
|
|
}
|
|
}
|
|
"22" => {
|
|
ssh_host_lines.push(data[0].to_string());
|
|
let mut existing = false;
|
|
project
|
|
.hosts
|
|
.iter_mut()
|
|
.filter(|h| h.ip == data[0].to_string())
|
|
.for_each(|h| {
|
|
h.add_port(data[1]);
|
|
existing = true;
|
|
});
|
|
if !existing {
|
|
let mut new_host = Host::new();
|
|
new_host.ip = data[0].trim().to_string();
|
|
new_host.add_port(data[1]);
|
|
project.add_host(new_host);
|
|
}
|
|
}
|
|
"5985" | "5986" => {
|
|
winrm_host_lines.push(data[0].to_string());
|
|
let mut existing = false;
|
|
project
|
|
.hosts
|
|
.iter_mut()
|
|
.filter(|h| h.ip == data[0].to_string())
|
|
.for_each(|h| {
|
|
h.add_port(data[1]);
|
|
existing = true;
|
|
});
|
|
if !existing {
|
|
let mut new_host = Host::new();
|
|
new_host.ip = data[0].trim().to_string();
|
|
new_host.add_port(data[1]);
|
|
project.add_host(new_host);
|
|
}
|
|
}
|
|
"3306" => {
|
|
mysql_host_lines.push(data[0].to_string());
|
|
let mut existing = false;
|
|
project
|
|
.hosts
|
|
.iter_mut()
|
|
.filter(|h| h.ip == data[0].to_string())
|
|
.for_each(|h| {
|
|
h.add_port(data[1]);
|
|
existing = true;
|
|
});
|
|
if !existing {
|
|
let mut new_host = Host::new();
|
|
new_host.ip = data[0].trim().to_string();
|
|
new_host.add_port(data[1]);
|
|
project.add_host(new_host);
|
|
}
|
|
}
|
|
"1433" => {
|
|
mssql_host_lines.push(data[0].to_string());
|
|
let mut existing = false;
|
|
project
|
|
.hosts
|
|
.iter_mut()
|
|
.filter(|h| h.ip == data[0].to_string())
|
|
.for_each(|h| {
|
|
h.add_port(data[1]);
|
|
existing = true;
|
|
});
|
|
if !existing {
|
|
let mut new_host = Host::new();
|
|
new_host.ip = data[0].trim().to_string();
|
|
new_host.add_port(data[1]);
|
|
project.add_host(new_host);
|
|
}
|
|
}
|
|
"161" => {
|
|
snmp_host_lines.push(data[0].to_string());
|
|
let mut existing = false;
|
|
project
|
|
.hosts
|
|
.iter_mut()
|
|
.filter(|h| h.ip == data[0].to_string())
|
|
.for_each(|h| {
|
|
h.add_port(data[1]);
|
|
existing = true;
|
|
});
|
|
if !existing {
|
|
let mut new_host = Host::new();
|
|
new_host.ip = data[0].trim().to_string();
|
|
new_host.add_port(data[1]);
|
|
project.add_host(new_host);
|
|
}
|
|
}
|
|
"88" => {
|
|
kerberos_host_lines.push(data[0].to_string());
|
|
let mut existing = false;
|
|
project
|
|
.hosts
|
|
.iter_mut()
|
|
.filter(|h| h.ip == data[0].to_string())
|
|
.for_each(|h| {
|
|
h.add_port(data[1]);
|
|
existing = true;
|
|
});
|
|
if !existing {
|
|
let mut new_host = Host::new();
|
|
new_host.ip = data[0].trim().to_string();
|
|
new_host.add_port(data[1]);
|
|
project.add_host(new_host);
|
|
}
|
|
}
|
|
"389" | "636" => {
|
|
ldap_host_lines.push(data[0].to_string());
|
|
let mut existing = false;
|
|
project
|
|
.hosts
|
|
.iter_mut()
|
|
.filter(|h| h.ip == data[0].to_string())
|
|
.for_each(|h| {
|
|
h.add_port(data[1]);
|
|
existing = true;
|
|
});
|
|
if !existing {
|
|
let mut new_host = Host::new();
|
|
new_host.ip = data[0].trim().to_string();
|
|
new_host.add_port(data[1]);
|
|
project.add_host(new_host);
|
|
}
|
|
}
|
|
"445" | "139" => {
|
|
smb_host_lines.push(data[0].to_string());
|
|
let mut existing = false;
|
|
project
|
|
.hosts
|
|
.iter_mut()
|
|
.filter(|h| h.ip == data[0].to_string())
|
|
.for_each(|h| {
|
|
h.add_port(data[1]);
|
|
existing = true;
|
|
});
|
|
if !existing {
|
|
let mut new_host = Host::new();
|
|
new_host.ip = data[0].trim().to_string();
|
|
new_host.add_port(data[1]);
|
|
project.add_host(new_host);
|
|
}
|
|
}
|
|
_ => {
|
|
unknonw_host_lines.push(data[0].to_string());
|
|
let mut existing = false;
|
|
project
|
|
.hosts
|
|
.iter_mut()
|
|
.filter(|h| h.ip == data[0].to_string())
|
|
.for_each(|h| {
|
|
h.add_port(data[1]);
|
|
existing = true;
|
|
});
|
|
if !existing {
|
|
let mut new_host = Host::new();
|
|
new_host.ip = data[0].trim().to_string();
|
|
new_host.add_port(data[1]);
|
|
project.add_host(new_host);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
});
|
|
if !url_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("urls.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", url_lines.join("\n"));
|
|
}
|
|
}
|
|
if !web_host_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("web_hosts.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", web_host_lines.join("\n"));
|
|
}
|
|
}
|
|
if !rdp_host_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("rdp_hosts.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", rdp_host_lines.join("\n"));
|
|
}
|
|
}
|
|
if !ftp_host_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("ftp_hosts.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", ftp_host_lines.join("\n"));
|
|
}
|
|
}
|
|
if !telnet_host_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("telnet_hosts.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", telnet_host_lines.join("\n"));
|
|
}
|
|
}
|
|
if !winrm_host_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("winrm_hosts.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", winrm_host_lines.join("\n"));
|
|
}
|
|
}
|
|
if !mysql_host_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("mysql_hosts.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", mysql_host_lines.join("\n"));
|
|
}
|
|
}
|
|
if !mssql_host_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("mssql_hosts.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", mssql_host_lines.join("\n"));
|
|
}
|
|
}
|
|
if !snmp_host_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("snmp_hosts.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", snmp_host_lines.join("\n"));
|
|
}
|
|
}
|
|
if !kerberos_host_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("kerberos_hosts.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", kerberos_host_lines.join("\n"));
|
|
}
|
|
}
|
|
if !ldap_host_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("ldap_hosts.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", ldap_host_lines.join("\n"));
|
|
}
|
|
}
|
|
if !smb_host_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("smb_hosts.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", smb_host_lines.join("\n"));
|
|
}
|
|
}
|
|
if !unknonw_host_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("unknown_hosts.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", unknonw_host_lines.join("\n"));
|
|
}
|
|
}
|
|
if !ssh_host_lines.is_empty() {
|
|
let mut file_path = project.files.clone();
|
|
file_path.push("ssh_hosts.txt");
|
|
if let Ok(mut file) = File::create(file_path) {
|
|
let _ = write!(file, "{}", ssh_host_lines.join("\n"));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
});
|
|
project.save_config()?;
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct Server {
|
|
pub address: String,
|
|
pub connected: bool,
|
|
pub timer: Duration,
|
|
pub config: Option<Arc<ClientConfig>>,
|
|
pub last_check: time::Instant,
|
|
pub message_que: Vec<String>,
|
|
pub action_que: Vec<String>,
|
|
pub client_id: usize,
|
|
pub selected_client: usize,
|
|
pub logged_in: bool,
|
|
pub password: String,
|
|
pub name: String,
|
|
pub cert_text: String,
|
|
}
|
|
|
|
impl Server {
|
|
pub fn donwload_certificate(&self) -> Result<String, Box<dyn Error>> {
|
|
let mut stream = TcpStream::connect(self.address.clone())?;
|
|
stream.write_all("CERT_REQ".as_bytes())?;
|
|
let mut buf = [0u8; 8192];
|
|
let bytes_read = stream.read(&mut buf)?;
|
|
let response = String::from_utf8_lossy(&buf[..bytes_read]);
|
|
if let Some((_, cert)) = response.split_once("|") {
|
|
return Ok(cert.to_string());
|
|
}
|
|
Err("Server did not send certificate".into())
|
|
}
|
|
|
|
pub fn connect(&mut self) -> Result<usize, Box<dyn Error>> {
|
|
if let Ok(cert) = self.donwload_certificate() {
|
|
self.config = Some(self.build_tls_config(cert.as_str())?);
|
|
self.cert_text = cert;
|
|
}
|
|
let mut stream = self.tls_connect()?;
|
|
stream.write("HELLO|attack".as_bytes())?;
|
|
let mut buf = [0; 8192];
|
|
let bytes_read = stream.read(&mut buf)?;
|
|
let response = String::from_utf8_lossy(&buf[..bytes_read]);
|
|
if let Some((_, id)) = response.split_once("|") {
|
|
self.client_id = id.trim().parse()?;
|
|
}
|
|
self.connected = true;
|
|
self.last_check = Instant::now();
|
|
return Ok(self.client_id);
|
|
}
|
|
|
|
pub fn login(&mut self, pass: String) -> Result<(), Box<dyn Error>> {
|
|
let mut stream = self.tls_connect()?;
|
|
stream.write(format!("NONE**{}|||LOGIN|{}||", self.client_id, pass).as_bytes())?;
|
|
let mut buf = [0; 8192];
|
|
let bytes_read = stream.read(&mut buf)?;
|
|
let response = String::from_utf8_lossy(&buf[..bytes_read]);
|
|
if response.contains("successful") {
|
|
self.password = pass;
|
|
self.last_check = Instant::now();
|
|
self.logged_in = true;
|
|
self.message_que.push(format!("SET_NAME|{}", self.name));
|
|
return Ok(());
|
|
} else {
|
|
return Err("[error] logging in!".into());
|
|
}
|
|
}
|
|
|
|
pub fn checkin(&mut self) {
|
|
if self.logged_in {
|
|
if let Ok(mut stream) = self.tls_connect() {
|
|
self.last_check = Instant::now();
|
|
let payload = format!(
|
|
"{}**{}|||{}\n",
|
|
self.password.clone(),
|
|
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}"));
|
|
}
|
|
self.message_que.clear();
|
|
let mut buffer = [0; 4096];
|
|
if let Ok(bytes_read) = stream.read(&mut buffer) {
|
|
if bytes_read > 0 {
|
|
let response = String::from_utf8_lossy(&buffer[..bytes_read]);
|
|
response.split("||").into_iter().for_each(|action| {
|
|
self.action_que.push(action.trim().to_string());
|
|
if action.contains("Not authenticated") {
|
|
self.logged_in = false;
|
|
}
|
|
});
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn build_tls_config(&self, pem: &str) -> Result<Arc<ClientConfig>, Box<dyn Error>> {
|
|
let mut roots = RootCertStore::empty();
|
|
let mut reader = BufReader::new(pem.as_bytes());
|
|
let certs = rustls_pemfile::certs(&mut reader).collect::<Result<Vec<_>, _>>()?;
|
|
for cert in certs {
|
|
roots.add(cert)?;
|
|
}
|
|
let config = ClientConfig::builder()
|
|
.with_root_certificates(roots)
|
|
.with_no_client_auth();
|
|
Ok(Arc::new(config))
|
|
}
|
|
|
|
fn tls_connect(&self) -> Result<StreamOwned<ClientConnection, TcpStream>, Box<dyn Error>> {
|
|
let config = self.config.clone().ok_or("TLS config not loaded")?;
|
|
let tcp = TcpStream::connect(&self.address)?;
|
|
let (host, _) = self.address.split_once(':').ok_or("Invalid address")?;
|
|
let server_name = ServerName::try_from(host.to_string())?;
|
|
let connection = ClientConnection::new(config, server_name)?;
|
|
let tls_stream = StreamOwned::new(connection, tcp);
|
|
Ok(tls_stream)
|
|
}
|
|
|
|
fn disconnect(&mut self) -> Result<(), Box<dyn Error>> {
|
|
let mut stream = self.tls_connect()?;
|
|
stream.write(format!("{}**{}|||DISCONNECT||", self.password, self.client_id).as_bytes())?;
|
|
self.connected = false;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[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>,
|
|
pub scope: Vec<String>,
|
|
}
|
|
|
|
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,
|
|
scope: Vec::new(),
|
|
}
|
|
}
|
|
|
|
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, mut host: Host) {
|
|
let mut new_id = 0;
|
|
self.hosts.iter().for_each(|h| {
|
|
if h.id > new_id {
|
|
new_id = h.id.clone();
|
|
}
|
|
});
|
|
new_id += 1;
|
|
if !self.hosts.iter().any(|h| h.ip == host.ip) {
|
|
host.id = new_id;
|
|
self.hosts.push(host);
|
|
}
|
|
}
|
|
|
|
pub fn config_folder(&mut self, path: PathBuf) {
|
|
self.config_folder = path;
|
|
}
|
|
|
|
pub fn add_scope(&mut self, host: String) {
|
|
if let Ok(ip) = host.parse::<IpAddr>() {
|
|
self.scope.push(ip.to_string())
|
|
} else if let Ok(net) = host.parse::<IpNet>() {
|
|
net.hosts().into_iter().for_each(|ip| {
|
|
self.scope.push(ip.to_string());
|
|
});
|
|
} else {
|
|
self.scope.push(host);
|
|
}
|
|
}
|
|
|
|
pub fn load_config(&mut self, display: bool) -> 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;
|
|
}
|
|
}
|
|
"scope" => {
|
|
let ips: Vec<&str> = parts[1].split(",").collect();
|
|
ips.iter().for_each(|ip| {
|
|
self.add_scope(ip.trim().to_string());
|
|
});
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
if display {
|
|
println!("{} | {} config loaded!", self.org_name, self.name);
|
|
println!("loading hosts...");
|
|
}
|
|
let mut conf_folder = self.config_folder.clone();
|
|
conf_folder.pop();
|
|
let mut hosts_folder = conf_folder.clone();
|
|
let mut users_folder = conf_folder.clone();
|
|
hosts_folder.push("hosts");
|
|
users_folder.push("users");
|
|
let mut users = HashMap::new();
|
|
if users_folder.exists() {
|
|
for entry in read_dir(users_folder)? {
|
|
let entry = entry?;
|
|
let user_config_string = read_to_string(entry.path())?;
|
|
let mut new_user = User::new();
|
|
for line in user_config_string.lines() {
|
|
if line.contains(": ") {
|
|
let (setting, data) = line.split_once(": ").unwrap();
|
|
match setting.trim() {
|
|
"username" => new_user.name = data.trim().to_string(),
|
|
"password" => new_user.password = Some(data.trim().to_string()),
|
|
"hash" => new_user.hash = Some(data.trim().to_string()),
|
|
"compromised" => new_user.compromised = data.trim().parse().unwrap(),
|
|
"ticket" => new_user.ticket = Some(PathBuf::from(data.trim())),
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
users.insert(new_user.name.clone(), new_user);
|
|
}
|
|
}
|
|
if hosts_folder.exists() {
|
|
read_dir(hosts_folder)?.into_iter().for_each(|res| {
|
|
if let Ok(entry) = res {
|
|
let host_path = entry.path();
|
|
let mut new_host = Host::new();
|
|
println!("loading {}...", host_path.display());
|
|
if let Ok(host_conf_text) = read_to_string(host_path) {
|
|
host_conf_text.lines().into_iter().for_each(|line| {
|
|
if line.contains(": ") {
|
|
let (setting, data) = line.split_once(": ").unwrap();
|
|
match setting.trim() {
|
|
"ip" => new_host.ip = data.trim().to_string(),
|
|
"hostname" => new_host.hostname = data.trim().to_string(),
|
|
"control_port" => {
|
|
new_host.control_port = data.trim().parse().unwrap();
|
|
}
|
|
"pwned" => {
|
|
new_host.pwned = data.trim().parse().unwrap();
|
|
}
|
|
"id" => new_host.id = data.trim().parse().unwrap(),
|
|
"findings" => {
|
|
new_host.findings = data
|
|
.trim()
|
|
.split(",")
|
|
.map(|finding| finding.to_string())
|
|
.collect::<Vec<String>>()
|
|
}
|
|
"open_ports" => {
|
|
new_host.open_ports = data
|
|
.trim()
|
|
.split(",")
|
|
.into_iter()
|
|
.map(|port| port.parse().unwrap_or_default())
|
|
.collect::<Vec<usize>>()
|
|
}
|
|
"users" => {
|
|
let names: Vec<&str> = data.trim().split(",").collect();
|
|
names.into_iter().for_each(|name| {
|
|
if let Some(user) = users.get(name) {
|
|
new_host.users.push(user.clone());
|
|
}
|
|
});
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
});
|
|
self.add_host(new_host);
|
|
}
|
|
}
|
|
});
|
|
}
|
|
if display {
|
|
println!("{} | {} loaded!", self.org_name, self.name);
|
|
}
|
|
return Ok(());
|
|
}
|
|
|
|
pub fn save_config(&mut self) -> Result<(), Box<dyn Error>> {
|
|
let mut conf_folder = self.config_folder.clone();
|
|
conf_folder.pop();
|
|
if !conf_folder.exists() {
|
|
match create_dir_all(&conf_folder) {
|
|
Ok(_) => {}
|
|
Err(e) => {
|
|
return Err(
|
|
format!("save_config: couldn't create project conf folder! {e}").into(),
|
|
);
|
|
}
|
|
}
|
|
}
|
|
let mut hosts_folder = conf_folder.clone();
|
|
hosts_folder.push("hosts");
|
|
let mut users_folder = conf_folder.clone();
|
|
users_folder.push("users");
|
|
if !hosts_folder.exists() {
|
|
create_dir_all(&hosts_folder)?;
|
|
}
|
|
if !users_folder.exists() {
|
|
create_dir_all(&users_folder)?;
|
|
}
|
|
let mut file = File::create(&self.config_folder)?;
|
|
let mut out_string = format!(
|
|
"org_name: {}\nname: {}\nnotes: {}\nfiles: {}\nscope: {}\n",
|
|
self.org_name,
|
|
self.name,
|
|
self.notes.display(),
|
|
self.files.display(),
|
|
self.scope.join(","),
|
|
);
|
|
if self.current {
|
|
out_string.push_str("stage: current");
|
|
} else {
|
|
out_string.push_str("stage: upcoming");
|
|
}
|
|
file.write_all(out_string.as_bytes())?;
|
|
self.hosts
|
|
.iter_mut()
|
|
.try_for_each(|h| -> Result<(), Box<dyn Error>> {
|
|
let mut host_path = hosts_folder.clone();
|
|
host_path.push(format!("{}.conf", h.id));
|
|
h.save_config_file(host_path)
|
|
})?;
|
|
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: &str) {
|
|
if let Ok(port) = port.trim().parse() {
|
|
if !self.open_ports.contains(&port) {
|
|
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 fn save_config_file(&mut self, path: PathBuf) -> Result<(), Box<dyn Error>> {
|
|
let mut host_file = File::create(path)?;
|
|
let mut out_string = format!(
|
|
"
|
|
ip: {}
|
|
hostname: {}
|
|
control_port: {}
|
|
pwned: {}
|
|
id: {}",
|
|
self.ip, self.hostname, self.control_port, self.pwned, self.id
|
|
);
|
|
if !self.findings.is_empty() {
|
|
out_string.push_str(&format!("\nfindings: "));
|
|
out_string.push_str(&self.findings.join(","));
|
|
}
|
|
if !self.open_ports.is_empty() {
|
|
out_string.push_str("\nopen_ports: ");
|
|
out_string.push_str(
|
|
&self
|
|
.open_ports
|
|
.iter()
|
|
.map(|port| port.to_string())
|
|
.collect::<Vec<String>>()
|
|
.join(","),
|
|
);
|
|
}
|
|
write!(host_file, "{}", out_string)?;
|
|
host_file.sync_all()?;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct User {
|
|
pub name: String,
|
|
pub password: Option<String>,
|
|
pub hash: Option<String>,
|
|
pub ticket: Option<PathBuf>,
|
|
pub compromised: bool,
|
|
}
|
|
|
|
impl User {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
name: String::new(),
|
|
password: None,
|
|
hash: None,
|
|
ticket: None,
|
|
compromised: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[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 calls: String,
|
|
}
|
|
|
|
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_get_set(
|
|
"notes",
|
|
|p: &mut Project| p.notes.clone(),
|
|
|p: &mut Project, val: PathBuf| p.notes = val,
|
|
)
|
|
.register_get_set(
|
|
"files",
|
|
|p: &mut Project| p.files.clone(),
|
|
|p: &mut Project, val: PathBuf| p.files = 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)
|
|
});
|
|
engine.register_fn("new_host", Host::new);
|
|
engine.register_fn("open_file", |path: &str| {
|
|
std::fs::read_to_string(path).unwrap_or_else(|_| String::new())
|
|
});
|
|
engine.register_fn("add_port", |host: &mut Host, port: i64| {
|
|
host.add_port(format!("{}", port).as_str());
|
|
});
|
|
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("filter_ips", |s: &str| -> rhai::Array {
|
|
let mut results = rhai::Array::new();
|
|
if let Ok(net) = s.trim().parse::<IpNet>() {
|
|
results.push(net.addr().to_string().trim().into());
|
|
} else if let Ok(ip) = s.trim().parse::<IpAddr>() {
|
|
results.push(ip.to_string().trim().into());
|
|
} else if s.contains(".") {
|
|
results.push(s.to_string().trim().into());
|
|
}
|
|
results
|
|
});
|
|
engine.register_fn("write_file", |path: &str, contents: &str| {
|
|
std::fs::write(path, contents).is_ok()
|
|
});
|
|
engine.register_fn("find_file", |path: PathBuf, filename: String| -> String {
|
|
for entry in walkdir::WalkDir::new(path).max_depth(50) {
|
|
if let Ok(entry) = entry {
|
|
if entry
|
|
.file_name()
|
|
.to_string_lossy()
|
|
.contains(filename.as_str())
|
|
{
|
|
return entry.path().display().to_string();
|
|
}
|
|
}
|
|
}
|
|
|
|
String::new()
|
|
});
|
|
engine.register_fn("execute", |cmd: &str| -> String {
|
|
let out_string;
|
|
if let Ok(output) = Command::new("sh").arg("-c").arg(cmd).output() {
|
|
if output.status.success() {
|
|
out_string = format!("[success] {}", String::from_utf8_lossy(&output.stdout));
|
|
} else {
|
|
out_string = format!("[error] {}", String::from_utf8_lossy(&output.stderr));
|
|
}
|
|
} else {
|
|
out_string = "[error] failed to run command".to_string();
|
|
}
|
|
return out_string;
|
|
});
|
|
engine.register_fn("db_run", |p: &mut Project, command: String| -> String {
|
|
let mut out_string = format!("No distrobox configured for selected project.");
|
|
if let Some(db) = p.db.clone() {
|
|
out_string = db.run_command(&command);
|
|
}
|
|
return out_string;
|
|
});
|
|
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();
|
|
let mut calls = String::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();
|
|
}
|
|
"calls" => {
|
|
calls = val.trim().to_string();
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
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,
|
|
calls,
|
|
};
|
|
|
|
Some((command, ast))
|
|
}
|
|
}
|
|
|
|
#[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>, rid: usize) -> Result<(), Box<dyn Error>> {
|
|
if let Err(e) = self.stop_template(tx.clone(), rid) {
|
|
tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("[error] stopping template box! {e}"),
|
|
)))?;
|
|
return Ok(());
|
|
}
|
|
let mut create_command = Command::new("distrobox");
|
|
create_command
|
|
.arg("create")
|
|
.arg("--clone")
|
|
.arg(self.template.clone())
|
|
.arg("--name")
|
|
.arg(self.name.clone());
|
|
let mut added_volumes = Vec::new();
|
|
for volume in &self.volumes {
|
|
let folder_name = volume.split("/").last().unwrap();
|
|
if !added_volumes.contains(&String::from(folder_name)) {
|
|
create_command
|
|
.arg("--volume")
|
|
.arg(format!("{}:/{}:rw", volume, folder_name));
|
|
added_volumes.push(folder_name.to_string());
|
|
}
|
|
}
|
|
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((rid, line.clone())));
|
|
println!("{}", 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((rid, line.clone())));
|
|
println!("{}", line);
|
|
}
|
|
});
|
|
}
|
|
let status = child.wait()?;
|
|
if !status.success() {
|
|
let _ = tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("[error] creating distrobox!"),
|
|
)));
|
|
}
|
|
|
|
self.created = true;
|
|
return Ok(());
|
|
}
|
|
|
|
pub fn stop_template(
|
|
&mut self,
|
|
tx: Sender<ToolMessage>,
|
|
rid: usize,
|
|
) -> Result<(), Box<dyn Error>> {
|
|
let mut stop_command = Command::new("distrobox");
|
|
stop_command
|
|
.arg("stop")
|
|
.arg(self.template.clone())
|
|
.arg("--yes");
|
|
stop_command.stdin(Stdio::piped());
|
|
stop_command.stdout(Stdio::piped());
|
|
stop_command.stderr(Stdio::piped());
|
|
let mut child = stop_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((rid, line.clone())));
|
|
println!("{}", 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((rid, line.clone())));
|
|
println!("{}", line);
|
|
}
|
|
});
|
|
}
|
|
let status = child.wait()?;
|
|
if !status.success() {
|
|
let _ = tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("[error] stopping template distrobox!"),
|
|
)));
|
|
}
|
|
return Ok(());
|
|
}
|
|
|
|
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(self.name.clone()).arg("--yes");
|
|
stop_command.stdin(Stdio::piped());
|
|
stop_command.stdout(Stdio::piped());
|
|
stop_command.stderr(Stdio::piped());
|
|
let mut child = stop_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((rid, line.clone())));
|
|
}
|
|
});
|
|
}
|
|
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((rid, line.clone())));
|
|
}
|
|
});
|
|
}
|
|
let status = child.wait()?;
|
|
if !status.success() {
|
|
let _ = tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("[error] stopping distrobox!"),
|
|
)));
|
|
}
|
|
return Ok(());
|
|
}
|
|
|
|
pub fn destroy(&mut self, tx: Sender<ToolMessage>, rid: usize) -> Result<(), Box<dyn Error>> {
|
|
let mut destroycmd = Command::new("distrobox");
|
|
destroycmd.arg("rm").arg(self.name.clone()).arg("-f");
|
|
let res = destroycmd.status()?;
|
|
if !res.success() {
|
|
let _ = tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("[error] destroying distrobox!"),
|
|
)));
|
|
let _ = tx.send(ToolMessage::Output((
|
|
rid,
|
|
String::from("pleas try destroying manually:"),
|
|
)));
|
|
let _ = tx.send(ToolMessage::Output((
|
|
rid,
|
|
format!("distrobox rm {} -f", self.name),
|
|
)));
|
|
}
|
|
self.created = false;
|
|
return Ok(());
|
|
}
|
|
|
|
pub fn launch_terminal(&self, cmd: String) {
|
|
if let Some((terminal_cmd, args)) = cmd.split_once(" ") {
|
|
let mut term_cmd = Command::new(terminal_cmd);
|
|
args.split(" ").into_iter().for_each(|arg| {
|
|
if arg == "|||COMMAND|||" {
|
|
term_cmd
|
|
.arg("distrobox")
|
|
.arg("enter")
|
|
.arg(self.name.clone());
|
|
} else if arg.contains("ENV_NAME=") {
|
|
let (_, env_name) = arg.split_once("=").unwrap();
|
|
unsafe { std::env::set_var(env_name, self.name.clone()) };
|
|
} else {
|
|
term_cmd.arg(arg);
|
|
}
|
|
});
|
|
let _ = term_cmd.spawn();
|
|
}
|
|
}
|
|
|
|
pub fn run_command(&self, cmd: &str) -> String {
|
|
let output = Command::new("distrobox")
|
|
.arg("enter")
|
|
.arg("-e")
|
|
.arg(format!("\"{}\"", cmd))
|
|
.output();
|
|
match output {
|
|
Ok(output) => {
|
|
if output.status.success() {
|
|
return String::from_utf8_lossy(&output.stdout).to_string();
|
|
} else {
|
|
return format!("[error] {}", String::from_utf8_lossy(&output.stdout));
|
|
}
|
|
}
|
|
Err(e) => {
|
|
return format!("[error] {e}");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct Prompt {
|
|
action: Option<ToolMessage>,
|
|
responses: Vec<PromptResponse>,
|
|
execute_command: String,
|
|
num_responses: usize,
|
|
prompts: Vec<PromptQuery>,
|
|
last_prompted: Option<PromptQuery>,
|
|
}
|
|
|
|
impl Prompt {
|
|
pub fn reset(&mut self) {
|
|
self.action = None;
|
|
self.responses.clear();
|
|
self.execute_command.clear();
|
|
self.num_responses = 0;
|
|
self.prompts.clear();
|
|
self.last_prompted = None;
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct PromptResponse {
|
|
query: PromptQuery,
|
|
response: String,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct PromptQuery {
|
|
id: usize,
|
|
query: String,
|
|
}
|