added code to generate victim binaries, implemented client -> server
authentication and fixed a bug in the ui where not all output would be displayed.
This commit is contained in:
+175
-79
@@ -82,6 +82,7 @@ pub fn run_tui(
|
||||
let mut history_index = state.history.len();
|
||||
let mut handles = Vec::new();
|
||||
loop {
|
||||
let mut text_area_width = 0;
|
||||
terminal.draw(|f| {
|
||||
let main_chunks = Layout::default()
|
||||
.direction(Direction::Vertical)
|
||||
@@ -95,6 +96,7 @@ pub fn run_tui(
|
||||
Constraint::Percentage(35),
|
||||
])
|
||||
.split(main_chunks[0]);
|
||||
text_area_width = top_chunks[1].width.saturating_sub(2);
|
||||
let projects: Vec<ListItem> = state
|
||||
.projects
|
||||
.iter()
|
||||
@@ -194,12 +196,25 @@ pub fn run_tui(
|
||||
.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 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;
|
||||
@@ -213,8 +228,7 @@ pub fn run_tui(
|
||||
.borders(Borders::ALL)
|
||||
.title(" Script Engine Output "),
|
||||
)
|
||||
.scroll((state.output_scroll, 0))
|
||||
.wrap(ratatui::widgets::Wrap { trim: false });
|
||||
.scroll((state.output_scroll, 0));
|
||||
|
||||
f.render_widget(output_paragraph, top_chunks[1]);
|
||||
let input_paragraph = Paragraph::new(state.curent_intput.as_str()).block(
|
||||
@@ -225,7 +239,7 @@ pub fn run_tui(
|
||||
|
||||
f.render_widget(input_paragraph, main_chunks[1]);
|
||||
})?;
|
||||
if let Ok(event) = event_rx.recv() {
|
||||
if let Ok(event) = event_rx.recv_timeout(Duration::from_millis(100)) {
|
||||
match event {
|
||||
AppEvent::Worker(msg) => match msg {
|
||||
ToolMessage::EndPrompt => {
|
||||
@@ -241,9 +255,26 @@ pub fn run_tui(
|
||||
}
|
||||
}
|
||||
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());
|
||||
state.output.push(txt.clone());
|
||||
state.output_scroll = u16::MAX;
|
||||
if rid != 0 {
|
||||
if let Some(server) = state.servers.get(state.selected_server) {
|
||||
if let Ok(mut server) = server.lock() {
|
||||
@@ -319,30 +350,38 @@ pub fn run_tui(
|
||||
}
|
||||
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) => {
|
||||
state.output.push(format!(
|
||||
"Server Connected! ID:{} ADDRESS:{}",
|
||||
id,
|
||||
locked_server.address.clone()
|
||||
));
|
||||
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) => {
|
||||
state
|
||||
.output
|
||||
.push(format!("Error connecting to server {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}"));
|
||||
.push(format!("[error] locking server state! {e}"));
|
||||
}
|
||||
}
|
||||
if connected {
|
||||
@@ -374,21 +413,33 @@ pub fn run_tui(
|
||||
);
|
||||
}
|
||||
"CMD" => {
|
||||
let _ = tx_clone.send(
|
||||
ToolMessage::Input((
|
||||
locked_server
|
||||
.client_id
|
||||
.clone(),
|
||||
data.trim().to_string(),
|
||||
)),
|
||||
);
|
||||
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" => {
|
||||
locked_server.connected = false;
|
||||
let _ =
|
||||
locked_server.disconnect();
|
||||
stop = true;
|
||||
}
|
||||
_ => {
|
||||
@@ -406,12 +457,19 @@ pub fn run_tui(
|
||||
locked_server.action_que.clear();
|
||||
if locked_server.message_que.is_empty() {
|
||||
let id = locked_server.client_id.clone();
|
||||
locked_server
|
||||
.message_que
|
||||
.push(format!("{}|||PING\n", id));
|
||||
let pass = locked_server.password.clone();
|
||||
locked_server.message_que.push(format!(
|
||||
"{}**{}|||PING\n",
|
||||
pass, id
|
||||
));
|
||||
}
|
||||
}
|
||||
Err(_) => {}
|
||||
Err(_) => {
|
||||
let _ = tx_clone.send(ToolMessage::Output((
|
||||
0,
|
||||
"Couldn't lock server!".to_string(),
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
if stop {
|
||||
@@ -423,6 +481,11 @@ pub fn run_tui(
|
||||
});
|
||||
handles.push(server_handle);
|
||||
}
|
||||
} else {
|
||||
let _ = state.main_tx.send(ToolMessage::Output((
|
||||
0,
|
||||
"Couldn't Find selected Server!".to_string(),
|
||||
)));
|
||||
}
|
||||
}
|
||||
ToolMessage::SendServer(data) => {
|
||||
@@ -457,9 +520,16 @@ pub fn run_tui(
|
||||
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 _ = 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(trimmed.clone());
|
||||
@@ -469,61 +539,77 @@ pub fn run_tui(
|
||||
None,
|
||||
0,
|
||||
)?;
|
||||
let _ = state.main_tx.send(ToolMessage::Output((
|
||||
0,
|
||||
format!("{}", prompt.execute_command),
|
||||
)));
|
||||
} else {
|
||||
state.execute_command("prompt", None, 0)?;
|
||||
}
|
||||
}
|
||||
let (command, args) =
|
||||
trimmed.split_once(' ').unwrap_or((&trimmed, ""));
|
||||
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()),
|
||||
} 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,
|
||||
);
|
||||
} else {
|
||||
state
|
||||
.output
|
||||
.push("Error: USAGE -> np <org> <name>".into());
|
||||
"Reloading Module Paths...".to_string(),
|
||||
)));
|
||||
}
|
||||
}
|
||||
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));
|
||||
"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(),
|
||||
)));
|
||||
}
|
||||
} else {
|
||||
if args == "" {
|
||||
let _ = state.execute_command(command, None, 0);
|
||||
} else {
|
||||
}
|
||||
"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) {
|
||||
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,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -601,8 +687,10 @@ pub fn run_tui(
|
||||
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 {
|
||||
@@ -612,6 +700,14 @@ pub fn run_tui(
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user