wrote the tool
This commit is contained in:
+10
@@ -0,0 +1,10 @@
|
||||
[package]
|
||||
name = "rustwitness"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
anyhow = "1.0.102"
|
||||
clap = {version = "4.6.1", features = ["derive"]}
|
||||
headless_chrome = "1.0.21"
|
||||
rayon = "1.12.0"
|
||||
+91
@@ -0,0 +1,91 @@
|
||||
use anyhow::{Context, Result};
|
||||
use clap::Parser;
|
||||
use headless_chrome::{Browser, LaunchOptions};
|
||||
use rayon::prelude::*;
|
||||
use std::fs::{self, File};
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(author, version, about = "Parallel web screenshot tool")]
|
||||
struct Args {
|
||||
#[arg(short, long)]
|
||||
input: PathBuf,
|
||||
|
||||
#[arg(short, long)]
|
||||
output: PathBuf,
|
||||
|
||||
#[arg(short, long)]
|
||||
proxy: Option<String>,
|
||||
|
||||
#[arg(short, long, default_value_t = 4)]
|
||||
workers: usize,
|
||||
}
|
||||
|
||||
fn main() -> Result<()> {
|
||||
let args = Args::parse();
|
||||
|
||||
if !args.output.exists() {
|
||||
fs::create_dir_all(&args.output).context("Failed to create output directory")?;
|
||||
}
|
||||
|
||||
rayon::ThreadPoolBuilder::new()
|
||||
.num_threads(args.workers)
|
||||
.build_global()
|
||||
.unwrap();
|
||||
|
||||
let file = File::open(&args.input).context("Failed to open input 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();
|
||||
|
||||
println!(
|
||||
"Starting capture of {} URLs using {} workers...",
|
||||
urls.len(),
|
||||
args.workers
|
||||
);
|
||||
|
||||
urls.par_iter().enumerate().for_each(|(index, url)| {
|
||||
if let Err(e) = capture_screenshot(url, index, &args) {
|
||||
eprintln!("Failed to process {}: {:?}", url, e);
|
||||
}
|
||||
});
|
||||
|
||||
println!("All tasks completed.");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn capture_screenshot(url: &str, index: usize, args: &Args) -> Result<()> {
|
||||
let mut launch_options = LaunchOptions::default();
|
||||
|
||||
if let Some(ref proxy_url) = args.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")?;
|
||||
let mut file_path = args.output.clone();
|
||||
file_path.push(format!("screenshot_{:03}.png", index + 1));
|
||||
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")?;
|
||||
println!("Successfully captured: {}", url);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user