extern crate markup5ever_rcdom as rcdom; extern crate html5ever; use std::rc::Rc; use db::{connect, Website}; use html5ever::{parse_document, tendril::TendrilSink, tree_builder::TreeBuilderOpts, ParseOpts}; use rcdom::{Node, RcDom}; use surrealdb::{engine::remote::ws::Client, Surreal}; use tracing::{debug, info, instrument}; use tracing_subscriber::EnvFilter; mod db; #[tokio::main] async fn main() { tracing_subscriber::fmt() .with_env_filter(EnvFilter::from_default_env()) .with_line_number(true) .without_time() .init(); debug!("Starting..."); // Would probably take these in as parameters from a cli let url = "https://oliveratkinson.net/"; let budget = 50; let mut crawled = 0; let db = connect().await.expect("Failed to connect to db, aborting."); // Kick off the whole machine - This Website object doesn't matter, it's just to allow for // get() to work. let mut site = Website::from(url); let dom = get(&mut site, &db).await.expect("Inital page returned None."); crawled += 1; walk(&dom, &db, &site).await; while crawled < budget { let uncrawled = get_uncrawled_links(&db).await; debug!("Crawling {} pages...", uncrawled.len()); for mut site in uncrawled { if let Some(dom) = get(&mut site, &db).await { walk(&dom, &db, &site).await; crawled += 1; let percent = format!("{:.2}%", (crawled as f32/budget as f32) * 100f32); info!("Crawled {crawled} out of {budget} pages. ({percent})"); } } } info!("Done"); } #[instrument(skip_all)] /// A quick helper function for downloading a url async fn get(site: &mut Website, db: &Surreal) -> Option> { if let Ok(response) = reqwest::get(site.href_str()).await { let data = response.text().await.unwrap(); let opts = ParseOpts { tree_builder: TreeBuilderOpts { drop_doctype: true, ..Default::default() }, ..Default::default() }; let dom = parse_document(RcDom::default(), opts) .from_utf8() .read_from(&mut data.as_bytes()) .unwrap(); site.crawled(); site.store(db).await; return Some(dom.document); } None } /// Walks the givin site, placing it's findings in the database async fn walk(node: &rcdom::Handle, db: &Surreal , site_name: &Website) { // Insert Or Update // create_root(site_name, db).await; match &node.data { rcdom::NodeData::Element { name, attrs, template_contents, mathml_annotation_xml_integration_point } => { for attr in attrs.borrow().clone() { if name.local.to_string() == "a" { if attr.value.starts_with("mailto") { // mailto link, lol let _created: Option = db.create("email").content(db::Email { email: attr.value.to_string() }).await.unwrap(); } else { let mut web = Website::new(&site_name.site(), &attr.value, false); web.store(db).await; } } }; }, _ => {}, }; for child in node.children.borrow().iter() { Box::pin(walk(child, db, site_name)).await; } } /// Returns 0-50 uncrawled links (LIMIT = 50) async fn get_uncrawled_links(db: &Surreal) -> Vec { let mut response = db.query("SELECT * FROM website WHERE crawled = false LIMIT 50").await.expect("Hard-coded query failed..?"); response.take(0).expect("Returned websites couldn't be parsed") }