Files
tetanus/src/funcs.rs
T

705 lines
33 KiB
Rust

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<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 server_broker(
rx: Receiver<ServerBrokerCmd>,
tx: Sender<ToolMessage>,
servers: Arc<Mutex<Vec<Server>>>,
) {
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::<usize>() {
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::<usize>() {
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<ToolMessage>,
) -> Result<(), Box<dyn std::error::Error>> {
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::<AppEvent>();
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<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(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<Line> = 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 <org> <name>".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(())
}