1 Commits

Author SHA1 Message Date
pyro 940673b381 made the healdess chromium window wait until the body and javascript
loaded before taking the screenshot.
2026-04-30 17:59:02 -05:00
2 changed files with 1689 additions and 30 deletions
Generated
+1625
View File
File diff suppressed because it is too large Load Diff
+64 -30
View File
@@ -2,7 +2,7 @@ use anyhow::{Context, Result};
use clap::Parser; use clap::Parser;
use headless_chrome::{Browser, LaunchOptions}; use headless_chrome::{Browser, LaunchOptions};
use rayon::prelude::*; use rayon::prelude::*;
use std::fs::{self, File}; use std::fs::{self, File, read_to_string};
use std::io::{BufRead, BufReader}; use std::io::{BufRead, BufReader};
use std::path::PathBuf; use std::path::PathBuf;
@@ -24,43 +24,68 @@ struct Args {
fn main() -> Result<()> { fn main() -> Result<()> {
let args = Args::parse(); let args = Args::parse();
if !args.output.exists() { if !args.output.exists() {
fs::create_dir_all(&args.output).context("Failed to create output directory")?; fs::create_dir_all(&args.output).context("Failed to create output directory")?;
} }
let mut log_path = args.output.clone();
log_path.push("log.txt");
let mut resume = None;
if log_path.exists() {
let log = read_to_string(log_path)?;
if let Some(last_line) = log.lines().last() {
resume = Some(last_line.to_string())
}
};
rayon::ThreadPoolBuilder::new() rayon::ThreadPoolBuilder::new()
.num_threads(args.workers) .num_threads(args.workers)
.build_global() .build_global()
.unwrap(); .unwrap();
let mut urls = Vec::new();
let file = File::open(&args.input).context("Failed to open input file")?; if resume.is_some() {
let urls: Vec<String> = BufReader::new(file) let url_resume = resume.clone().unwrap();
.lines() let url_text = read_to_string(&args.input)?;
.filter_map(|line| line.ok()) let mut resume_met = false;
.map(|s| s.trim().to_string()) for line in url_text.lines() {
.filter(|s| !s.is_empty()) if !resume_met {
.collect(); if line.trim().to_string() == url_resume {
resume_met = true;
println!( }
"Starting capture of {} URLs using {} workers...", } else {
urls.len(), urls.push(line.trim().to_string());
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);
} }
}); }
if resume.is_none() {
println!("All tasks completed."); 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, &args) {
eprintln!("Failed to process {}: {:?}", url, e);
}
});
println!("All tasks completed.");
} else {
urls.par_iter().enumerate().for_each(|(_index, url)| {
if let Err(e) = capture_screenshot(url, &args) {
eprintln!("failed to process {}: {:?}", url, e);
}
});
}
Ok(()) Ok(())
} }
fn capture_screenshot(url: &str, args: &Args) -> Result<()> {
fn capture_screenshot(url: &str, index: usize, args: &Args) -> Result<()> {
let mut launch_options = LaunchOptions::default(); let mut launch_options = LaunchOptions::default();
let file_name = url.split("://").collect::<Vec<&str>>()[1].to_string();
if let Some(ref proxy_url) = args.proxy { if let Some(ref proxy_url) = args.proxy {
let proxy_arg = format!("--proxy-server={}", proxy_url); let proxy_arg = format!("--proxy-server={}", proxy_url);
launch_options launch_options
@@ -73,8 +98,12 @@ fn capture_screenshot(url: &str, index: usize, args: &Args) -> Result<()> {
tab.navigate_to(url).context("Navigation failed")?; tab.navigate_to(url).context("Navigation failed")?;
tab.wait_until_navigated() tab.wait_until_navigated()
.context("Waiting for load failed")?; .context("Waiting for load failed")?;
tab.wait_for_element("body")?;
tab.evaluate("document.readyState === 'complete'", false)?;
let mut file_path = args.output.clone(); let mut file_path = args.output.clone();
file_path.push(format!("screenshot_{:03}.png", index + 1)); let mut log_path = args.output.clone();
log_path.push("log.txt");
file_path.push(format!("{}.png", file_name));
let png_data = tab let png_data = tab
.capture_screenshot( .capture_screenshot(
headless_chrome::protocol::cdp::Page::CaptureScreenshotFormatOption::Png, headless_chrome::protocol::cdp::Page::CaptureScreenshotFormatOption::Png,
@@ -84,8 +113,13 @@ fn capture_screenshot(url: &str, index: usize, args: &Args) -> Result<()> {
) )
.context("Screenshot capture failed")?; .context("Screenshot capture failed")?;
fs::write(&file_path, png_data).context("Failed to write file")?; 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); println!("Successfully captured: {}", url);
fs::write(&log_path, format!("{}\n", url).as_bytes()).context("failed to write log file!")?;
Ok(()) Ok(())
} }