Add src/main.rs

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<String>,
+
+    #[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<String>),
+    UrlSet(Vec<String>),
+}
+
+#[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<Url> = 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<String> = HashSet::new();
+    let mut extracted_urls: HashSet<String> = 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::<Vec<_>>()
+        .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<String> = 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::<Vec<_>>();
+
+    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<Url> {
+    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<String> {
+    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")?)
+}