cyberstorm/src/render/mdbook.rs

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::{io, mem};

use async_trait::async_trait;
use lazy_static::lazy_static;
use mdbook::config::{Config as MDBookConfig, HtmlConfig};
use regex::Regex;
use serde::Serialize;
use tera::{Context, Map, Tera, Value};
use thiserror::Error;
use tokio::fs;

include!(concat!(env!("OUT_DIR"), "/mdbook-templates.rs"));

use crate::document::Document;
use crate::domains::common::{DomainId, DomainModel, ModelId};
use crate::domains::GenericDocument;
use crate::registry::{Registry, RegistryConfig};

use crate::render::{Engine, RendererError};

type DocumentMap = HashMap<DomainId, Map<String, Value>>;

#[derive(Debug, Error)]
pub enum MDBookEngineError {
    #[error("{0}")]
    Io(#[from] io::Error),
    #[error("{0}")]
    Tera(#[from] tera::Error),
}

impl From<MDBookEngineError> for RendererError {
    fn from(err: MDBookEngineError) -> Self {
        RendererError::Engine(err.into())
    }
}

#[derive(Debug, Serialize)]
struct SummaryItem {
    id: DomainId,
    name: String,
}

pub struct MDBookEngine {
    inner: Inner,
}

#[derive(Debug)]
#[allow(clippy::large_enum_variant)]
enum Inner {
    Start {
        src: PathBuf,
        global: Map<String, Value>,
    },
    LoadAndRender {
        src: PathBuf,
        global: Map<String, Value>,
        templates: Tera,
        summary: Vec<SummaryItem>,
        documents: DocumentMap,
    },
    Finish,
}

impl MDBookEngine {
    pub fn new(root: &Path, config: &MDBookConfig) -> Self {
        let mut global = Map::new();
        if let Ok(Some(HtmlConfig { site_url, .. })) = config.get_deserialized_opt("output.html") {
            global.insert(
                "site_url".to_owned(),
                site_url.as_deref().unwrap_or("/").to_owned().into(),
            );
        }
        Self {
            inner: Inner::Start {
                src: root.join(&config.book.src),
                global,
            },
        }
    }
}

// (?Send)
#[async_trait]
impl Engine for MDBookEngine {
    async fn start<R>(&mut self, _registry: &R) -> Result<(), RendererError>
    where
        R: Registry,
    {
        if let Inner::Start { src, global } = mem::replace(&mut self.inner, Inner::Finish) {
            let templates = load_templates(&src).await?;
            let documents = HashMap::new();
            let summary = Vec::new();
            self.inner = Inner::LoadAndRender {
                src,
                global,
                templates,
                documents,
                summary,
            };
        }
        Ok(())
    }

    async fn push_document<R>(
        &mut self,
        model_id: ModelId,
        doc: GenericDocument,
        _registry: &R,
    ) -> Result<(), RendererError>
    where
        R: Registry,
    {
        if let Inner::LoadAndRender {
            ref mut summary,
            ref mut documents,
            ..
        } = self.inner
        {
            macro_rules! push_documents {
                ($($kind:ident),+) => {
                    match doc {
                        $(
                            GenericDocument::$kind(doc) => {
                                save_domain_model(summary, documents, model_id, &doc);
                            }
                        )+
                    }
                };
            }
            push_documents!(
                SourceIntelligence,
                SourceRequirement,
                SourceProvider,
                ThreatTactic,
                ThreatTechnique,
                ThreatSoftware,
                ThreatSimulation,
                ObserveEvent,
                ObserveDetection,
                ObserveProvider,
                ObserveConfiguration,
                ReactStage,
                ReactAction,
                ReactPlaybook,
                MitigateStrategy,
                MitigatePlatform,
                MitigateConfiguration
            );
        }

        Ok(())
    }

    async fn finish<R>(&mut self, registry: &R) -> Result<(), RendererError>
    where
        R: Registry,
    {
        if let Inner::LoadAndRender {
            src,
            global,
            mut templates,
            documents,
            summary,
        } = mem::replace(&mut self.inner, Inner::Finish)
        {
            let summary = Arc::new(summary);
            let documents = Arc::new(documents);
            let registry_config = registry.get_config().await?;
            register_filters(&mut templates, documents.clone(), registry_config, global);
            render_indexes(&templates, &src, &summary).await?;
            render_summary(&templates, &src, &summary).await?;
            render_documents(&templates, &src, documents).await?;
        }
        Ok(())
    }
}

async fn load_templates(src: &Path) -> Result<Tera, MDBookEngineError> {
    let mut files = Vec::new();

    for filename in BASE_TEMPLATES.file_names() {
        let name = filename.strip_prefix("book/").unwrap();
        let external = src.join(name);

        let content = if external.is_file() {
            fs::read_to_string(external).await?
        } else {
            let mut content = String::new();
            BASE_TEMPLATES
                .read(filename)?
                .read_to_string(&mut content)?;
            content
        };

        files.push((name, content));
    }

    let mut templates = Tera::default();
    templates.add_raw_templates(files)?;
    Ok(templates)
}

fn save_domain_model<M>(
    summary: &mut Vec<SummaryItem>,
    documents: &mut DocumentMap,
    model_id: ModelId,
    doc: &Document<M>,
) where
    M: DomainModel,
{
    let id = DomainId::new(M::kind(), model_id);
    if let Value::Object(mut document) = tera::to_value(&doc).unwrap() {
        document.insert("id".to_owned(), tera::to_value(&id).unwrap());
        documents.insert(id, document);
    }
    summary.push(SummaryItem {
        id,
        name: doc.meta.name().to_owned(),
    });
}

async fn render_documents(
    templates: &Tera,
    src: &Path,
    documents: Arc<DocumentMap>,
) -> Result<(), MDBookEngineError> {
    for (id, document) in documents.iter() {
        let (domain, model) = id.kind.parts();
        let model_path = format!("{}/{}", domain, model);
        let template_path = format!("{}.instance.tera", model_path);
        let output_path = src
            .join(model_path)
            .join(id.model_id.to_string())
            .with_extension("md");
        let mut context = Context::new();
        context.insert("doc", &document);
        let contents = templates.render(&template_path, &context)?;
        fs::write(output_path, contents).await?;
    }
    Ok(())
}

async fn render_summary(
    templates: &Tera,
    src: &Path,
    summary: &[SummaryItem],
) -> Result<(), MDBookEngineError> {
    let mut context = Context::new();
    context.insert("summary", summary);

    let contents = templates.render("summary.tera", &context)?;
    let summary_path = src.join("SUMMARY.md");
    fs::write(summary_path, contents).await?;
    Ok(())
}

async fn render_indexes(
    templates: &Tera,
    src: &Path,
    summary: &[SummaryItem],
) -> Result<(), MDBookEngineError> {
    let mut last_i = 0;
    let mut last_item = if let Some(item) = summary.get(0) {
        item
    } else {
        return Ok(());
    };
    for (i, item) in summary.iter().enumerate() {
        if item.id.kind != last_item.id.kind || i + 1 == summary.len() {
            let (domain, model) = last_item.id.kind.parts();

            let domain_dir = src.join(domain);
            let model_dir = domain_dir.join(model);

            if !domain_dir.is_dir() {
                fs::create_dir(&domain_dir).await?;
            }
            if !model_dir.is_dir() {
                fs::create_dir(&model_dir).await?;
            }

            let mut context = Context::new();
            context.insert("instances", &summary[last_i..i]);

            let model_index = model_dir.with_extension("md");
            let template_path = format!("{}/{}.tera", domain, model);
            let contents = templates.render(&template_path, &context)?;
            fs::write(model_index, contents).await?;

            last_i = i;
            last_item = item;
        }
    }

    Ok(())
}

fn register_filters(
    tera: &mut Tera,
    documents: Arc<DocumentMap>,
    registry_config: Arc<RegistryConfig>,
    global: Map<String, Value>,
) {
    use tera::{try_get_value, Error};

    tera.register_function("global", move |args: &HashMap<String, Value>| {
        if let Some(Value::String(key)) = args.get("key") {
            Ok(global.get(key).cloned().unwrap_or(Value::Null))
        } else {
            Err(Error::msg("expected `key` for `global` function"))
        }
    });

    fn parse_domain_id(value: &Value) -> Result<DomainId, Error> {
        let id = try_get_value!("link", "value", String, value);
        id.parse()
            .map_err(|_| Error::msg(format_args!("failed to parse domain id `{}`", id)))
    }

    tera.register_filter(
        "domain_id_link",
        move |value: &Value, args: &HashMap<String, Value>| {
            let id = parse_domain_id(value)?;
            let (domain, model) = id.kind.parts();
            let link = match args.get("for").and_then(|v| v.as_str()) {
                None => format!("{}/{}/{}.md", domain, model, id.model_id),
                Some("model") => format!("{}/{}.md", domain, model),
                Some("edit") => registry_config.edit_link(id).map_err(Error::msg)?,
                Some(_) => {
                    return Err(Error::msg(
                        "domain_id_link `for` arg can only be `model|edit`",
                    ))
                }
            };
            Ok(Value::String(link))
        },
    );

    tera.register_filter(
        "domain_id",
        move |value: &Value, _: &HashMap<String, Value>| {
            let id = parse_domain_id(value)?;
            let (domain, model) = id.kind.parts();
            let mut parts = Map::with_capacity(4);
            parts.insert("domain".to_owned(), domain.into());
            parts.insert("model".to_owned(), model.into());
            parts.insert("model_id".to_owned(), id.model_id.to_string().into());
            Ok(Value::Object(parts))
        },
    );

    tera.register_filter(
        "get_doc",
        move |value: &Value, _: &HashMap<String, Value>| {
            let id = parse_domain_id(value)?;
            documents
                .get(&id)
                .cloned()
                .map(Value::Object)
                .ok_or_else(|| Error::msg(format_args!("unknown document id {}", id)))
        },
    );

    tera.register_filter("autolink", |value: &Value, _: &HashMap<String, Value>| {
        lazy_static! {
            static ref URL: Regex = Regex::new(
                r"(?ix)
                \b(([\w-]+://?|www[.])[^\s()<>]+(?:\([\w\d]+\)|([^[:punct:]\s]|/)))
            ",
            )
            .unwrap();
        }
        let text = try_get_value!("autolink", "value", String, value);
        let text = if text.is_empty() {
            text
        } else {
            URL.replace_all(text.as_str(), "[$0]($0)").into_owned()
        };
        Ok(Value::String(text))
    });
}