paperoni/src/extractor.rs

use async_std::fs::File;
use async_std::io::prelude::*;
use async_std::task;
use kuchiki::{traits::*, ElementData, NodeDataRef, NodeRef};
use url::Url;

pub type ResourceInfo = (String, Option<String>);

pub struct Extractor {
    pub root_node: NodeRef,
    pub content: Option<NodeDataRef<ElementData>>,
    pub img_urls: Vec<ResourceInfo>,
}

impl Extractor {
    /// Create a new instance of an HTML extractor given an HTML string
    pub fn from_html(html_str: &str) -> Self {
        Extractor {
            content: None,
            img_urls: Vec::new(),
            root_node: kuchiki::parse_html().one(html_str),
        }
    }

    /// Extract the value of an attribute
    fn extract_attr_val<T: Fn(&str) -> U, U>(
        &self,
        css_selector: &str,
        attr_target: &str,
        mapper: T,
    ) -> Option<U> {
        self.root_node
            .select_first(css_selector)
            .ok()
            .and_then(|data| data.attributes.borrow().get(attr_target).map(mapper))
    }

    /// Extract the text of a DOM node given its CSS selector
    fn extract_inner_text(&self, css_selector: &str) -> Option<String> {
        let node_ref = self.root_node.select_first(css_selector).ok()?;
        extract_text_from_node(node_ref.as_node())
    }

    /// Locates and extracts the HTML in a document which is determined to be
    /// the source of the content
    pub fn extract_content(&mut self) {
        // Extract the useful parts of the head section
        let author: Option<String> =
            self.extract_attr_val("meta[name='author']", "content", |author| {
                author.to_string()
            });

        let description =
            self.extract_attr_val("meta[name='description']", "content", |description| {
                description.to_string()
            });

        let tags = self.extract_attr_val("meta[name='keywords']", "content", |tags| {
            tags.split(",")
                .map(|tag| tag.trim().to_string())
                .collect::<Vec<String>>()
        });

        let title = self.extract_inner_text("title").unwrap_or("".to_string());
        let lang = self
            .extract_attr_val("html", "lang", |lang| lang.to_string())
            .unwrap_or("en".to_string());

        // Extract the article

        let article_ref = self.root_node.select_first("article").unwrap();

        for node_ref in article_ref.as_node().descendants() {
            match node_ref.data() {
                kuchiki::NodeData::Element(..) | kuchiki::NodeData::Text(..) => (),
                _ => node_ref.detach(),
            }
        }
        self.content = Some(article_ref);
    }

    /// Traverses the DOM tree of the content and retrieves the IMG URLs
    fn extract_img_urls(&mut self) {
        if let Some(content_ref) = &self.content {
            for img_ref in content_ref.as_node().select("img").unwrap() {
                img_ref.as_node().as_element().map(|img_elem| {
                    img_elem.attributes.borrow().get("src").map(|img_url| {
                        if !img_url.is_empty() {
                            self.img_urls.push((img_url.to_string(), None))
                        }
                    })
                });
            }
        }
    }

    pub async fn download_images(&mut self, article_origin: &Url) -> async_std::io::Result<()> {
        let mut async_download_tasks = Vec::with_capacity(self.img_urls.len());
        self.extract_img_urls();
        println!("Downloading images to res/");
        for img_url in &self.img_urls {
            let img_url = img_url.0.clone();
            let abs_url = get_absolute_url(&img_url, article_origin);
            async_download_tasks.push(task::spawn(async move {
                let mut img_response = surf::get(&abs_url).await.expect("Unable to retrieve file");
                let img_content: Vec<u8> = img_response.body_bytes().await.unwrap();
                let img_mime = img_response
                    .header("Content-Type")
                    .map(|header| header.to_string());
                let img_ext = img_response
                    .header("Content-Type")
                    .and_then(map_mime_type_to_ext)
                    .unwrap();

                let img_path = format!("res/{}{}", hash_url(&abs_url), &img_ext);
                let mut img_file = File::create(&img_path)
                    .await
                    .expect("Unable to create file");
                img_file
                    .write_all(&img_content)
                    .await
                    .expect("Unable to save to file");

                (img_url, img_path, img_mime)
            }));
        }

        self.img_urls.clear();

        for async_task in async_download_tasks {
            let (img_url, img_path, img_mime) = async_task.await;
            // Update the image sources
            let img_ref = self
                .content
                .as_mut()
                .expect("Unable to get mutable ref")
                .as_node()
                .select_first(&format!("img[src='{}']", img_url))
                .expect("Image node does not exist");
            let mut img_node = img_ref.attributes.borrow_mut();
            *img_node.get_mut("src").unwrap() = img_path.clone();
            self.img_urls.push((img_path, img_mime));
        }
        Ok(())
    }
}

fn extract_text_from_node(node: &NodeRef) -> Option<String> {
    node.first_child()
        .map(|child_ref| child_ref.text_contents())
}

/// Utility for hashing URLs. This is used to help store files locally with unique values
fn hash_url(url: &str) -> String {
    format!("{:x}", md5::compute(url.as_bytes()))
}

/// Handles getting the extension from a given MIME type. The extension starts with a dot
fn map_mime_type_to_ext(mime_type: &str) -> Option<String> {
    mime_type
        .split("/")
        .last()
        .map(|format| {
            if format == ("svg+xml") {
                return "svg";
            } else if format == "x-icon" {
                "ico"
            } else {
                format
            }
        })
        .map(|format| String::from(".") + format)
}

