internet_mapper/src/main.rs
Oliver Atkinson 3a46dd937b updates
2024-10-31 15:09:48 -06:00

90 lines
2.7 KiB
Rust

extern crate markup5ever_rcdom as rcdom;
extern crate html5ever;
use std::rc::Rc;
use db::connect;
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, warn};
mod db;
#[tokio::main]
async fn main() {
tracing_subscriber::fmt::init();
debug!("Starting...");
let url = "https://oliveratkinson.net";
let db = connect().await.expect("Failed to connect to db, aborting.");
let dom = get(url).await;
walk(&dom, &db, url).await;
info!("Done");
}
async fn get(url: &str) -> Rc<Node> {
let response = reqwest::get(url).await.unwrap();
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();
dom.document
}
async fn walk(node: &rcdom::Handle, db: &Surreal<Client> , site_name: &str) {
// Insert Or Update
let _created: Vec<db::Record> = match db.upsert("website").content(db::Website { href: String::from("/"), crawled: true, site: site_name.to_string() } ).await {
Ok(e) => {
// Return this for type coercion
e
},
Err(e) => {
unimplemented!("{}", e);
}
};
match &node.data {
rcdom::NodeData::Element { name, attrs, template_contents, mathml_annotation_xml_integration_point } => {
for attr in attrs.borrow().clone() {
let name = name.local.to_string();
if name == "a" {
if attr.value.starts_with("mailto") {
// mailto link, lol
let created: Option<db::Record> = db.create("email").content(db::Email {
email: attr.value.to_string()
}).await.unwrap();
warn!("{:?}", created)
} else {
let href = attr.value.to_string();
info!("{}", href);
// Every not-mailto link
let _created: Option<db::Record> = db.create("website").content(db::Website {
href,
crawled: false,
site: site_name.to_string()
}).await.unwrap();
}
}
};
},
_ => {},
};
for child in node.children.borrow().iter() {
Box::pin(walk(child, db, site_name)).await;
}
}