avitex/cyberstorm
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

refactor usage

Browse Source
This commit is contained in:
avitex 2021-03-03 13:38:15 +11:00
parent b07d70c769
commit 5c25d628f5
Signed by: avitex
GPG Key ID: 38C76CBF3749D62C
13 changed files with 248 additions and 118 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
Cargo.toml
View File

@ -6,24 +6,37 @@ edition = "2018"
publish = false
build = "build.rs"
[features]
default = ["cli"]
cli = ["structopt", "tokio/macros", "tokio/rt-multi-thread"]
mdbook-renderer = ["mdbook", "includedir", "includedir_codegen", "glob", "phf"]
[[bin]]
name = "cyberstorm"
required-features = ["cli"]
[dependencies]
zc = "0.4"
log = "0.4"
serde = { version = "1.0", features = ["derive"] }
toml = "0.5"
async-trait = "0.1"
strum = { version = "0.20", features = ["derive"] }
roxmltree = "0.14"
tokio = { version = "1", features = ["rt-multi-thread", "fs", "macros"] }
validator = { version = "0.12", features = ["derive"] }
tokio = { version = "1", features = ["fs"] }
thiserror = "1.0"
tera = "1.6"
anyhow = "1.0"
regex = "1"
lazy_static = "1"
structopt = "0.3"
phf = "0.8.0"
includedir = "0.6"
env_logger = "0.8"
chrono = "0.4"
structopt = { version = "0.3", optional = true }
phf = { version = "0.8.0", optional = true }
includedir = { version = "0.6", optional = true }
mdbook = { version = "0.4", optional = true }
[build-dependencies]
includedir_codegen = "0.6"
glob = "0.3"
includedir_codegen = { version = "0.6", optional = true }
glob = { version = "0.3", optional = true }

56
book/macros.tera
View File

@ -1,3 +1,8 @@
{% macro content(content) %}
{%- set trimmed = content | trim -%}
{%- if trimmed | length == 0 %}No description{% else %}{{ trimmed }}{% endif -%}
{% endmacro content %}
{% macro references(refs) %}
{%- if refs | length == 0 %}No references{% endif -%}
{%- for ref in refs -%}
@ -5,46 +10,41 @@
{% endfor -%}
{% endmacro references %}
{% macro title(title, id) -%}
{{ title }} <small>([edit]({{ id | domain_id_link(for="edit") }}))</small>
{%- endmacro details_next %}
{% macro doc_title(doc) -%}
{{ doc.name }} <small>([edit]({{ doc.id | domain_id_link(for="edit") }}))</small>
{%- endmacro doc_title %}
{% macro details(id, name) -%}
| Title | {{ name }} |
{% macro doc_details(doc) -%}
| Title | {{ doc.name }} |
|:---------------------------:|:------------------------|
{{ self::details_next(title="ID", value=id)}}
{%- endmacro details %}
{{ self::doc_details_next(title="ID", value=doc.id)}}
{%- endmacro doc_details %}
{% macro details_next(title, value) -%}
{% macro doc_details_next(title, value) -%}
| **{{ title }}** | {{ value }} |
{%- endmacro details_next %}
{%- endmacro doc_details_next %}
{% macro details_authors(authors) -%}
{{ self::details_next(title="Authors", value=authors | join )}}
{%- endmacro details_next %}
{% macro doc_details_authors(authors) -%}
{{ self::doc_details_next(title="Authors", value=authors | join )}}
{%- endmacro doc_details_authors %}
{% macro details_tags(tags) -%}
{% macro doc_details_tags(tags) -%}
{% if tags | length == 0 -%}
{{ self::details_next(title="Tags", value="No tags") }}
{{ self::doc_details_next(title="Tags", value="No tags") }}
{%- else -%}
{{ self::details_next(title="Tags", value=tags | join) }}
{{ self::doc_details_next(title="Tags", value=tags | join) }}
{%- endif %}
{%- endmacro details_next %}
{%- endmacro doc_details_tags %}
{% macro content(content) %}
{%- set trimmed = content | trim -%}
{%- if trimmed | length == 0 %}No description{% else %}{{ trimmed }}{% endif -%}
{% endmacro details_next %}
{% macro name_and_id_link(value) -%}
{{ value["doc"]["name"] }} ([{{ value["id"] }}]({{ value["id"] | domain_id_link }}))
{%- endmacro name_and_id_link %}
{% macro doc_rich_link(doc) -%}
{{ doc.name }} ([{{ doc.id }}]({{ global(key="site_url")}}{{ doc.id | domain_id_link }}))
{%- endmacro doc_rich_link %}
{% macro summary_table(instances) -%}
| ID | Name |
|:---------------------------:|:------------------------|
{% for item in instances -%}
| [{{ item.id }}]({{ item.id | domain_id_link }}) | {{ item.name }} |
| [{{ item.id }}]({{ global(key="site_url") }}{{ item.id | domain_id_link }}) | {{ item.name }} |
{% endfor %}
{%- endmacro summary_table %}
@ -59,8 +59,8 @@
{%- endif -%}
{%- if last_model != id_parts.model -%}
{%- set_global last_model = id_parts.model %}
- [{{ id_parts.model | capitalize }}](.{{ item.id | domain_id_link(for="model") }})
- [{{ id_parts.model | capitalize }}](./{{ item.id | domain_id_link(for="model") }})
{%- endif %}
- [{{ item.name }}](.{{ item.id | domain_id_link }})
- [{{ item.name }}](./{{ item.id | domain_id_link }})
{%- endfor -%}
{% endmacro summary_list %}
{% endmacro summary_list %}

