use anyhow::{Context, Result};
use reqwest::header::{HeaderMap, HeaderValue, USER_AGENT};
use scraper::{Html, Selector};
use serde::{Deserialize, Serialize};
use std::path::Path;
use std::sync::LazyLock;
use std::time::Duration;
static OG_TITLE_SEL: LazyLock<Selector> =
LazyLock::new(|| Selector::parse("meta[property='og:title']").unwrap());
static TITLE_SEL: LazyLock<Selector> = LazyLock::new(|| Selector::parse("title").unwrap());
static H1_SEL: LazyLock<Selector> = LazyLock::new(|| Selector::parse("h1").unwrap());
static MAIN_SELECTORS: LazyLock<Vec<Selector>> = LazyLock::new(|| {
[
"article",
"main",
"[role='main']",
".mw-parser-output",
"#mw-content-text",
"#content",
".post-content",
".article-content",
".entry-content",
"body",
]
.into_iter()
.filter_map(|s| Selector::parse(s).ok())
.collect()
});
const REFERENCE_HEADINGS: &[&str] = &[
"references",
"## references",
"=== references ===",
"== references ==",
"external links",
"## external links",
"=== external links ===",
"notes and references",
"## notes",
"## bibliography",
"see also",
"## see also",
];
pub const MAX_SCRAPE_BYTES: usize = 10 * 1024 * 1024;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScrapedPage {
pub url: String,
pub title: Option<String>,
pub content: String,
pub word_count: usize,
}
#[derive(Clone)]
pub struct WebScraper {
client: reqwest::Client,
}
impl WebScraper {
pub fn new() -> Result<Self> {
let mut headers = HeaderMap::new();
headers.insert(
USER_AGENT,
HeaderValue::from_static(
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) \
AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
),
);
headers.insert(
reqwest::header::ACCEPT,
HeaderValue::from_static(
"text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8",
),
);
headers.insert(
reqwest::header::ACCEPT_LANGUAGE,
HeaderValue::from_static("en-US,en;q=0.9"),
);
headers.insert(
reqwest::header::HeaderName::from_static("sec-ch-ua"),
HeaderValue::from_static(
"\"Chromium\";v=\"124\", \"Google Chrome\";v=\"124\", \"Not-A.Brand\";v=\"99\"",
),
);
headers.insert(
reqwest::header::HeaderName::from_static("sec-ch-ua-mobile"),
HeaderValue::from_static("?0"),
);
headers.insert(
reqwest::header::HeaderName::from_static("sec-ch-ua-platform"),
HeaderValue::from_static("\"macOS\""),
);
headers.insert(
reqwest::header::HeaderName::from_static("sec-fetch-dest"),
HeaderValue::from_static("document"),
);
headers.insert(
reqwest::header::HeaderName::from_static("sec-fetch-mode"),
HeaderValue::from_static("navigate"),
);
headers.insert(
reqwest::header::HeaderName::from_static("sec-fetch-site"),
HeaderValue::from_static("none"),
);
headers.insert(
reqwest::header::HeaderName::from_static("sec-fetch-user"),
HeaderValue::from_static("?1"),
);
headers.insert(
reqwest::header::HeaderName::from_static("upgrade-insecure-requests"),
HeaderValue::from_static("1"),
);
let client = reqwest::Client::builder()
.default_headers(headers)
.connect_timeout(Duration::from_secs(5))
.timeout(Duration::from_secs(20))
.redirect(reqwest::redirect::Policy::limited(10))
.build()
.context("Failed to initialize HTTP client")?;
Ok(Self { client })
}
pub async fn scrape(&self, url: &str) -> Result<ScrapedPage> {
let mut response = self
.client
.get(url)
.send()
.await
.with_context(|| format!("Failed to fetch URL: {}", url))?;
if !response.status().is_success() {
anyhow::bail!(
"Failed to fetch webpage (HTTP status: {})",
response.status()
);
}
if let Some(content_len) = response.content_length()
&& content_len as usize > MAX_SCRAPE_BYTES
{
anyhow::bail!(
"Webpage content length ({} bytes) exceeds security limit of 10 MB",
content_len
);
}
let mut bytes = Vec::new();
while let Some(chunk) = response.chunk().await.context("Failed to read response body")? {
if bytes.len() + chunk.len() > MAX_SCRAPE_BYTES {
anyhow::bail!(
"Webpage response body exceeded security limit of 10 MB"
);
}
bytes.extend_from_slice(&chunk);
}
let html = String::from_utf8_lossy(&bytes);
self.normalize(url, &html)
}
pub fn normalize(&self, url: &str, raw_html: &str) -> Result<ScrapedPage> {
let document = Html::parse_document(raw_html);
let title = self.extract_title(&document);
let target_html = self.extract_main_html(&document);
let rendered_text = html2text::config::plain()
.link_footnotes(false)
.no_table_borders()
.string_from_read(target_html.as_bytes(), 10_000)
.context("Failed to convert HTML to text")?;
let (cleaned_content, word_count) = clean_text(&rendered_text, true, 2);
Ok(ScrapedPage {
url: url.to_string(),
title,
content: cleaned_content,
word_count,
})
}
fn extract_title(&self, document: &Html) -> Option<String> {
if let Some(elem) = document.select(&OG_TITLE_SEL).next()
&& let Some(content) = elem.value().attr("content")
{
let trimmed = content.trim();
if !trimmed.is_empty() {
return Some(trimmed.to_string());
}
}
for sel in [&*TITLE_SEL, &*H1_SEL] {
if let Some(elem) = document.select(sel).next() {
let text: String = elem.text().collect();
let trimmed = text.trim();
if !trimmed.is_empty() {
return Some(trimmed.to_string());
}
}
}
None
}
fn extract_main_html(&self, document: &Html) -> String {
for sel in MAIN_SELECTORS.iter() {
if let Some(element) = document.select(sel).next() {
let html = element.html();
if html.len() > 200 {
return html;
}
}
}
document.html()
}
}
fn is_reference_heading(line: &str) -> bool {
REFERENCE_HEADINGS
.iter()
.any(|heading| line.eq_ignore_ascii_case(heading))
}
pub fn is_pdf_path(path: &Path) -> bool {
path.extension().is_some_and(|e| e.eq_ignore_ascii_case("pdf"))
}
pub fn is_epub_path(path: &Path) -> bool {
path.extension().is_some_and(|e| e.eq_ignore_ascii_case("epub"))
}
#[derive(Debug, Clone)]
pub struct ExtractedDocument {
pub title: String,
pub author: Option<String>,
pub content: String,
pub word_count: usize,
}
pub fn extract_pdf(path: &Path) -> Result<ExtractedDocument> {
let text = pdf_extract::extract_text(path)
.with_context(|| format!("Failed to extract text from PDF file: {:?}", path))?;
let title = path
.file_stem()
.map(|s| s.to_string_lossy().to_string())
.unwrap_or_else(|| "PDF Document".to_string());
let (content, word_count) = clean_raw_text(&text);
if content.trim().is_empty() {
anyhow::bail!(
"PDF document {:?} contained no extractable text (it may be scanned/image-only).",
path
);
}
Ok(ExtractedDocument {
title,
author: None,
content,
word_count,
})
}
pub fn extract_epub(path: &Path) -> Result<ExtractedDocument> {
let mut doc = epub::doc::EpubDoc::new(path)
.map_err(|e| anyhow::anyhow!("Failed to open EPUB file: {:?}", e))?;
let title = doc
.mdata("title")
.map(|v| v.value.clone())
.or_else(|| {
path.file_stem()
.map(|s| s.to_string_lossy().to_string())
})
.unwrap_or_else(|| "EPUB Book".to_string());
let author = doc.mdata("creator").map(|v| v.value.clone());
let mut full_text = String::new();
let num_chapters = doc.get_num_chapters();
for _ in 0..num_chapters {
if let Some((page_html, _mime)) = doc.get_current_str() {
let rendered = html2text::config::plain()
.link_footnotes(false)
.no_table_borders()
.string_from_read(page_html.as_bytes(), 10_000)
.unwrap_or_default();
if !rendered.trim().is_empty() {
if !full_text.is_empty() {
full_text.push_str("\n\n");
}
full_text.push_str(rendered.trim());
}
}
if !doc.go_next() {
break;
}
}
let (content, word_count) = clean_raw_text(&full_text);
if content.trim().is_empty() {
anyhow::bail!("EPUB file {:?} contained no extractable text.", path);
}
Ok(ExtractedDocument {
title,
author,
content,
word_count,
})
}
pub fn clean_raw_text(text: &str) -> (String, usize) {
clean_text(text, false, 1)
}
fn clean_text(text: &str, strip_references: bool, max_empty_lines: usize) -> (String, usize) {
let mut out = String::with_capacity(text.len());
let mut word_count = 0;
let mut empty_line_count = 0;
for raw_line in text.lines() {
let line = raw_line.trim();
if strip_references && is_reference_heading(line) {
break;
}
if line.is_empty() {
empty_line_count += 1;
if empty_line_count <= max_empty_lines && !out.is_empty() {
out.push('\n');
}
} else {
empty_line_count = 0;
word_count += line.split_whitespace().count();
if !out.is_empty() {
out.push('\n');
}
out.push_str(line);
}
}
let trimmed_len = out.trim_end().len();
out.truncate(trimmed_len);
(out, word_count)
}