diff --git a/Cargo.lock b/Cargo.lock index 92a59e6..f35af97 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -267,6 +267,12 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "clap" version = "4.5.20" @@ -357,6 +363,16 @@ version = "0.8.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" +[[package]] +name = "ctrlc" +version = "3.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90eeab0aa92f3f9b4e87f258c72b139c207d251f9cbc1080a0086b86a8870dd3" +dependencies = [ + "nix", + "windows-sys 0.59.0", +] + [[package]] name = "encode_unicode" version = "0.3.6" @@ -892,6 +908,18 @@ dependencies = [ "tempfile", ] +[[package]] +name = "nix" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "number_prefix" version = "0.4.0" @@ -1043,6 +1071,7 @@ version = "1.0.0" dependencies = [ "async-std", "clap", + "ctrlc", "futures", "indicatif", "regex", diff --git a/Cargo.toml b/Cargo.toml index fe13618..bae06eb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ edition = "2021" [dependencies] async-std = "1" clap = { version = "4", features = ["derive"] } +ctrlc = "3" futures = "0" indicatif = "0" regex = "1" diff --git a/src/main.rs b/src/main.rs index 1e63617..397a3e0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,6 +10,8 @@ use tokio::time::{sleep, Duration}; use std::io::Write; use std::process::ExitCode; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; const BAR_LENGTH: u64 = 8; @@ -20,12 +22,23 @@ async fn main() -> ExitCode { let uri_tags = &args.tags.join("+"); let _ = std::fs::create_dir(uri_tags); + let running = Arc::new(AtomicBool::new(true)); + let running_t = running.clone(); + ctrlc::set_handler(move || { + running_t.store(false, Ordering::SeqCst); + }) + .unwrap(); + let client = Client::builder() .user_agent(&args.user_agent) .build() .unwrap(); for page in args.page - 1.. { + if !running.load(Ordering::SeqCst) { + return ExitCode::FAILURE; + } + println!("now scraping page {} (https://rule34.xxx/index.php?page=post&s=list&tags={uri_tags}&pid={})", page + 1, page * 42); let post_html = async |client: &Client| { @@ -51,6 +64,10 @@ async fn main() -> ExitCode { println!("no urls found, retrying in 5 seconds..."); sleep(Duration::from_millis(5000)).await; + if !running.load(Ordering::SeqCst) { + return ExitCode::FAILURE; + } + urls = post_html(&client).await; if !urls.is_empty() { println!("urls found! continuing..."); @@ -71,12 +88,17 @@ async fn main() -> ExitCode { .map(|(i, url)| { let i = i + 1; let client = &client; + let running_t = running.clone(); let this_bar = indicatif::ProgressBar::new(BAR_LENGTH); this_bar.set_style(indicatif::ProgressStyle::with_template("[{bar}] {msg}").unwrap().progress_chars("=> ")); let this_bar = multi_prog.insert(i, this_bar); async move { // "thread" loop { + if !running_t.load(Ordering::SeqCst) { + return; + } + this_bar.set_message(format!("\x1b[37m[{i: >4}/{urls_amount}] \x1b[36mscraping {url}\x1b[0m")); let resp = client.get(&url).send().await.unwrap(); if let Ok(img_url) = extract_img_url(&resp.text().await.unwrap()) { @@ -85,7 +107,7 @@ async fn main() -> ExitCode { "\x1b[37m[{i: >4}/{urls_amount}] \x1b[1;31mimage url not found\x1b[0m" )); } else { - download_file(&img_url, this_bar, i, urls_amount, uri_tags).await; + download_file(running_t, &img_url, this_bar, i, urls_amount, uri_tags).await; } break; } @@ -129,6 +151,7 @@ fn extract_img_url(html: &str) -> Result { } async fn download_file( + running: Arc, img_url: &str, this_bar: ProgressBar, i: usize, @@ -151,7 +174,7 @@ async fn download_file( )); return; } else { - std::fs::File::create(file_path).unwrap() + std::fs::File::create(&file_path).unwrap() }; let mut res = Client::new() @@ -163,6 +186,14 @@ async fn download_file( let file_length = res.content_length().unwrap(); let mut written = 0; while let Some(img_chunk) = res.chunk().await.unwrap() { + if !running.load(Ordering::SeqCst) { + this_bar.abandon_with_message(format!( + "\x1b[37m[{i: >4}/{urls_amount}] \x1b[33mcancelling {img_url}\x1b[0m" + )); + drop(file); + std::fs::remove_file(&file_path).unwrap(); + return; + } file.write_all(&img_chunk).unwrap(); written += img_chunk.len(); this_bar.set_position((written as f64 / file_length as f64 * BAR_LENGTH as f64) as u64);