diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..297f850 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,201 @@ +use std::{ + collections::{HashSet, VecDeque}, + fs, + io::Read, + sync::Arc, + time::Duration, +}; + +use anyhow::{bail, Context, Result}; +use clap::{Parser, ValueEnum}; +use flate2::read::GzDecoder; +use futures::stream::{self, StreamExt}; +use mimalloc::MiMalloc; +use reqwest::{header, Client}; +use roxmltree::Document; +use tokio::time::timeout; +use url::Url; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +#[derive(Parser, Debug)] +#[command(author, version, about)] +struct Args { + #[arg(long = "sitemap", required = true)] + sitemaps: Vec, + + #[arg(long, default_value = "urls.txt")] + output: String, + + #[arg(long, value_enum, default_value_t = SchemeOpt::Https)] + scheme: SchemeOpt, + + #[arg(long, default_value_t = 8)] + workers: usize, + + #[arg(long, default_value = "riptide 0.1 https://codeberg.org/gigirassy/riptide/")] + user_agent: String, + + #[arg(long, default_value_t = 30)] + timeout_secs: u64, +} + +#[derive(Copy, Clone, Debug, ValueEnum)] +enum SchemeOpt { + Http, + Https, +} + +impl SchemeOpt { + fn as_str(self) -> &'static str { + match self { + SchemeOpt::Http => "http", + SchemeOpt::Https => "https", + } + } +} + +enum ParsedSitemap { + Index(Vec), + UrlSet(Vec), +} + +#[tokio::main] +async fn main() -> Result<()> { + let args = Args::parse(); + + let client = Client::builder() + .user_agent(args.user_agent) + .timeout(Duration::from_secs(args.timeout_secs)) + .redirect(reqwest::redirect::Policy::limited(10)) + .build() + .context("failed to build HTTP client")?; + + let mut queue: VecDeque = VecDeque::new(); + for s in &args.sitemaps { + queue.push_back(Url::parse(s).with_context(|| format!("invalid sitemap URL: {s}"))?); + } + + let mut visited_sitemaps: HashSet = HashSet::new(); + let mut extracted_urls: HashSet = HashSet::new(); + let workers = args.workers.max(1); + + while !queue.is_empty() { + let mut batch = Vec::new(); + + while batch.len() < workers { + let Some(next) = queue.pop_front() else { + break; + }; + + let key = next.as_str().to_string(); + if visited_sitemaps.insert(key) { + batch.push(next); + } + } + + if batch.is_empty() { + continue; + } + + let results = stream::iter(batch.into_iter().map(|sitemap_url| { + let client = client.clone(); + async move { fetch_and_parse(&client, sitemap_url).await } + })) + .buffer_unordered(workers) + .collect::>() + .await; + + for result in results { + let (parsed, base_url) = result?; + + match parsed { + ParsedSitemap::Index(children) => { + for child in children { + if let Ok(url) = resolve_url(&base_url, &child) { + queue.push_back(url); + } + } + } + ParsedSitemap::UrlSet(urls) => { + for raw in urls { + if let Ok(url) = resolve_url(&base_url, &raw) { + let normalized = force_scheme(url, args.scheme); + extracted_urls.insert(normalized); + } + } + } + } + } + } + + let mut out: Vec = extracted_urls.into_iter().collect(); + out.sort_unstable(); + + fs::write(&args.output, out.join("\n") + "\n") + .with_context(|| format!("failed to write {}", args.output))?; + + Ok(()) +} + +async fn fetch_and_parse(client: &Client, sitemap_url: Url) -> Result<(ParsedSitemap, Url)> { + let resp = timeout(Duration::from_secs(30), client.get(sitemap_url.clone()).send()) + .await + .context("request timed out")? + .context("request failed")? + .error_for_status() + .context("server returned an error status")?; + + let final_url = resp.url().clone(); + let bytes = resp.bytes().await.context("failed to read response body")?; + + let xml = decode_sitemap_body(&bytes, &final_url)?; + let doc = Document::parse(&xml).context("failed to parse sitemap XML")?; + let root = doc.root_element().tag_name().name(); + + let locs = doc + .descendants() + .filter(|n| n.has_tag_name("loc")) + .filter_map(|n| n.text()) + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .collect::>(); + + let parsed = match root { + "sitemapindex" => ParsedSitemap::Index(locs), + "urlset" => ParsedSitemap::UrlSet(locs), + other => bail!("unexpected root element: {other}"), + }; + + Ok((parsed, final_url)) +} + +fn resolve_url(base: &Url, raw: &str) -> Result { + if let Ok(url) = Url::parse(raw) { + return Ok(url); + } + Ok(base.join(raw)?) +} + +fn force_scheme(mut url: Url, scheme: SchemeOpt) -> String { + let _ = url.set_scheme(scheme.as_str()); + url.to_string() +} + +fn decode_sitemap_body(bytes: &[u8], url: &Url) -> Result { + if let Ok(s) = std::str::from_utf8(bytes) { + return Ok(s.to_string()); + } + + if url.path().ends_with(".gz") || bytes.starts_with(&[0x1f, 0x8b]) { + let mut decoder = GzDecoder::new(bytes); + let mut s = String::new(); + decoder + .read_to_string(&mut s) + .context("failed to decompress gzip sitemap")?; + return Ok(s); + } + + Ok(String::from_utf8(bytes.to_vec()).context("sitemap body is not valid UTF-8")?) +} \ No newline at end of file