internet_mapper/src/main.rs

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<Client>) -> Option<Rc<Node>> {
    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<Client> , 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::Record> = 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<Client>) -> Vec<Website> {
    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")
}