8
book/observe/event.instance.tera
View File

@ -1,12 +1,12 @@
{% import "macros.tera" as macros %}
# {{ macros::title(title=doc.name, id=id) }}
# {{ macros::doc_title(doc=doc) }}
{{ macros::details(id=id, name=doc.name) }}
{{ macros::doc_details(doc=doc) }}
{%- if doc.WindowsEvent %}
{{ macros::details_next(title="Type", value="Windows event") }}
{{ macros::doc_details_next(title="Type", value="Windows event") }}
{% endif -%}
{{ macros::details_next(title="Description", value=doc.description) }}
{{ macros::doc_details_next(title="Description", value=doc.description) }}
## Description

10
book/react/action.instance.tera
View File

@ -1,12 +1,12 @@
{% import "macros.tera" as macros %}
{% set stage = doc.stage | get_doc %}
# {{ macros::title(title=doc.name, id=id) }}
# {{ macros::doc_title(doc=doc) }}
{{ macros::details(id=id, name=doc.name) }}
{{ macros::details_next(title="Stage", value=macros::name_and_id_link(value=stage)) }}
{{ macros::details_next(title="Description", value=doc.description) }}
{{ macros::details_tags(tags=doc.tags) }}
{{ macros::doc_details(doc=doc) }}
{{ macros::doc_details_next(title="Stage", value=macros::doc_rich_link(doc=stage)) }}
{{ macros::doc_details_next(title="Description", value=doc.description) }}
{{ macros::doc_details_tags(tags=doc.tags) }}
## Description

6
book/react/stage.instance.tera
View File

@ -1,9 +1,9 @@
{% import "macros.tera" as macros %}
# {{ macros::title(title=doc.name, id=id) }}
# {{ macros::doc_title(doc=doc) }}
{{ macros::details(id=id, name=doc.name) }}
{{ macros::details_next(title="Description", value=doc.description) }}
{{ macros::doc_details(doc=doc) }}
{{ macros::doc_details_next(title="Description", value=doc.description) }}
## Description

6
book/source/intelligence.instance.tera
View File

@ -1,9 +1,9 @@
{% import "macros.tera" as macros %}
# {{ macros::title(title=doc.name, id=id) }}
# {{ macros::doc_title(doc=doc) }}
{{ macros::details(id=id, name=doc.name) }}
{{ macros::details_next(title="Provider", value=doc.provider) }}
{{ macros::doc_details(doc=doc) }}
{{ macros::doc_details_next(title="Provider", value=doc.provider) }}
## Description

4
book/source/provider.instance.tera
View File

@ -1,8 +1,8 @@
{% import "macros.tera" as macros %}
# {{ macros::title(title=doc.name, id=id) }}
# {{ macros::doc_title(doc=doc) }}
{{ macros::details(id=id, name=doc.name) }}
{{ macros::doc_details(doc=doc) }}
## Description

4
book/source/requirement.instance.tera
View File

