use crate::{ToolMessage::Input, *}; 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::sync::Arc; use std::sync::Mutex; use std::time::Duration; use std::{error::Error, time::Instant}; use std::{ io::{Read, Write}, usize, }; use std::{net::TcpStream, thread}; pub enum ServerBrokerCmd { ConnectServer(usize), RegisterActionOutput((Server, String)), Disconnect(Server), DisconnectAll, Exit, } 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 server_broker( rx: Receiver, tx: Sender, servers: Arc>>, ) { loop { if let Ok(msg) = rx.try_recv() { match msg { ServerBrokerCmd::ConnectServer(id) => { let _ = tx.send(ToolMessage::Output(( 0, format!("attempting to connect to server {}", id), ))); if let Ok(mut lock) = servers.lock() { lock.iter_mut().for_each(|s| { let _ = tx.send(ToolMessage::Output((0, format!("got server list lock!")))); if s.id == id { let _ = tx.send(ToolMessage::Output(( 0, format!("Found the server for {}: {}", s.id, s.address), ))); if let Ok(mut stream) = TcpStream::connect(s.address.clone()) { let _ = tx.send(ToolMessage::Output(( 0, String::from("TCP successful to server!"), ))); let _ = stream.write(format!("HELLO|{}\n", s.id).as_bytes()); let _ = tx.send(ToolMessage::Output(( 0, String::from("Hello message sent!"), ))); let mut buf = [0; 1024]; if let Ok(bytes_read) = stream.read(&mut buf) { let response = String::from_utf8_lossy(&buf[..bytes_read]); if let Some((_, id)) = response.split_once("|") { if let Ok(id) = id.trim().parse::() { s.client_id = id; s.connected = true; let _ = tx.send(ToolMessage::Output(( 0, String::from("got client ID from server!"), ))); } else { let _ = tx.send(ToolMessage::Output(( 0, String::from( "failed to get response from server!", ), ))); } } } } } }); } } ServerBrokerCmd::Disconnect(server) => { if let Ok(mut lock) = servers.lock() { let mut rid = 0; let mut found = false; lock.iter().enumerate().for_each(|(id, s)| { if s.id == server.id { rid = id; found = true; } }); if found { lock.remove(rid); } } } ServerBrokerCmd::DisconnectAll => { println!("todo"); } ServerBrokerCmd::Exit => { println!("todo"); break; } ServerBrokerCmd::RegisterActionOutput(text) => { if let Ok(mut lock) = servers.lock() { lock.iter_mut().for_each(|s| { if s.id == text.0.id { s.message_que.push(text.1.clone()); } }); } } } } if let Ok(mut lock) = servers.try_lock() { lock.iter_mut().for_each(|s| { if s.connected { if Instant::now().duration_since(s.last_check) >= s.timer { if let Ok(mut stream) = TcpStream::connect(s.address.clone()) { s.last_check = Instant::now(); let mut success_ids = Vec::new(); if s.message_que.len() > 0 { s.message_que.iter().enumerate().for_each(|(id, msg)| { if let Ok(_) = stream.write( format!("{}|{}\n", s.client_id.clone(), msg).as_bytes(), ) { success_ids.push(id); } }); } else { let _ = stream.write(format!("{}|PING", s.client_id).as_bytes()); } for id in success_ids { s.message_que.remove(id); } let mut buffer = [0; 1024]; if let Ok(bytes_read) = stream.read(&mut buffer) { let response = String::from_utf8_lossy(&buffer[..bytes_read]).to_string(); let (_, action_string) = response.split_once("|").unwrap_or(("ACTIONS", "NONE")); let actions: Vec<&str> = action_string.split(",").collect(); for action in actions { let (action, data) = action.split_once("|").unwrap_or(("PING", "NONE")); match action { "OUTPUT" => { let _ = tx.send(ToolMessage::Output(( 0, format!("FROM SERVER: {}", data), ))); } "CMD" => { let (rid_str, cmd) = data.split_once("|").unwrap_or(("0", "NONE")); if let Ok(rid) = rid_str.parse::() { let _ = tx.send(ToolMessage::Input(( rid, cmd.to_string(), ))); } } _ => {} } } } } } } }); } thread::sleep(Duration::from_millis(10)); } } pub fn run_tui( mut state: AppState, main_rx: Receiver, ) -> Result<(), Box> { enable_raw_mode()?; let mut stdout = io::stdout(); execute!( stdout, EnterAlternateScreen, cursor::Hide, EnableMouseCapture )?; let (server_tx, server_rx) = channel(); let main_tx = state.main_tx.clone(); let server_clone = state.servers.clone(); std::thread::spawn(move || { server_broker(server_rx, main_tx, server_clone); }); 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(); 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(); loop { 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]); 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(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 output_lines: Vec = state .output .iter() .map(|text| Line::from(Span::raw(text))) .collect(); let text_area_height = top_chunks[1].height.saturating_sub(2) as usize; 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)) .wrap(ratatui::widgets::Wrap { trim: false }); 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() { match event { AppEvent::Worker(msg) => match msg { ToolMessage::EndPrompt => { state.prompt.reset(); } ToolMessage::Input(cmd) => { state.log.push(cmd.1.clone()); } ToolMessage::Output((rid, txt)) => { state.log.push(txt.clone()); state.output.push(txt.clone()); state.output_scroll = u16::MAX; if rid != 0 { if let Some(server) = state.selected_server.clone() { let _ = server_tx.send(ServerBrokerCmd::RegisterActionOutput(( server, format!("OUTPUT|{}|{}", rid, 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.selected_server.clone() { let _ = server_tx.send(ServerBrokerCmd::ConnectServer(server.id)); } } ToolMessage::SendServer(data) => { let _ = server_tx.send(ServerBrokerCmd::RegisterActionOutput(data)); } ToolMessage::DisconnectAllServers => { println!("todo"); } ToolMessage::DisconnectServer => { println!("todo"); } 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()); state.output.push("\n".to_string()); state.output.push(format!("[user input] > {}", trimmed)); state.output.push("\n".to_string()); let prompt = state.prompt.clone(); if prompt.action.is_some() { state.prompt.responses.push(trimmed.clone()); if state.prompt.responses.len() == state.prompt.num_responses { state.execute_command( prompt.execute_command.as_str(), None, 0, )?; } else { state.execute_command("prompt", None, 0)?; } } let (command, args) = trimmed.split_once(' ').unwrap_or((&trimmed, "")); if state.selected_client != 0 { if let Some(server) = state.selected_server.clone() { let _ = server_tx.send( ServerBrokerCmd::RegisterActionOutput(( server.clone(), format!( "CMD|{}|{}|{} {}", server.client_id, state.selected_client, command, args ), )), ); } } else { match command { "exit" | "quit" => break, "reload-modules" => { state.initialize_modules(); state.output.push("Reloading module paths...".into()); } "help" => { let help_text = state.help.clone().join("\n"); for line in help_text.lines() { state.output.push(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()); } } command_name => { if state.module_loader.commands.contains_key(command) { state.output.push(format!( "[Worker] Executing script '{}'...", command_name )); if let Err(e) = state.execute_command( command, Some(args.to_string()), 0, ) { state.output.push(format!( "[Error] Pipeline fail: {}", 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_scroll = state.output_scroll.saturating_sub(1); } else if mouse.kind == crossterm::event::MouseEventKind::ScrollDown { 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); } } } } } if !state.app_state_running && state.server_broker_running { break; } } } disable_raw_mode()?; execute!( terminal.backend_mut(), LeaveAlternateScreen, cursor::Show, DisableMouseCapture )?; Ok(()) }