use std::env; use std::fs::{self, File}; use std::io::{BufRead, BufReader, BufWriter, Read, Seek, Write}; use std::path::Path; use std::sync::{Arc, Mutex}; use std::thread; use std::time::Duration; use regex::Regex; use warc::{Record, RecordType, WarcHeader, WarcReader, RecordBuilder}; use flate2::read::GzDecoder; use zstd::stream::read::Decoder; use rayon::ThreadPoolBuilder; use rayon::prelude::*; fn main() -> Result<(), Box> { let args: Vec = env::args().collect(); if args.len() < 3 { eprintln!("Usage: "); std::process::exit(1); } let warc_directory = &args[1]; let output_file_path = &args[2]; let regex_pattern = r"\S*imgur\S*"; let pattern = Regex::new(regex_pattern)?; loop { process_warcs_in_directory(warc_directory, output_file_path, &pattern)?; thread::sleep(Duration::from_secs(60)); } } fn read_skippable_frame(reader: &mut R) -> Result>, Box> { let mut magic_number = [0; 4]; reader.read_exact(&mut magic_number)?; if magic_number != [0x28, 0xB5, 0x2F, 0xFD] { return Ok(None); } let mut frame_size_bytes = [0; 4]; reader.read_exact(&mut frame_size_bytes)?; let frame_size = u32::from_le_bytes(frame_size_bytes) as usize; let mut frame_data = vec![0; frame_size]; reader.read_exact(&mut frame_data)?; Ok(Some(frame_data)) } fn process_warcs_in_directory( warc_directory: &str, output_file_path: &str, pattern: &Regex, ) -> Result<(), Box> { let output_file = Arc::new(Mutex::new(BufWriter::new(File::create(output_file_path)?))); let paths: Vec<_> = fs::read_dir(warc_directory)? .filter_map(|entry| { let entry = entry.ok()?; let path = entry.path(); let file_name = path.file_name()?.to_str()?.to_string(); if file_name.ends_with(".warc.gz") || file_name.ends_with(".warc.zst") { Some(path) } else { None } }) .collect(); let thread_pool = ThreadPoolBuilder::new() .num_threads(12) .build() .unwrap(); thread_pool.install(|| { paths.par_iter().for_each(|path| { if let Err(e) = process_warc(path, &output_file, &pattern) { eprintln!("Error processing {}: {}", path.display(), e); } }); }); Ok(()) } fn process_warc( path: &Path, output_file: &Arc>>, pattern: &Regex, ) -> Result<(), Box> { let file = File::open(path)?; let mut warc_stream: Box = if path.extension().unwrap() == "gz" { Box::new(GzDecoder::new(file)) } else { let mut file = BufReader::new(file); if let Some(dict_data) = read_skippable_frame(&mut file)? { Box::new(Decoder::with_dictionary(file, &dict_data)?) } else { Box::new(file) } }; let warc_reader = warc::WarcReader::new(BufReader::new(warc_stream)); let mut record_builder = RecordBuilder::default(); for result in warc_reader.iter_records() { match result { Ok((record_header, body)) => { let record = record_builder.build(); let is_response = record.warc_type() == &warc::RecordType::from("response"); if is_response { let content = String::from_utf8_lossy(&raw_body); if pattern.is_match(&content) { let mut output_file = output_file.lock().unwrap(); writeln!(output_file, "{}", content)?; } } } Err(e) => eprintln!("Error reading record: {}", e), } } Ok(()) }