From d14f02ad53838a26b1542c97fa8b8f350b0a1955 Mon Sep 17 00:00:00 2001 From: pyro Date: Fri, 28 Aug 2026 14:38:45 -0500 Subject: [PATCH] made the rustwetness stuff work! --- src/lib.rs | 64 +++++++++++++++++++++++++++++++++++++++++------ src/rust_tools.rs | 45 ++++++++------------------------- 2 files changed, 66 insertions(+), 43 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 6e93930..552dd7c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,6 +4,7 @@ use fs_extra::dir::{CopyOptions, copy}; use ipnet::IpNet; use keyring::Entry; use ratatui::crossterm::event; +use rayon::prelude::*; use rhai::{AST, Dynamic, Engine, Scope}; use rustls::{ClientConfig, ClientConnection, RootCertStore, StreamOwned, pki_types::ServerName}; use std::collections::HashMap; @@ -994,10 +995,19 @@ impl AppState { if let Some(project) = self.projects.get(self.selected_project) { if let Some(args) = command_args { if let Some((outfile, proxy)) = args.split_once(" ") { + let outfile = format!("{}/{}", project.files.display(), outfile); let mut proxy_string = None; if let Some((ip, port)) = proxy.split_once(" ") { proxy_string = Some(format!("socks5://{}:{}", ip, port)); } + let _ = self.main_tx.send(ToolMessage::Output(( + rid, + format!( + "[success] rustwitness running will save screenshots to {}", + outfile + ), + ))); + let mut urls = Vec::new(); WalkDir::new(&project.files) .into_iter() .filter(|res| res.is_ok()) @@ -1005,17 +1015,55 @@ impl AppState { if let Ok(entry) = res { if entry.file_name().to_string_lossy().contains("urls.txt") { - let input = PathBuf::from(entry.path()); - let mut output = project.files.clone(); - output.push(outfile.trim()); - let proxy_intput = proxy_string.clone(); - let tx = self.main_tx.clone(); - self.workers.spawn(move || { - rustwitness(input, output, proxy_intput, tx) - }); + if let Ok(input) = read_to_string(entry.path()) { + input.lines().into_iter().for_each(|url| { + urls.push(url.to_string()); + }); + } } } }); + let proxy_string_clone = proxy_string.clone(); + let out_folder = PathBuf::from(outfile); + let tx_clone = self.main_tx.clone(); + self.workers.spawn(move || { + urls.par_iter().for_each(|url| { + if url.contains("://"){ + match rustwitness( + url.clone(), + out_folder.clone(), + proxy_string_clone.clone(), + ) { + Ok(_) => { + let _ = tx_clone.send(ToolMessage::Output(( + rid, + format!("[success] {} captured!", url), + ))); + } + Err(e) => { + let _ = tx_clone.send(ToolMessage::Output(( + rid, + format!( + "[error] capturing screenshot of {url}: {e}" + ), + ))); + } + } + } + else{ + let _ = tx_clone.send(ToolMessage::Output(( + rid, + format!( + "[error] capturing screenshot of {url}: not a valid url." + ), + ))); + } + }); + let _ = tx_clone.send(ToolMessage::Output(( + rid, + "[success] all urls scanned for screenshots!".to_string(), + ))); + }); } } else { let _ = self.main_tx.send(ToolMessage::Output(( diff --git a/src/rust_tools.rs b/src/rust_tools.rs index 300da64..60d6a8b 100644 --- a/src/rust_tools.rs +++ b/src/rust_tools.rs @@ -1,14 +1,12 @@ use crate::*; use anyhow::Context; use headless_chrome::{Browser, LaunchOptions}; -use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator}; pub fn rustwitness( - input: PathBuf, + input: String, output: PathBuf, proxy: Option, - tx: Sender, -) { +) -> Result<(), Box> { fn capture_screenshot( url: &str, proxy: Option, @@ -18,7 +16,12 @@ pub fn rustwitness( create_dir_all(&output)?; } let mut launch_options = LaunchOptions::default(); - let file_name = url.split("://").collect::>()[1].to_string(); + let file_name; + if url.contains("://") { + file_name = url.split("://").collect::>()[1].to_string(); + } else { + file_name = url.to_string(); + } if let Some(ref proxy_url) = proxy { let proxy_arg = format!("--proxy-server={}", proxy_url); launch_options @@ -51,38 +54,10 @@ pub fn rustwitness( "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 = 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}")))); - } - } + capture_screenshot(&input, proxy.clone(), output.clone())?; + Ok(()) }