fn get_absolute_url(url: &str, request_url: &Url) -> String {
    if Url::parse(url).is_ok() {
        url.to_owned()
    } else if url.starts_with("/") {
        Url::parse(&format!(
            "{}://{}",
            request_url.scheme(),
            request_url.host_str().unwrap()
        ))
        .unwrap()
        .join(url)
        .unwrap()
        .into_string()
    } else {
        request_url.join(url).unwrap().into_string()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    const TEST_HTML: &'static str = r#"
        <!doctype html>
        <html lang="en">
            <head>
                <meta charset="utf-8">
                <meta name="description" content="A sample document">
                <meta name="keywords" content="test,Rust">
                <meta name="author" content="Paperoni">                
                <title>Testing Paperoni</title>
            </head>
            <body>
                <header>
                <!-- Unimportant information -->
                    <h1>Testing Paperoni</h1>
                </header>
                <article>
                    <h1>Starting out</h1>
                    <p>Some Lorem Ipsum text here</p>
                    <p>Observe this picture</p>
                    <img src="./img.jpg" alt="Random image">
                </article>
                <footer>
                    <p>Made in HTML</p>
                </footer>
            </body>
        </html>
        "#;

    #[test]
    fn test_extract_attr_val() {
        let extractor = Extractor::from_html(TEST_HTML);
        let ext_author =
            extractor.extract_attr_val("meta[name='author']", "content", |val| val.to_string());
        assert!(ext_author.is_some());
        assert_eq!("Paperoni", &ext_author.unwrap());
        let ext_author =
            extractor.extract_attr_val("meta[name='invalid-name']", "content", |val| {
                val.to_string()
            });
        assert!(ext_author.is_none());
        let lang_attr = extractor.extract_attr_val("html", "lang", |lang| lang.to_string());
        assert!(lang_attr.is_some());
        assert_eq!("en".to_string(), lang_attr.unwrap());
    }

    #[test]
    fn test_extract_inner_text() {
        let extractor = Extractor::from_html(TEST_HTML);
        let title_text = extractor.extract_inner_text("title");
        assert!(title_text.is_some());
        assert_eq!("Testing Paperoni".to_string(), title_text.unwrap());

        let title_text = extractor.extract_inner_text("titln");
        assert!(title_text.is_none());
    }
    #[test]
    fn test_extract_text() {
        let extractor = Extractor::from_html(TEST_HTML);
        let h1_node = extractor.root_node.select_first("h1").unwrap();
        let h1_text = extract_text_from_node(h1_node.as_node());
        assert!(h1_text.is_some());
        assert_eq!("Testing Paperoni".to_string(), h1_text.unwrap());
    }

    #[test]
    fn test_extract_content() {
        let extracted_html: String = r#"
            <article>
                <h1>Starting out</h1>
                <p>Some Lorem Ipsum text here</p>
                <p>Observe this picture</p>
                <img alt="Random image" src="./img.jpg">
            </article>
        "#
        .lines()
        .map(|line| line.trim())
        .collect();

        let mut extractor = Extractor::from_html(
            &TEST_HTML
                .lines()
                .map(|line| line.trim())
                .collect::<String>(),
        );

        extractor.extract_content();
        let mut output = Vec::new();
        assert!(extractor.content.is_some());

        extractor
            .content
            .unwrap()
            .as_node()
            .serialize(&mut output)
            .expect("Unable to serialize output HTML");
        let output = std::str::from_utf8(&output).unwrap();

        assert_eq!(extracted_html, output);
    }

    #[test]
    fn test_extract_img_urls() {
        let mut extractor = Extractor::from_html(TEST_HTML);
        extractor.extract_content();
        extractor.extract_img_urls();

        assert!(extractor.img_urls.len() > 0);
        assert_eq!(vec![("./img.jpg".to_string(), None)], extractor.img_urls);
    }

    #[test]
    fn test_map_mime_type_to_ext() {
        let mime_types = vec![
            "image/apng",
            "image/bmp",
            "image/gif",
            "image/x-icon",
            "image/jpeg",
            "image/png",
            "image/svg+xml",
            "image/tiff",
            "image/webp",
        ];
        let exts = mime_types
            .into_iter()
            .map(|mime_type| map_mime_type_to_ext(mime_type).unwrap())
            .collect::<Vec<_>>();
        assert_eq!(
            vec![".apng", ".bmp", ".gif", ".ico", ".jpeg", ".png", ".svg", ".tiff", ".webp"],
            exts
        );
    }

    #[test]
    fn test_get_absolute_url() {
        let absolute_url = "https://example.image.com/images/1.jpg";
        let relative_url = "../../images/2.jpg";
        let relative_from_host_url = "/images/3.jpg";
        let host_url = Url::parse("https://example.image.com/blog/how-to-test-resolvers/").unwrap();
        let abs_url = get_absolute_url(&absolute_url, &host_url);
        assert_eq!("https://example.image.com/images/1.jpg", abs_url);
        let abs_url = get_absolute_url(&relative_url, &host_url);
        assert_eq!("https://example.image.com/images/2.jpg", abs_url);
        let relative_url = "2-1.jpg";
        let abs_url = get_absolute_url(&relative_url, &host_url);
        assert_eq!(
            "https://example.image.com/blog/how-to-test-resolvers/2-1.jpg",
            abs_url
        );
        let abs_url = get_absolute_url(&relative_from_host_url, &host_url);
        assert_eq!("https://example.image.com/images/3.jpg", abs_url);
    }
}