made the rustwetness stuff work!

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());
+                                            if let Ok(input) = read_to_string(entry.path()) {
-                                            let mut output = project.files.clone();
+                                                input.lines().into_iter().for_each(|url| {
-                                            output.push(outfile.trim());
+                                                    urls.push(url.to_string());
-                                            let proxy_intput = proxy_string.clone();
+                                                });
-                                            let tx = self.main_tx.clone();
+                                            }
-                                            self.workers.spawn(move || {
-                                                rustwitness(input, output, proxy_intput, tx)
-                                            });
                                         }
                                     }
                                 });
+                            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((

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<String>,
-    tx: Sender<ToolMessage>,
+) -> Result<(), Box<dyn Error>> {
-) {
     fn capture_screenshot(
         url: &str,
         proxy: Option<String>,
@@ -18,7 +16,12 @@ pub fn rustwitness(
             create_dir_all(&output)?;
         }
         let mut launch_options = LaunchOptions::default();
-        let file_name = url.split("://").collect::<Vec<&str>>()[1].to_string();
+        let file_name;
+        if url.contains("://") {
+            file_name = url.split("://").collect::<Vec<&str>>()[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) {
+    capture_screenshot(&input, proxy.clone(), output.clone())?;
-        Ok(file) => {
+    Ok(())
-            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}"))));
-        }
-    }
 }