@ -1,8 +1,8 @@
{% import "macros.tera" as macros %}
# {{ macros::title(title=doc.name, id=id) }}
# {{ macros::doc_title(doc=doc) }}
{{ macros::details(id=id, name=doc.name) }}
{{ macros::doc_details(doc=doc) }}
## Description

12
build.rs
View File

@ -1,7 +1,13 @@
use glob::glob;
use includedir_codegen::Compression;
fn main() {
#[cfg(feature = "mdbook-renderer")]
include_mdbook_templates();
}
#[cfg(feature = "mdbook-renderer")]
fn include_mdbook_templates() {
use glob::glob;
use includedir_codegen::Compression;
let mut templates = includedir_codegen::start("BASE_TEMPLATES");
for path_result in glob("book/**/*.tera").unwrap() {

2
src/lib.rs
View File

@ -2,5 +2,5 @@ mod util;
pub mod document;
pub mod domains;
pub mod generator;
pub mod registry;
pub mod render;

114
src/main.rs
View File

@ -1,12 +1,18 @@
use std::path::PathBuf;
use std::env;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use anyhow::{Context, Error};
use anyhow::{anyhow, Context, Error};
use chrono::Local;
use log::LevelFilter;
use structopt::StructOpt;
use tokio::fs;
use cyberstorm::generator::mdbook::MDBookEngine;
use cyberstorm::generator::Generator;
use cyberstorm::registry::{DirectoryRegistry, SupportedRegistry};
#[cfg(feature = "mdbook-renderer")]
use cyberstorm::render::mdbook::MDBookEngine;
use cyberstorm::render::Renderer;
#[derive(Debug, StructOpt)]
#[structopt(name = "cyberstorm")]
@ -20,13 +26,19 @@ struct Opt {
#[derive(Debug, StructOpt)]
enum Cmd {
BuildMdbook(BuildMdbookOpt),
/// Renders the content into a mdBook.
#[cfg(feature = "mdbook-renderer")]
MdbookRender(MdbookRenderOpt),
/// Renders the content into a mdBook and builds it.
#[cfg(feature = "mdbook-renderer")]
MdbookBuild(MdbookRenderOpt),
}
#[derive(Debug, StructOpt)]
struct BuildMdbookOpt {
/// Path to the output MDBook directory.
#[structopt(parse(from_os_str))]
#[cfg(feature = "mdbook-renderer")]
struct MdbookRenderOpt {
/// Path to the output mdBook root or `book.toml`.
#[structopt(parse(from_os_str), default_value = "./book.toml")]
book: PathBuf,
}
@ -48,25 +60,97 @@ async fn run(opt: Opt) -> Result<(), Error> {
RegistryOpt::Directory(path) => SupportedRegistry::Directory(DirectoryRegistry::new(path)),
};
match &opt.cmd {
Cmd::BuildMdbook(build_opt) => build_mdbook(&opt, build_opt, &registry).await,
#[cfg(feature = "mdbook-renderer")]
Cmd::MdbookRender(render_opt) => {
let (root, config) = load_mdbook_config(&render_opt.book).await?;
render_mdbook(&registry, &root, &config).await
}
#[cfg(feature = "mdbook-renderer")]
Cmd::MdbookBuild(render_opt) => {
let (root, config) = load_mdbook_config(&render_opt.book).await?;
render_mdbook(&registry, &root, &config).await?;
mdbook::MDBook::load_with_config(root, config)?.build()
}
}
}
async fn build_mdbook(
_opt: &Opt,
build_opt: &BuildMdbookOpt,
#[cfg(feature = "mdbook-renderer")]
async fn load_mdbook_config(book: &Path) -> Result<(PathBuf, mdbook::Config), Error> {
let (root, config_file) = book
.canonicalize()
.ok()
.and_then(|book| {
let (root, config) = if book.extension().is_some() {
(book.parent().unwrap().to_owned(), book.to_owned())
} else {
let config_file = book.join("book.toml");
(book, config_file)
};
if config.is_file() {
Some((root, config))
} else {
None
}
})
.ok_or_else(|| {
anyhow!(
"expected mdbook config path (eg. `book.toml`) at `{}`",
book.display()
)
})?;
let config = fs::read_to_string(config_file)
.await
.context("failed to read mdbook configuration")?
.parse()?;
Ok((root, config))
}
#[cfg(feature = "mdbook-renderer")]
async fn render_mdbook(
registry: &SupportedRegistry,
root: &Path,
config: &mdbook::Config,
) -> Result<(), Error> {
let engine = MDBookEngine::new(&build_opt.book).context("failed to load mdbook engine")?;
Generator::new(engine, registry)
let engine = MDBookEngine::new(root, config);
Renderer::new(engine, registry)
.generate()
.await
.context("failed to generate mdbook content")
}
fn init_logger() {
let mut builder = env_logger::Builder::new();
builder.format(|formatter, record| {
writeln!(
formatter,
"{} [{}] ({}): {}",
Local::now().format("%Y-%m-%d %H:%M:%S"),
record.level(),
record.target(),
record.args()
)
});
if let Ok(var) = env::var("RUST_LOG") {
builder.parse_filters(&var);
} else {
// if no RUST_LOG provided, default to logging at the Info level
builder.filter(None, LevelFilter::Info);
// Filter extraneous html5ever not-implemented messages
builder.filter(Some("html5ever"), LevelFilter::Error);
}
builder.init();
}
#[tokio::main]
async fn main() {
init_logger();
if let Err(err) = run(Opt::from_args()).await {
eprintln!("{:#}", err);
log::error!("{:#}", err);
}
}

93
src/generator/mdbook.rs → src/render/mdbook.rs
View File

@ -5,6 +5,7 @@ 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};
@ -18,7 +19,9 @@ use crate::domains::common::{DomainId, DomainModel, ModelId};
use crate::domains::GenericDocument;
use crate::registry::{Registry, RegistryConfig};
use crate::generator::{Engine, GeneratorError};
use crate::render::{Engine, RendererError};
type DocumentMap = HashMap<DomainId, Map<String, Value>>;
#[derive(Debug, Error)]
pub enum MDBookEngineError {
@ -28,9 +31,9 @@ pub enum MDBookEngineError {
Tera(#[from] tera::Error),
}
impl From<MDBookEngineError> for GeneratorError {
impl From<MDBookEngineError> for RendererError {
fn from(err: MDBookEngineError) -> Self {
GeneratorError::Engine(err.into())
RendererError::Engine(err.into())
}
}
@ -49,39 +52,50 @@ pub struct MDBookEngine {
enum Inner {
Start {
src: PathBuf,
global: Map<String, Value>,
},
LoadAndRender {
src: PathBuf,
global: Map<String, Value>,
templates: Tera,
summary: Vec<SummaryItem>,
documents: HashMap<DomainId, Value>,
documents: DocumentMap,
},
Finish,
}
impl MDBookEngine {
pub fn new(src: &Path) -> Result<Self, MDBookEngineError> {
Ok(Self {
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: src.canonicalize()?,
src: root.join(&config.book.src),
global,
},
})
}
}
}
// (?Send)
#[async_trait]
impl Engine for MDBookEngine {
async fn start<R>(&mut self, _registry: &R) -> Result<(), GeneratorError>
async fn start<R>(&mut self, _registry: &R) -> Result<(), RendererError>
where
R: Registry,
{
if let Inner::Start { src } = mem::replace(&mut self.inner, Inner::Finish) {
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,
@ -95,7 +109,7 @@ impl Engine for MDBookEngine {
model_id: ModelId,
doc: GenericDocument,
_registry: &R,
) -> Result<(), GeneratorError>
) -> Result<(), RendererError>
where
R: Registry,
{
@ -140,12 +154,13 @@ impl Engine for MDBookEngine {
Ok(())
}
async fn finish<R>(&mut self, registry: &R) -> Result<(), GeneratorError>
async fn finish<R>(&mut self, registry: &R) -> Result<(), RendererError>
where
R: Registry,
{
if let Inner::LoadAndRender {
src,
global,
mut templates,
documents,
summary,
@ -154,7 +169,7 @@ impl Engine for MDBookEngine {
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);
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?;
@ -190,17 +205,17 @@ async fn load_templates(src: &Path) -> Result<Tera, MDBookEngineError> {
fn save_domain_model<M>(
summary: &mut Vec<SummaryItem>,
documents: &mut HashMap<DomainId, Value>,
documents: &mut DocumentMap,
model_id: ModelId,
doc: &Document<M>,
) where
M: DomainModel,
{
let id = DomainId::new(M::kind(), model_id);
let mut value = Map::new();
value.insert("id".to_owned(), tera::to_value(&id).unwrap());
value.insert("doc".to_owned(), tera::to_value(&doc).unwrap());
documents.insert(id, value.into());
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(),
@ -210,9 +225,9 @@ fn save_domain_model<M>(
async fn render_documents(
templates: &Tera,
src: &Path,
documents: Arc<HashMap<DomainId, Value>>,
documents: Arc<DocumentMap>,
) -> Result<(), MDBookEngineError> {
for (id, value) in documents.iter() {
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);
@ -220,7 +235,8 @@ async fn render_documents(
.join(model_path)
.join(id.model_id.to_string())
.with_extension("md");
let context = Context::from_value(value.clone())?;
let mut context = Context::new();
context.insert("doc", &document);
let contents = templates.render(&template_path, &context)?;
fs::write(output_path, contents).await?;
}
@ -284,11 +300,20 @@ async fn render_indexes(
fn register_filters(
tera: &mut Tera,
registry: Arc<HashMap<DomainId, Value>>,
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()
@ -301,8 +326,8 @@ fn register_filters(
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),
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(
@ -331,20 +356,15 @@ fn register_filters(
"get_doc",
move |value: &Value, _: &HashMap<String, Value>| {
let id = parse_domain_id(value)?;
registry
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>| {
let text = try_get_value!("autolink", "value", String, value);
if text.is_empty() {
return Ok(Value::String(String::new()));
}
lazy_static! {
static ref URL: Regex = Regex::new(
r"(?ix)
@ -353,9 +373,12 @@ fn register_filters(
)
.unwrap();
}
let replaced = URL.replace_all(text.as_str(), "[$0]($0)").into_owned();
Ok(Value::String(replaced))
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))
});
}

24
src/generator/mod.rs → src/render/mod.rs
View File

@ -1,3 +1,4 @@
#[cfg(feature = "mdbook-renderer")]
pub mod mdbook;
use async_trait::async_trait;
@ -9,19 +10,19 @@ use crate::domains::GenericDocument;
use crate::registry::{Registry, RegistryError};
#[derive(Debug, Error)]
pub enum GeneratorError {
pub enum RendererError {
#[error("{0}")]
Engine(#[from] anyhow::Error),
#[error("{0}")]
Registry(#[from] RegistryError),
}
pub struct Generator<E, R> {
pub struct Renderer<E, R> {
engine: E,
registry: R,
}
impl<E, R> Generator<E, R>
impl<E, R> Renderer<E, R>
where
E: Engine,
R: Registry,
@ -30,7 +31,8 @@ where
Self { engine, registry }
}
pub async fn generate(mut self) -> Result<(), GeneratorError> {
pub async fn generate(mut self) -> Result<(), RendererError> {
log::info!("Rendering of content has started");
self.engine.start(&self.registry).await?;
macro_rules! push_documents {
($($kind:ident),+) => {
@ -62,10 +64,12 @@ where
MitigatePlatform,
MitigateConfiguration
);
self.engine.finish(&self.registry).await
self.engine.finish(&self.registry).await?;
log::info!("Rendering of content has finished");
Ok(())
}
async fn get_documents<M>(&self) -> Result<Vec<(DomainId, Document<M>)>, GeneratorError>
async fn get_documents<M>(&self) -> Result<Vec<(DomainId, Document<M>)>, RendererError>
where
M: DomainModel,
{
@ -76,13 +80,13 @@ where
docs.sort_by_key(|(id, _)| *id);
docs
})
.map_err(GeneratorError::Registry)
.map_err(RendererError::Registry)
}
}
#[async_trait]
pub trait Engine {
async fn start<R>(&mut self, registry: &R) -> Result<(), GeneratorError>
async fn start<R>(&mut self, registry: &R) -> Result<(), RendererError>
where
R: Registry;
@ -91,11 +95,11 @@ pub trait Engine {
id: ModelId,
doc: GenericDocument,
registry: &R,
) -> Result<(), GeneratorError>
) -> Result<(), RendererError>
where
R: Registry;
async fn finish<R>(&mut self, registry: &R) -> Result<(), GeneratorError>
async fn finish<R>(&mut self, registry: &R) -> Result<(), RendererError>
where
R: Registry;
}