internet_mapper/src/main.rs

extern crate markup5ever_rcdom as rcdom;
extern crate html5ever;

use std::{env, 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, error, info, trace, warn};

mod db;

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt::init();
    debug!("Starting...");

    let url = "https://oliveratkinson.net";
    let budget = "10";
    
    let dom = get(url).await;
    let db = connect().await.expect("Failed to connect to db, aborting.");

    warn!("Walking...");
    walk(&dom, &db, url).await;
}

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) {
    // Should protentailly just check for the record first.
    let created: Option<db::Record> = match db.create("website").content(db::Website { href: String::from("/"), crawled: true, site: site_name.to_string() } ).await {
        Ok(e) => e,
        Err(e) => {
            match e {
                surrealdb::Error::Db(_) => todo!(),
                surrealdb::Error::Api(api) => {
                    match api {
                        surrealdb::error::Api::Query(query) => {
                            error!(query);  
                            None
                        },
                        _ => todo!(),
                    }
                },
            }
        }
    };
    info!{"{:?}", created};

    match &node.data {
        rcdom::NodeData::Document => (),
        rcdom::NodeData::ProcessingInstruction { target, contents } => debug!("ProcessingInstruction"),
        rcdom::NodeData::Doctype { name, public_id, system_id } => debug!("doctype"),
        rcdom::NodeData::Text { contents } => {},
        rcdom::NodeData::Comment { contents } => debug!("comment"),
        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();
                        info!("{:?}", created)
                    } else {
                        // Every not-mailto link
                        let created: Option<db::Record> = db.create("website").content(db::Website { href: attr.value.to_string(), crawled: false, site: site_name.to_string() } ).await.unwrap();
                        info!{"{:?}", created};
                    }
                }
            };
        },
    };

    for child in node.children.borrow().iter() {
        Box::pin(walk(child, db, site_name)).await; 
    }
}