439 lines
13 KiB
Rust
439 lines
13 KiB
Rust
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, Instance};
|
|
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,
|
|
site_url: Option<String>,
|
|
},
|
|
LoadAndRender {
|
|
src: PathBuf,
|
|
site_url: Option<String>,
|
|
templates: Tera,
|
|
summary: Vec<SummaryItem>,
|
|
documents: DocumentMap,
|
|
},
|
|
Finish,
|
|
}
|
|
|
|
impl MDBookEngine {
|
|
pub fn new(root: &Path, config: &MDBookConfig) -> Self {
|
|
let site_url = if let Ok(Some(HtmlConfig { site_url, .. })) =
|
|
config.get_deserialized_opt("output.html")
|
|
{
|
|
site_url
|
|
} else {
|
|
None
|
|
};
|
|
Self {
|
|
inner: Inner::Start {
|
|
src: root.join(&config.book.src),
|
|
site_url,
|
|
},
|
|
}
|
|
}
|
|
}
|
|
|
|
// (?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, site_url } = 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,
|
|
site_url,
|
|
templates,
|
|
documents,
|
|
summary,
|
|
};
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn push_document<R>(
|
|
&mut self,
|
|
instance: Instance,
|
|
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, instance, &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,
|
|
site_url,
|
|
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, site_url);
|
|
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,
|
|
instance: Instance,
|
|
doc: &Document<M>,
|
|
) where
|
|
M: DomainModel,
|
|
{
|
|
let id = DomainId::new(M::kind(), instance);
|
|
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.instance.to_string())
|
|
.with_extension("md");
|
|
let mut context = Context::new();
|
|
context.insert("document", &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>,
|
|
site_url: Option<String>,
|
|
) {
|
|
use tera::{try_get_value, Error};
|
|
|
|
let site_url = site_url.as_deref().unwrap_or("/").to_owned();
|
|
|
|
let mut global: HashMap<&str, Value> = HashMap::new();
|
|
global.insert("site_url", tera::to_value(&site_url).unwrap());
|
|
global.insert("registry", tera::to_value(&*registry_config).unwrap());
|
|
|
|
tera.register_function("global", move |args: &HashMap<String, Value>| {
|
|
if let Some(Value::String(full_key)) = args.get("key") {
|
|
let mut keys = full_key.split('.');
|
|
let mut value = global.get(keys.next().unwrap());
|
|
for key in keys {
|
|
match value {
|
|
Some(Value::Object(map)) => {
|
|
value = map.get(key);
|
|
}
|
|
None => {
|
|
return Err(Error::msg(format_args!(
|
|
"no key `{}` section of `{}`",
|
|
key, full_key
|
|
)))
|
|
}
|
|
Some(_) => {
|
|
return Err(Error::msg(format_args!(
|
|
"value cannot be indexed with `{}`",
|
|
key
|
|
)))
|
|
}
|
|
}
|
|
}
|
|
Ok(value.cloned().unwrap_or(Value::Null))
|
|
} else {
|
|
Err(Error::msg("expected valid `key` arg for `global` function"))
|
|
}
|
|
});
|
|
|
|
tera.register_function("link_for", move |args: &HashMap<String, Value>| {
|
|
let link_for = if let Some(Value::String(link_for)) = args.get("for") {
|
|
link_for
|
|
} else {
|
|
return Err(Error::msg(
|
|
"expected string `for` arg for `link_for` function",
|
|
));
|
|
};
|
|
let edit = match args.get("edit") {
|
|
Some(Value::Bool(edit)) => *edit,
|
|
Some(_) => return Err(Error::msg("`edit` arg for `link_for` must be a boolean")),
|
|
None => false,
|
|
};
|
|
let link = match link_for.as_str() {
|
|
"document" => match args.get("id") {
|
|
Some(value) => {
|
|
let id = parse_domain_id("link_for", value)?;
|
|
if edit {
|
|
registry_config.document_edit_link(id).map_err(Error::msg)?
|
|
} else {
|
|
let (domain, model) = id.kind.parts();
|
|
format!("{}{}/{}/{}.md", site_url, domain, model, id.instance)
|
|
}
|
|
}
|
|
_ => {
|
|
return Err(Error::msg(
|
|
"`link_for(for=\"document\")` `` expects `id` arg",
|
|
))
|
|
}
|
|
},
|
|
"page" => match args.get("path") {
|
|
Some(Value::String(page)) if edit => {
|
|
registry_config.page_edit_link(page).map_err(Error::msg)?
|
|
}
|
|
Some(Value::String(page)) => {
|
|
format!("{}{}.md", site_url, page)
|
|
}
|
|
_ => return Err(Error::msg("`link_for(for=\"path\")` `` expects `path` arg")),
|
|
},
|
|
_ => {
|
|
return Err(Error::msg(
|
|
"`link_for` `for` arg can only be `document|page`",
|
|
))
|
|
}
|
|
};
|
|
Ok(Value::String(link))
|
|
});
|
|
|
|
tera.register_function("document", move |args: &HashMap<String, Value>| {
|
|
if let Some(value) = args.get("id") {
|
|
let id = parse_domain_id("document", value)?;
|
|
documents
|
|
.get(&id)
|
|
.cloned()
|
|
.map(Value::Object)
|
|
.ok_or_else(|| Error::msg(format_args!("unknown document id {}", id)))
|
|
} else {
|
|
Err(Error::msg("expected `id` arg for `document` function"))
|
|
}
|
|
});
|
|
|
|
fn parse_domain_id(within: &str, value: &Value) -> Result<DomainId, Error> {
|
|
let id = try_get_value!(within, "id", String, value);
|
|
id.parse()
|
|
.map_err(|_| Error::msg(format_args!("failed to parse domain id `{}`", id)))
|
|
}
|
|
|
|
tera.register_filter(
|
|
"domain_id",
|
|
move |value: &Value, _: &HashMap<String, Value>| {
|
|
let id = parse_domain_id("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("instance".to_owned(), id.instance.to_string().into());
|
|
Ok(Value::Object(parts))
|
|
},
|
|
);
|
|
|
|
tera.register_tester("url", |value: Option<&Value>, _: &[Value]| {
|
|
lazy_static! {
|
|
static ref URL: Regex = Regex::new(
|
|
r"(?ix)
|
|
^\b(([\w-]+://?|www[.])[^\s()<>]+(?:\([\w\d]+\)|([^[:punct:]\s]|/)))$
|
|
",
|
|
)
|
|
.unwrap();
|
|
}
|
|
if let Some(Value::String(text)) = value {
|
|
Ok(URL.is_match(text))
|
|
} else {
|
|
Err(Error::msg("expected string for `url` test"))
|
|
}
|
|
});
|
|
}
|