pyro/tetanus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

made the rustwetness stuff work!

Browse Source
This commit is contained in:
pyro
2026-08-28 14:38:45 -05:00
parent 9da3472c59
commit d14f02ad53
2 changed files with 66 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lib.rs
+56 -8
View File
@@ -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((
src/rust_tools.rs
+10 -35
View File
@@ -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) {
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}"))));
}
}
capture_screenshot(&input, proxy.clone(), output.clone())?;
Ok(())
}