Files
tetanus/src/funcs.rs
T
pyro 9ed69b03a0 added a bit of breakup between the startup output and the prompt to
select a config, also made the switch_config command clear the
current_input so that it's not left over when you come back.
2026-08-19 17:04:39 -05:00

823 lines
41 KiB
Rust

use crate::*;
use crossterm::{
cursor,
event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode, KeyEventKind},
execute,
terminal::{EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode},
};
use ratatui::{
prelude::*,
widgets::{Block, Borders, List, ListItem, ListState, Paragraph},
};
use std::thread::spawn;
use std::thread::{JoinHandle, sleep};
use std::time::Duration;
use std::{error::Error, time::Instant};
use std::{io::Write, usize};
use sysinfo::System;
pub fn get_user_input(prompt: &str) -> Result<String, Box<dyn Error>> {
println!("{}", prompt);
let mut response = String::new();
std::io::stdin().read_line(&mut response)?;
return Ok(response.trim().to_string());
}
pub fn startup(
mut state: AppState,
nested: bool,
handles: Arc<Mutex<Vec<JoinHandle<()>>>>,
) -> Result<(), Box<dyn std::error::Error>> {
let mut app = AppState::new();
let (tx, rx) = channel();
app.main_tx = tx.clone();
if let Some(line) = state.config.get("alternate_configs") {
let mut selections = HashMap::new();
println!("\n");
line.split(",").into_iter().for_each(|line| {
if let Some((name, path)) = line.split_once("|") {
selections.insert(name.to_string(), path.to_string());
println!("{} | {} ", name, path);
}
});
selections.insert(
"default".to_string(),
state.config_file.display().to_string(),
);
println!("default | {}", state.config_file.display());
println!("\n");
let response = get_user_input("which config would you like to load?")?;
if let Some(path) = selections.get(&response) {
app.load_config(PathBuf::from(path), true)?;
run_tui(app, rx, handles.clone())?;
} else {
eprintln!("error invalid selection! exiting...");
return Err("invalid config selected.".into());
}
} else {
app.load_config(state.config_file.clone(), true)?;
run_tui(app, rx, handles.clone())?;
}
if nested {
let (tx, rx) = channel();
state.main_tx = tx;
run_tui(state, rx, handles.clone())?;
}
Ok(())
}
pub fn run_tui(
mut state: AppState,
main_rx: Receiver<ToolMessage>,
main_handles: Arc<Mutex<Vec<JoinHandle<()>>>>,
) -> Result<(), Box<dyn std::error::Error>> {
enable_raw_mode()?;
let mut stdout = io::stdout();
execute!(
stdout,
EnterAlternateScreen,
cursor::Hide,
EnableMouseCapture
)?;
let backend = CrosstermBackend::new(&stdout);
let mut terminal = Terminal::new(backend)?;
let (event_tx, event_rx) = channel::<AppEvent>();
let input_tx = event_tx.clone();
let mut project_list_state = ListState::default();
let mut system = System::new_all();
let system_timer = Duration::from_secs(1);
let mut last_print = Instant::now();
let mut switch_config = false;
const BYTES_PER_GB: u64 = 1024 * 1024 * 1024;
if !state.projects.is_empty() {
project_list_state.select(Some(0));
state.selected_project = 0;
}
std::thread::spawn(move || {
loop {
if !state.server_broker_running && !state.app_state_running {
break;
}
if event::poll(Duration::from_millis(100)).unwrap_or(false) {
match event::read() {
Ok(Event::Key(key)) => {
if key.kind == KeyEventKind::Press {
if input_tx.send(AppEvent::Key(key)).is_err() {
break;
}
}
}
Ok(Event::Mouse(mouse)) => {
if input_tx.send(AppEvent::Mouse(mouse)).is_err() {
break;
}
}
_ => {}
}
}
}
});
let worker_tx = event_tx.clone();
std::thread::spawn(move || {
while let Ok(msg) = main_rx.recv() {
if worker_tx.send(AppEvent::Worker(msg)).is_err() {
break;
}
}
});
let mut project_list_state = ListState::default();
if !state.projects.is_empty() {
project_list_state.select(Some(0));
}
let mut history_index = state.history.len();
let mut handles = Vec::new();
loop {
if switch_config {}
if Instant::now().duration_since(last_print) > system_timer {
system.refresh_cpu_usage();
system.refresh_memory();
last_print = Instant::now();
}
let mut text_area_width = 0;
terminal.draw(|f| {
let main_chunks = Layout::default()
.direction(Direction::Vertical)
.constraints([Constraint::Min(0), Constraint::Max(3)])
.split(f.area());
let top_chunks = Layout::default()
.direction(Direction::Horizontal)
.constraints([
Constraint::Percentage(20),
Constraint::Percentage(45),
Constraint::Percentage(35),
])
.split(main_chunks[0]);
text_area_width = top_chunks[1].width.saturating_sub(2);
let projects: Vec<ListItem> = state
.projects
.iter()
.map(|p| ListItem::new(format!(" {} | {}", p.org_name, p.name)))
.collect();
let projects_list = List::new(projects)
.block(
Block::default()
.borders(Borders::ALL)
.title(" Projects (ctrl + arrow keys to select) "),
)
.highlight_style(
Style::default()
.bg(Color::Blue)
.fg(Color::White)
.add_modifier(Modifier::BOLD),
)
.highlight_symbol(">> ");
f.render_stateful_widget(projects_list, top_chunks[0], &mut project_list_state);
let mut info_lines: Vec<Line> = Vec::new();
let project = state.projects[state.selected_project].clone();
info_lines.push(Line::from(format!(
"|CPU: {}% | RAM: {}GB / {}GB |",
system.global_cpu_usage().floor(),
system.used_memory() / BYTES_PER_GB,
system.total_memory() / BYTES_PER_GB,
)));
info_lines.push(Line::from(Span::styled(
"--- PROJECT INFORMATION ---",
Style::default().fg(Color::Green),
)));
info_lines.push(Line::from(format!("ORG: {}", project.org_name)));
info_lines.push(Line::from(format!("NAME: {}", project.name)));
info_lines.push(Line::from(format!("NOTES: {}", project.notes.display())));
info_lines.push(Line::from(format!("Files: {}", project.files.display())));
if let Some(db) = project.db {
info_lines.push(Line::from(format!("DISTROBOX: {}", db.name)));
}
if project.current {
info_lines.push(Line::from("STATUS: CURRENT"));
} else {
info_lines.push(Line::from("STATUS: UPCOMING"));
}
if !project.hosts.is_empty() {
info_lines.push(Line::from(Span::styled(
"--- PROJECT HOST INFORMATION ---",
Style::default().fg(Color::Green),
)));
for host in &project.hosts {
info_lines.push(Line::from(format!("- {}:{}", host.hostname, host.ip)));
info_lines.push(Line::from(format!(" - PWNED: {}", host.pwned)));
info_lines.push(Line::from(format!(
" - CONTROL PORT: {}",
host.control_port
)));
if !host.open_ports.is_empty() {
info_lines.push(Line::from(format!(" - PORTS:")));
for port in &host.open_ports {
info_lines.push(Line::from(format!(" - {}", port)));
}
}
if !host.users.is_empty() {
info_lines.push(Line::from(format!(" - USERS:")));
for user in &host.users {
info_lines.push(Line::from(format!(
" - {} - PWNED: {}",
user.name, user.compromised
)));
}
}
}
}
let spacer = format!("{}", "+".repeat(top_chunks[2].width as usize - 2));
info_lines.push(Line::from(Span::styled(
&spacer,
Style::default().fg(Color::Green),
)));
info_lines.push(Line::from(Span::styled(
"--- Tool Information ---",
Style::default().fg(Color::Green),
)));
info_lines.push(Line::from(format!(
"CONFIG FILE: {}",
state.config_file.display()
)));
for setting in state.config.keys() {
info_lines.push(Line::from(format!(
"{}: {}",
setting.to_uppercase(),
state.config.get(setting).unwrap()
)));
}
let info_area_height = top_chunks[1].height.saturating_sub(2) as usize;
if state.info_scroll as usize > info_lines.len().saturating_sub(info_area_height) {
state.info_scroll = info_lines.len().saturating_sub(info_area_height) as u16;
}
let info_paragraph = Paragraph::new(info_lines)
.block(
Block::default()
.borders(Borders::ALL)
.title("Selected Project Information"),
)
.scroll((state.info_scroll, 0))
.wrap(ratatui::widgets::Wrap { trim: false });
f.render_widget(info_paragraph, top_chunks[2]);
let mut output_lines = Vec::new();
state.output.iter().for_each(|text| {
let line = Line::from(Span::raw(text));
if text.to_lowercase().contains("error") {
let line = line.style(Style::default().fg(Color::Red));
output_lines.push(line);
} else if text.to_lowercase().contains("success") {
let line = line.style(Style::default().fg(Color::Green));
output_lines.push(line);
} else {
output_lines.push(line);
}
});
let text_area_height = top_chunks[1].height.saturating_sub(2) as usize;
let max_scroll = state.output.len().saturating_sub(text_area_height);
if state.output_follow {
state.output_scroll = max_scroll as u16;
}
state.output_scroll = state.output_scroll.min(max_scroll as u16);
if state.output_scroll == u16::MAX {
if state.output.len() > text_area_height {
state.output_scroll = (state.output.len() - text_area_height) as u16;
} else {
state.output_scroll = 0;
}
}
let output_paragraph = Paragraph::new(output_lines)
.block(
Block::default()
.borders(Borders::ALL)
.title(" Script Engine Output "),
)
.scroll((state.output_scroll, 0));
f.render_widget(output_paragraph, top_chunks[1]);
let input_paragraph = Paragraph::new(state.curent_intput.as_str()).block(
Block::default()
.borders(Borders::ALL)
.title(" What is thy bidding, my master? "),
);
f.render_widget(input_paragraph, main_chunks[1]);
})?;
if let Ok(event) = event_rx.recv_timeout(Duration::from_millis(100)) {
match event {
AppEvent::Worker(msg) => match msg {
ToolMessage::EndPrompt => {
state.prompt.reset();
}
ToolMessage::Input(cmd) => {
state.log.push(cmd.1.clone());
if cmd.1.contains(" ") {
let (gcmd, args) = cmd.1.split_once(" ").unwrap();
state.execute_command(gcmd, Some(args.to_string()), cmd.0)?;
} else {
state.execute_command(cmd.1.as_str(), None, cmd.0)?;
}
}
ToolMessage::Output((rid, txt)) => {
if txt.to_lowercase().contains("[error]") {
let (_, output) = txt.split_once("]").unwrap();
for line in
textwrap::wrap(output, text_area_width.saturating_sub(6) as usize)
{
state.output.push(format!("error: {}", line.to_string()));
}
} else if txt.contains("[success]") {
let (_, output) = txt.split_once("]").unwrap();
for line in
textwrap::wrap(output, text_area_width.saturating_sub(9) as usize)
{
state.output.push(format!("success: {}", line.to_string()));
}
} else {
for line in textwrap::wrap(&txt, text_area_width as usize) {
state.output.push(line.to_string());
}
}
state.log.push(txt.clone());
if rid != 0 {
if let Some(server) = state.servers.get(state.selected_server) {
if let Ok(mut server) = server.lock() {
server.message_que.push(format!("OUTPUT|{}", txt));
}
}
}
}
ToolMessage::UpdateProject(index, project) => {
if index < state.projects.len() {
state.projects[index] = project;
}
}
ToolMessage::RebuildDB => {
state.main_tx.send(ToolMessage::Output((
0,
String::from("rebuilddb message recieved..."),
)))?;
let project = state.projects[state.selected_project].clone();
if let Some(mut db) = project.db.clone() {
state.main_tx.send(ToolMessage::Output((
0,
String::from("distrobox detected..."),
)))?;
disable_raw_mode()?;
execute!(
&stdout,
LeaveAlternateScreen,
cursor::Show,
DisableMouseCapture
)?;
std::io::stdout().flush()?;
if db.created {
if let Err(_) = db.destroy(state.main_tx.clone(), 0) {
state.main_tx.send(ToolMessage::Output((
0,
String::from("failed to destroy the project box!"),
)))?;
}
}
if let Err(_) = db.create(state.main_tx.clone(), 0) {
state.main_tx.send(ToolMessage::Output((
0,
String::from("failed to create the project box!"),
)))?;
}
enable_raw_mode()?;
execute!(
&stdout,
EnterAlternateScreen,
EnableMouseCapture,
cursor::Hide
)?;
terminal.clear()?;
}
}
ToolMessage::DestroyDB(mut db) => {
execute!(
&stdout,
LeaveAlternateScreen,
cursor::Show,
DisableMouseCapture
)?;
db.destroy(state.main_tx.clone(), 0)?;
enable_raw_mode()?;
execute!(
&stdout,
EnterAlternateScreen,
EnableMouseCapture,
cursor::Hide
)?;
terminal.clear()?;
}
ToolMessage::ConnectServer => {
if let Some(server) = state.servers.get_mut(state.selected_server) {
let _ = state.main_tx.send(ToolMessage::Output((
0,
"[success] Found Server in list!".to_string(),
)));
let mut address = String::new();
let lock_res = server.try_lock();
let mut connected = false;
match lock_res {
Ok(mut locked_server) => match locked_server.connect() {
Ok(id) => {
let _ = state.main_tx.send(ToolMessage::Output((
0,
format!(
"[success] Server Connected! ID: {} ADDRESS:{}",
id,
locked_server.address.clone()
),
)));
address = locked_server.address.clone();
connected = true;
}
Err(e) => {
let _ = state.main_tx.send(ToolMessage::Output((
0,
format!("[error] connecting to server {e}"),
)));
}
},
Err(e) => {
state
.output
.push(format!("[error] locking server state! {e}"));
}
}
if connected {
let server_clone = server.clone();
let tx_clone = state.main_tx.clone();
let mut timer = Duration::from_secs(5);
let mut last_check = Instant::now();
let server_handle = spawn(move || {
let server = server_clone;
loop {
let mut stop = false;
if Instant::now().duration_since(last_check) >= timer {
let lock_res = server.try_lock();
match lock_res {
Ok(mut locked_server) => {
locked_server.checkin();
last_check = locked_server.last_check.clone();
for action in locked_server.action_que.clone() {
if action.contains("|") {
let (action, data) =
action.split_once("|").unwrap();
match action.trim() {
"OUTPUT" => {
let _ = tx_clone.send(
ToolMessage::Output((
0,
data.trim().to_string(),
)),
);
}
"CMD" => {
if data.contains("SET_NAME") {
if let Some((_, name)) =
data.split_once("|")
{
locked_server.name =
name.trim()
.to_string();
}
} else {
let _ = tx_clone.send(
ToolMessage::Input((
locked_server
.client_id
.clone(),
data.trim()
.to_string(),
)),
);
}
}
_ => {}
}
} else {
match action.as_str() {
"disconnect" => {
let _ =
locked_server.disconnect();
stop = true;
}
_ => {
let _ = tx_clone.send(
ToolMessage::Output((
0,
action.clone(),
)),
);
}
}
}
}
timer = locked_server.timer.clone();
locked_server.action_que.clear();
if locked_server.message_que.is_empty() {
let id = locked_server.client_id.clone();
let pass = locked_server.password.clone();
locked_server.message_que.push(format!(
"{}**{}|||PING\n",
pass, id
));
}
}
Err(_) => {
let _ = tx_clone.send(ToolMessage::Output((
0,
"Couldn't lock server!".to_string(),
)));
}
}
}
if stop {
println!("{} disconnected!", address);
break;
}
sleep(Duration::from_secs(1));
}
});
handles.push(server_handle);
}
} else {
let _ = state.main_tx.send(ToolMessage::Output((
0,
"Couldn't Find selected Server!".to_string(),
)));
}
}
ToolMessage::SendServer(data) => {
if let Some(server) = state.servers.get(state.selected_server) {
if let Ok(mut server) = server.lock() {
server
.message_que
.push(format!("{}|||{}\n", data.0, data.1));
}
}
}
ToolMessage::DisconnectAllServers => {
println!("todo");
}
ToolMessage::DisconnectServer => {
if let Some(server) = state.servers.get(state.selected_server) {
if let Ok(mut lock) = server.lock() {
lock.action_que.push("disconnect".to_string());
let _ = state.main_tx.send(ToolMessage::Input((
0,
format!("server_remove_confim {}", state.selected_server),
)));
}
}
}
ToolMessage::ServerBrokerExit => {
state.server_broker_running = false;
}
ToolMessage::AppStateExit => {
println!("todo");
}
_ => {}
},
AppEvent::Key(key) => {
match key.code {
KeyCode::Esc => {
break;
}
KeyCode::Enter => {
let trimmed = state.curent_intput.trim().to_string();
if !trimmed.is_empty() {
state.history.push(trimmed.clone());
let _ = state
.main_tx
.send(ToolMessage::Output((0, "\n".to_string())));
let _ = state.main_tx.send(ToolMessage::Output((
0,
format!("[user input] > {}", trimmed),
)));
let _ = state
.main_tx
.send(ToolMessage::Output((0, "\n".to_string())));
let prompt = state.prompt.clone();
if prompt.action.is_some() {
state.prompt.responses.push(PromptResponse {
query: prompt.last_prompted.clone().unwrap(),
response: trimmed.clone(),
});
if state.prompt.responses.len() == state.prompt.num_responses {
state.execute_command(
prompt.execute_command.as_str(),
None,
0,
)?;
let _ = state.main_tx.send(ToolMessage::Output((
0,
format!("{}", prompt.execute_command),
)));
} else {
state.execute_command("prompt", None, 0)?;
}
} else {
let (command, args) =
trimmed.split_once(' ').unwrap_or((&trimmed, ""));
match command {
"exit" | "quit" => break,
"reload-modules" => {
state.initialize_modules();
let _ = state.main_tx.send(ToolMessage::Output((
0,
"Reloading Module Paths...".to_string(),
)));
}
"help" => {
let help_text = state.help.clone().join("\n");
for line in help_text.lines() {
let _ = state.main_tx.send(ToolMessage::Output((
0,
line.to_string(),
)));
}
}
"new_project" | "np" => {
if args.split_once(' ').is_some() {
let _ = state.execute_command(
command,
Some(args.to_string()),
0,
);
} else {
state.output.push(
"[error]: USAGE -> np <org> <name>".into(),
);
}
}
"switch_config" => {
switch_config = true;
state.curent_intput.clear();
break;
}
command_name => {
if state.module_loader.commands.contains_key(command) {
let _ = state.main_tx.send(ToolMessage::Output((
0,
format!(
"[Worker] Executing script {}...",
command_name
),
)));
if let Err(e) = state.execute_command(
command,
Some(args.to_string()),
0,
) {
let _ =
state.main_tx.send(ToolMessage::Output((
0,
format!("[error] Pipeline faile: {e}"),
)));
}
} else {
if args == "" {
let _ = state.execute_command(command, None, 0);
} else {
let _ = state.execute_command(
command,
Some(args.to_string()),
0,
);
}
}
}
}
}
state.curent_intput.clear();
}
history_index = state.history.len();
}
KeyCode::Char(c) => {
state.curent_intput.push(c);
}
KeyCode::Backspace => {
state.curent_intput.pop();
}
KeyCode::Up if key.modifiers.is_empty() => {
if !state.history.is_empty() && history_index > 0 {
history_index -= 1;
state.curent_intput = state.history[history_index].clone();
}
}
KeyCode::Down if key.modifiers.is_empty() => {
if history_index < state.history.len() {
history_index += 1;
if history_index == state.history.len() {
state.curent_intput.clear(); // Clear back to a fresh prompt
} else {
state.curent_intput = state.history[history_index].clone();
}
}
}
KeyCode::Up
if key
.modifiers
.contains(crossterm::event::KeyModifiers::CONTROL) =>
{
if let Some(selected) = project_list_state.selected() {
if selected > 0 {
let new_index = selected - 1;
project_list_state.select(Some(new_index));
state.selected_project = new_index;
}
}
}
KeyCode::Down
if key
.modifiers
.contains(crossterm::event::KeyModifiers::CONTROL) =>
{
if let Some(selected) = project_list_state.selected() {
if selected + 1 < state.projects.len() {
let new_index = selected + 1;
project_list_state.select(Some(new_index));
state.selected_project = new_index;
}
}
}
_ => {}
}
}
AppEvent::Mouse(mouse) => {
if let Ok(size) = terminal.size() {
let main_chunks = Layout::default()
.direction(Direction::Vertical)
.constraints([Constraint::Min(0), Constraint::Max(3)])
.split(ratatui::layout::Rect::from(size));
let top_chunks = Layout::default()
.direction(Direction::Horizontal)
.constraints([
Constraint::Percentage(20),
Constraint::Percentage(45),
Constraint::Percentage(35),
])
.split(main_chunks[0]);
let left_side_cutoff = top_chunks[0].width + top_chunks[1].width;
if mouse.column < left_side_cutoff {
if mouse.kind == crossterm::event::MouseEventKind::ScrollUp {
state.output_follow = false;
state.output_scroll = state.output_scroll.saturating_sub(1);
} else if mouse.kind == crossterm::event::MouseEventKind::ScrollDown {
state.output_follow = false;
state.output_scroll = state.output_scroll.saturating_add(1);
}
} else {
if mouse.kind == crossterm::event::MouseEventKind::ScrollUp {
state.info_scroll = state.info_scroll.saturating_sub(1);
} else if mouse.kind == crossterm::event::MouseEventKind::ScrollDown {
state.info_scroll = state.info_scroll.saturating_add(1);
}
}
let max_scroll = state
.output
.len()
.saturating_sub(top_chunks[1].height.saturating_sub(2) as usize)
as u16;
if state.output_scroll >= max_scroll {
state.output_follow = true;
}
}
}
}
if !state.app_state_running && state.server_broker_running {
break;
}
}
}
disable_raw_mode()?;
execute!(
terminal.backend_mut(),
LeaveAlternateScreen,
cursor::Show,
DisableMouseCapture
)?;
println!("disconnecting from servers...");
state.servers.iter().for_each(|server| {
if let Ok(mut server_lock) = server.lock() {
server_lock.action_que.push("disconnect".to_string());
}
});
for handle in handles {
handle.join().unwrap();
}
if switch_config {
let handles_clone = main_handles.clone();
if let Ok(mut lock) = main_handles.lock() {
lock.push(spawn(move || {
let _ = startup(state, true, handles_clone);
}));
}
}
Ok(())
}