fixed a bug on new project creation where the config file path wasn't
being passed correctly. Also started the groundwork for adding more of my rust based tooling to tetanus.
This commit is contained in:
+56
-4
@@ -49,6 +49,7 @@ pub fn startup(
|
||||
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)?;
|
||||
println!("entering tui...");
|
||||
run_tui(app, rx, handles.clone())?;
|
||||
} else {
|
||||
eprintln!("error invalid selection! exiting...");
|
||||
@@ -56,11 +57,13 @@ pub fn startup(
|
||||
}
|
||||
} else {
|
||||
app.load_config(state.config_file.clone(), true)?;
|
||||
println!("entering tui...");
|
||||
run_tui(app, rx, handles.clone())?;
|
||||
}
|
||||
if nested {
|
||||
let (tx, rx) = channel();
|
||||
state.main_tx = tx;
|
||||
println!("returning to prevous tui...");
|
||||
run_tui(state, rx, handles.clone())?;
|
||||
}
|
||||
Ok(())
|
||||
@@ -475,6 +478,35 @@ pub fn run_tui(
|
||||
);
|
||||
}
|
||||
"CMD" => {
|
||||
if let Some((cmd, data)) =
|
||||
data.split_once("|")
|
||||
{
|
||||
match cmd.trim() {
|
||||
"SET_NAME" => {
|
||||
locked_server
|
||||
.name = data
|
||||
.trim()
|
||||
.to_string();
|
||||
}
|
||||
"LIST_PROJECTS" => {
|
||||
let _ = ToolMessage::Input((locked_server.client_id.clone(), "list_projects".to_string()));
|
||||
}
|
||||
_ => {
|
||||
let msg = format!(
|
||||
"MSG FROM: {}",
|
||||
data.trim()
|
||||
);
|
||||
let _ = tx_clone.send(
|
||||
ToolMessage::Output((
|
||||
locked_server
|
||||
.client_id
|
||||
.clone(),
|
||||
msg,
|
||||
)),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
if data.contains("SET_NAME") {
|
||||
if let Some((_, name)) =
|
||||
data.split_once("|")
|
||||
@@ -485,7 +517,7 @@ pub fn run_tui(
|
||||
}
|
||||
} else {
|
||||
let _ = tx_clone.send(
|
||||
ToolMessage::Input((
|
||||
ToolMessage::Output((
|
||||
locked_server
|
||||
.client_id
|
||||
.clone(),
|
||||
@@ -680,13 +712,33 @@ pub fn run_tui(
|
||||
}
|
||||
} else {
|
||||
if args == "" {
|
||||
let _ = state.execute_command(command, None, 0);
|
||||
match state.execute_command(command, None, 0) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
let _ = state.main_tx.send(
|
||||
ToolMessage::Output((
|
||||
0,
|
||||
e.to_string(),
|
||||
)),
|
||||
);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let _ = state.execute_command(
|
||||
match state.execute_command(
|
||||
command,
|
||||
Some(args.to_string()),
|
||||
0,
|
||||
);
|
||||
) {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
let _ = state.main_tx.send(
|
||||
ToolMessage::Output((
|
||||
0,
|
||||
e.to_string(),
|
||||
)),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+696
-148
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,88 @@
|
||||
use crate::*;
|
||||
use anyhow::Context;
|
||||
use headless_chrome::{Browser, LaunchOptions};
|
||||
use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
|
||||
|
||||
pub fn rustwitness(
|
||||
input: PathBuf,
|
||||
output: PathBuf,
|
||||
proxy: Option<String>,
|
||||
tx: Sender<ToolMessage>,
|
||||
) {
|
||||
fn capture_screenshot(
|
||||
url: &str,
|
||||
proxy: Option<String>,
|
||||
output: PathBuf,
|
||||
) -> Result<(), Box<dyn Error>> {
|
||||
if !output.exists() {
|
||||
create_dir_all(&output)?;
|
||||
}
|
||||
let mut launch_options = LaunchOptions::default();
|
||||
let file_name = url.split("://").collect::<Vec<&str>>()[1].to_string();
|
||||
if let Some(ref proxy_url) = proxy {
|
||||
let proxy_arg = format!("--proxy-server={}", proxy_url);
|
||||
launch_options
|
||||
.args
|
||||
.push(std::ffi::OsStr::new(Box::leak(proxy_arg.into_boxed_str())));
|
||||
}
|
||||
let browser = Browser::new(launch_options).context("Browser launch failed")?;
|
||||
let tab = browser.new_tab().context("Failed to open tab")?;
|
||||
tab.navigate_to(url).context("Navigation failed")?;
|
||||
tab.wait_until_navigated()
|
||||
.context("Waiting for load failed")?;
|
||||
tab.wait_for_element("body")?;
|
||||
tab.evaluate("document.readyState === 'complete'", false)?;
|
||||
let mut file_path = output.clone();
|
||||
let mut log_path = output.clone();
|
||||
log_path.push("log.txt");
|
||||
file_path.push(format!("{}.png", file_name));
|
||||
let png_data = tab
|
||||
.capture_screenshot(
|
||||
headless_chrome::protocol::cdp::Page::CaptureScreenshotFormatOption::Png,
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.context("Screenshot capture failed")?;
|
||||
|
||||
fs::write(&file_path, png_data)
|
||||
.context("Failed to write file")
|
||||
.context(format!(
|
||||
"failed to write png file! {}",
|
||||
&file_path.display()
|
||||
))?;
|
||||
println!("Successfully captured: {}", url);
|
||||
fs::write(&log_path, format!("{}\n", url).as_bytes())
|
||||
.context("failed to write log file!")?;
|
||||
Ok(())
|
||||
}
|
||||
match File::open(&input) {
|
||||
Ok(file) => {
|
||||
let urls: Vec<String> = BufReader::new(file)
|
||||
.lines()
|
||||
.filter_map(|line| line.ok())
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.collect();
|
||||
let _ = tx.send(ToolMessage::Output((
|
||||
0,
|
||||
format!("Starting capture of {} urls...", urls.len()),
|
||||
)));
|
||||
urls.par_iter().enumerate().for_each(|(_index, url)| {
|
||||
if let Err(e) = capture_screenshot(url, proxy.clone(), output.clone()) {
|
||||
let _ = tx.send(ToolMessage::Output((
|
||||
0,
|
||||
format!("[error] rustwitness: Failed to process {}: {:?}", url, e),
|
||||
)));
|
||||
}
|
||||
});
|
||||
let _ = tx.send(ToolMessage::Output((
|
||||
0,
|
||||
"[success] all urls scanned for screenshots!".to_string(),
|
||||
)));
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = tx.send(ToolMessage::Output((0, format!("[error] rustwitness {e}"))));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -22,6 +22,7 @@ use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer};
|
||||
enum ClientAction {
|
||||
Output(String),
|
||||
Cmd(String),
|
||||
Collab(String),
|
||||
}
|
||||
|
||||
pub struct Client {
|
||||
@@ -31,6 +32,7 @@ pub struct Client {
|
||||
actions: Vec<ClientAction>,
|
||||
controlling: usize,
|
||||
controlled: usize,
|
||||
collabing: Vec<usize>,
|
||||
output_que: Vec<String>,
|
||||
logged_in: bool,
|
||||
timer: Duration,
|
||||
@@ -201,6 +203,7 @@ where
|
||||
actions: Vec::new(),
|
||||
controlled: 0,
|
||||
controlling: 0,
|
||||
collabing: Vec::new(),
|
||||
output_que: Vec::new(),
|
||||
logged_in: false,
|
||||
timer: Duration::from_secs(5),
|
||||
@@ -312,6 +315,26 @@ where
|
||||
dest_client.id, cmd
|
||||
)));
|
||||
}
|
||||
else{
|
||||
let mut colabs = Vec::new();
|
||||
lock.clients.iter_mut().filter(|c| c.id == source_id).for_each(|c|{
|
||||
if c.collabing.is_empty(){
|
||||
c.actions.push(ClientAction::Output("Error no controlled or collabed clients associated with you!".to_string()));
|
||||
}
|
||||
else{
|
||||
c.collabing.iter().for_each(|id|{
|
||||
colabs.push(id.clone());
|
||||
});
|
||||
}
|
||||
});
|
||||
if !colabs.is_empty(){
|
||||
colabs.iter().for_each(|id|{
|
||||
lock.clients.iter_mut().filter(|c| &c.id == id).for_each(|c|{
|
||||
c.actions.push(ClientAction::Cmd(data.to_string()));
|
||||
});
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
"CONTROL" => {
|
||||
let new_control_id = data.trim().parse::<usize>().unwrap();
|
||||
@@ -429,6 +452,22 @@ where
|
||||
c.remove = true;
|
||||
});
|
||||
}
|
||||
"COLLAB" => {
|
||||
let dest_id: usize = data.trim().parse().unwrap();
|
||||
lock.clients
|
||||
.iter_mut()
|
||||
.filter(|c| c.id == source_id)
|
||||
.for_each(|c| {
|
||||
c.collabing.push(dest_id.clone());
|
||||
});
|
||||
lock.clients
|
||||
.iter_mut()
|
||||
.filter(|d| d.id == dest_id)
|
||||
.for_each(|d| {
|
||||
d.collabing.push(dest_id.clone());
|
||||
});
|
||||
println!("{} is collabing with {}", source_id, dest_id);
|
||||
}
|
||||
"SLEEP" => {
|
||||
if let Ok(secs) = data.parse::<u64>() {
|
||||
lock.clients
|
||||
@@ -455,6 +494,9 @@ where
|
||||
println!("adding {} to output", text);
|
||||
messages.push(format!("OUTPUT|{}", text));
|
||||
}
|
||||
ClientAction::Collab(text) => {
|
||||
println!("todo");
|
||||
}
|
||||
}
|
||||
}
|
||||
let full_message = messages.join("||");
|
||||
@@ -523,6 +565,9 @@ where
|
||||
println!("adding {} to output", text);
|
||||
messages.push(format!("OUTPUT|{}", text));
|
||||
}
|
||||
ClientAction::Collab(text) => {
|
||||
println!("todo");
|
||||
}
|
||||
}
|
||||
}
|
||||
let full_message = messages.join("||");
|
||||
|
||||
Reference in New Issue
Block a user