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> { 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>>>, ) -> Result<(), Box> { 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, main_handles: Arc>>>, ) -> Result<(), Box> { 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::(); 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 = 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 = 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 ".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(()) }