made the healdess chromium window wait until the body and javascript
loaded before taking the screenshot.
This commit is contained in:
Generated
+1625
File diff suppressed because it is too large
Load Diff
+49
-15
@@ -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,16 +24,38 @@ 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();
|
||||||
|
if resume.is_some() {
|
||||||
|
let url_resume = resume.clone().unwrap();
|
||||||
|
let url_text = read_to_string(&args.input)?;
|
||||||
|
let mut resume_met = false;
|
||||||
|
for line in url_text.lines() {
|
||||||
|
if !resume_met {
|
||||||
|
if line.trim().to_string() == url_resume {
|
||||||
|
resume_met = true;
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
urls.push(line.trim().to_string());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if resume.is_none() {
|
||||||
let file = File::open(&args.input).context("Failed to open input file")?;
|
let file = File::open(&args.input).context("Failed to open input file")?;
|
||||||
let urls: Vec<String> = BufReader::new(file)
|
let urls: Vec<String> = BufReader::new(file)
|
||||||
.lines()
|
.lines()
|
||||||
@@ -41,26 +63,29 @@ fn main() -> Result<()> {
|
|||||||
.map(|s| s.trim().to_string())
|
.map(|s| s.trim().to_string())
|
||||||
.filter(|s| !s.is_empty())
|
.filter(|s| !s.is_empty())
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
println!(
|
println!(
|
||||||
"Starting capture of {} URLs using {} workers...",
|
"Starting capture of {} URLs using {} workers...",
|
||||||
urls.len(),
|
urls.len(),
|
||||||
args.workers
|
args.workers
|
||||||
);
|
);
|
||||||
|
urls.par_iter().enumerate().for_each(|(_index, url)| {
|
||||||
urls.par_iter().enumerate().for_each(|(index, url)| {
|
if let Err(e) = capture_screenshot(url, &args) {
|
||||||
if let Err(e) = capture_screenshot(url, index, &args) {
|
|
||||||
eprintln!("Failed to process {}: {:?}", url, e);
|
eprintln!("Failed to process {}: {:?}", url, e);
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
println!("All tasks completed.");
|
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(())
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user