use std::sync::Arc; use color_eyre::eyre::{OptionExt as _, WrapErr as _}; use foldhash::HashSetExt as _; #[cfg(feature = "tui")] use crate::event::{AppEvent, Event}; use crate::{ HashSet, config::{Config, Source}, parsers::PROXY_REGEX, proxy::{Proxy, ProxyType}, utils::pretty_error, }; async fn scrape_one( config: Arc, http_client: reqwest_middleware::ClientWithMiddleware, proto: ProxyType, proxies: Arc>>, source: Arc, #[cfg(feature = "tui")] tx: tokio::sync::mpsc::UnboundedSender, ) -> crate::Result<()> { let text_result = if let Ok(u) = url::Url::parse(&source.url) { match u.scheme() { "http" | "https" => { let mut request = http_client.get(u); drop(http_client); if let Some(auth) = &source.basic_auth { request = request .basic_auth(&auth.username, auth.password.as_ref()); } if let Some(headers) = &source.headers { for (k, v) in headers { request = request.header(k, v); } } match request.send().await { Ok(resp) => resp.text().await.map_err(Into::into), Err(e) => Err(e.into()), } } _ => { drop(http_client); match u.to_file_path() { Ok(path) => tokio::fs::read_to_string(path) .await .wrap_err_with(move || { format!("failed to read file to string: {u}") }), Err(()) => tokio::fs::read_to_string(&source.url) .await .wrap_err_with(move || { format!("failed to read file to string: {u}") }), } } } } else { drop(http_client); tokio::fs::read_to_string(&source.url).await.wrap_err_with(|| { format!("failed to read file to string: {}", source.url) }) }; #[cfg(feature = "tui")] drop(tx.send(Event::App(AppEvent::SourceScraped(proto)))); let text = match text_result { Ok(text) => text, Err(e) => { tracing::warn!("{}: {}", source.url, pretty_error(&e)); return Ok(()); } }; #[cfg(feature = "tui")] let mut seen_protocols = HashSet::new(); let mut new_proxies = HashSet::new(); for maybe_capture in PROXY_REGEX.captures_iter(&text) { if config.scraping.max_proxies_per_source != 0 && new_proxies.len() >= config.scraping.max_proxies_per_source { tracing::warn!( "{}: too many proxies (> {}) - skipped", source.url, config.scraping.max_proxies_per_source ); return Ok(()); } let capture = maybe_capture?; let protocol = match capture.name("protocol") { Some(m) => m.as_str().parse()?, None => proto, }; if config.protocol_is_enabled(protocol) { #[cfg(feature = "tui")] seen_protocols.insert(protocol); new_proxies.insert(Proxy { protocol, host: capture .name("host") .ok_or_eyre("failed to match \"host\" regex capture group")? .as_str() .to_owned(), port: capture .name("port") .ok_or_eyre("failed to match \"port\" regex capture group")? .as_str() .parse()?, username: capture .name("username") .map(|m| m.as_str().to_owned()), password: capture .name("password") .map(|m| m.as_str().to_owned()), timeout: None, exit_ip: None, }); } } drop(config); drop(text); if new_proxies.is_empty() { tracing::warn!("{}: no proxies found", source.url); return Ok(()); } drop(source); let mut proxies = proxies.lock(); proxies.extend(new_proxies); #[cfg(feature = "tui")] for proto in seen_protocols { let count = proxies.iter().filter(move |p| p.protocol == proto).count(); drop(tx.send(Event::App(AppEvent::TotalProxies(proto, count)))); } drop(proxies); Ok(()) } pub async fn scrape_all( config: Arc, http_client: reqwest_middleware::ClientWithMiddleware, token: tokio_util::sync::CancellationToken, #[cfg(feature = "tui")] tx: tokio::sync::mpsc::UnboundedSender, ) -> crate::Result> { let proxies = Arc::new(parking_lot::Mutex::new(HashSet::new())); let mut join_set = tokio::task::JoinSet::new(); for (&proto, sources) in &config.scraping.sources { #[cfg(feature = "tui")] drop(tx.send(Event::App(AppEvent::SourcesTotal(proto, sources.len())))); for source in sources { let config = Arc::clone(&config); let http_client = http_client.clone(); let proxies = Arc::clone(&proxies); let token = token.clone(); let source = Arc::clone(source); #[cfg(feature = "tui")] let tx = tx.clone(); join_set.spawn(async move { tokio::select! { biased; res = scrape_one( config, http_client, proto, proxies, source, #[cfg(feature = "tui")] tx, ) => res, () = token.cancelled() => Ok(()), } }); } } drop(config); drop(http_client); drop(token); #[cfg(feature = "tui")] drop(tx); while let Some(res) = join_set.join_next().await { res??; } drop(join_set); Ok(Arc::into_inner(proxies) .ok_or_eyre("failed to unwrap Arc")? .into_inner() .into_iter() .collect()) }