diff --git a/Cargo.toml b/Cargo.toml index e2c2474..af1c835 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ edition = "2024" calamine = "0.36.0" clap = { version = "4.6.1", features = ["derive"] } docx-rust = "0.1.11" +indicatif = "0.18.6" lopdf = "0.44.0" rayon = "1.12.0" reqwest = { version = "0.13.3", features = ["blocking"] } diff --git a/src/main.rs b/src/main.rs index 26bab52..04a8dd9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ use calamine::{Reader, Xlsx}; use clap::Parser; use docx_rust::DocxFile; +use indicatif::ProgressBar; use lopdf::{Document, Object}; use rayon::iter::IntoParallelRefMutIterator; use rayon::{ThreadPoolBuilder, iter::ParallelIterator}; @@ -24,7 +25,11 @@ use urlencoding::decode; struct Args { ///company name #[arg(short, long)] - name: String, + name: Option, + + ///url to search files on + #[arg(short, long)] + url: Option, ///serpapi key #[arg(short, long)] @@ -69,16 +74,13 @@ struct FoundFile { impl FoundFile { fn download(&mut self) { - println!("downloading {}...", self.filename); if let Ok(res) = get(&self.url) { if res.status().is_success() { if let Ok(data) = res.bytes() { if let Ok(mut local_file) = File::create(&self.filepath) { if let Err(e) = local_file.write(&data) { - eprintln!("Error downloading file! {e}"); } else { self.downloaded = Some(true); - println!("{} Downloaded!", self.filename); return; } } @@ -89,7 +91,6 @@ impl FoundFile { } fn parse(&mut self) -> Result<(), Box> { - println!("parsing {}", self.filename); match self.filetype { FileType::Pdf => { let doc = Document::load(&self.filepath)?; @@ -109,7 +110,6 @@ impl FoundFile { FileType::Xlsx => {} } self.parsed = Some(true); - println!("{} parsed!", self.filename); return Ok(()); } } @@ -146,48 +146,95 @@ async fn main() { settings.insert("engine".to_string(), "google".to_string()); let client = Client::new(settings).unwrap(); let mut params = HashMap::new(); - params.insert("q".to_string(), format!("Filetype: PDF \"{}\"", args.name)); + match &args.url { + None => { + if let Some(name) = args.name { + params.insert("q".to_string(), format!("Filetype: PDF \"{}\"", name)); + } else { + eprintln!("error no url or name given!"); + } + } + Some(url) => { + params.insert("q".to_string(), format!("Site: {} Filetype: PDF", url)); + } + } params.insert("h1".to_string(), "en".to_string()); params.insert("g1".to_string(), "us".to_string()); params.insert("google_domain".to_string(), "google.com".to_string()); params.insert("start".to_string(), "0".to_string()); let mut files = Vec::new(); - println!("beginning search"); + let search_pb = ProgressBar::new(pages.len() as u64); + search_pb.println("starting search..."); for page in pages { params.insert("start".to_string(), page.to_string()); if let Ok(results) = client.search(params.clone()).await { + search_pb.inc(1); if let Some(arr) = results["organic_results"].as_array() { for item in arr { if let Some(link) = item["link"].as_str() { - if link.contains(".pdf") { - let link_vec: Vec<&str> = link.split("/").collect(); - let link = link.to_string(); - if let Some(filename) = link_vec.last() { - let encoded_filename = filename.to_string(); - let filename = decode(&encoded_filename).unwrap().to_string(); - let mut filepath = download_path.clone(); - filepath.push(&filename); - let new_file = FoundFile { - filename: filename, - filepath, - url: link, - downloaded: None, - parsed: None, - metadata: HashMap::new(), - filetype: FileType::Pdf, - }; - files.push(new_file); + match &args.url { + None => { + if link.contains(".pdf") { + let link_vec: Vec<&str> = link.split("/").collect(); + let link = link.to_string(); + if let Some(filename) = link_vec.last() { + let encoded_filename = filename.to_string(); + let filename = + decode(&encoded_filename).unwrap().to_string(); + let mut filepath = download_path.clone(); + filepath.push(&filename); + let new_file = FoundFile { + filename: filename, + filepath, + url: link, + downloaded: None, + parsed: None, + metadata: HashMap::new(), + filetype: FileType::Pdf, + }; + files.push(new_file); + } + } + } + Some(url) => { + if link.contains(url) { + if link.contains(".pdf") { + let link_vec: Vec<&str> = link.split("/").collect(); + let link = link.to_string(); + if let Some(filename) = link_vec.last() { + let encoded_filename = filename.to_string(); + let filename = + decode(&encoded_filename).unwrap().to_string(); + let mut filepath = download_path.clone(); + filepath.push(&filename); + let new_file = FoundFile { + filename: filename, + filepath, + url: link, + downloaded: None, + parsed: None, + metadata: HashMap::new(), + filetype: FileType::Pdf, + }; + files.push(new_file); + } + } + } } } } } } + } else { + break; } } - println!("Search finished!"); - println!("downloding files..."); + search_pb.finish_and_clear(); + let download_pb = ProgressBar::new(files.len() as u64); + download_pb.println("downloading files..."); files.par_iter_mut().for_each(|file| { file.download(); + download_pb.inc(1); if file.downloaded == Some(true) { match file.parse() { Ok(_) => {}