use crate::{
condition::Condition, entity::Entity, feature::Feature, helpers::Connection,
infrastructure::Infrastructure, vulnerability::Vulnerability, Formalize,
};
use anyhow::{anyhow, Result};
use bytesize::ByteSize;
use serde::{Deserialize, Deserializer, Serialize};
use serde_aux::prelude::*;
use std::collections::HashMap;
use crate::common::{HelperSource, Source};
fn parse_bytesize<'de, D>(deserializer: D) -> Result<u64, D::Error>
where
D: Deserializer<'de>,
{
let s = String::deserialize(deserializer)?;
Ok(s.parse::<ByteSize>()
.map_err(|_| serde::de::Error::custom("Failed"))?
.0)
}
#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
pub enum NodeType {
#[serde(alias = "vm", alias = "Vm")]
VM,
#[serde(alias = "switch", alias = "SWITCH")]
Switch,
}
#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
pub struct Resources {
#[serde(deserialize_with = "parse_bytesize")]
pub ram: u64,
pub cpu: u32,
}
#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
pub struct Role {
#[serde(alias = "Username", alias = "USERNAME")]
pub username: String,
#[serde(alias = "Entity", alias = "ENTITY")]
pub entities: Option<Vec<String>>,
}
pub type Roles = HashMap<String, Role>;
#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
pub struct Node {
#[serde(rename = "type", alias = "Type", alias = "TYPE")]
pub type_field: NodeType,
#[serde(alias = "Description", alias = "DESCRIPTION")]
pub description: Option<String>,
#[serde(
default,
alias = "Resources",
alias = "RESOURCES",
deserialize_with = "deserialize_struct_case_insensitive"
)]
pub resources: Option<Resources>,
#[serde(
default,
rename = "source",
alias = "Source",
alias = "SOURCE",
skip_serializing
)]
_source_helper: Option<HelperSource>,
#[serde(default, skip_deserializing)]
pub source: Option<Source>,
#[serde(default, alias = "Features", alias = "FEATURES")]
pub features: Option<HashMap<String, String>>,
#[serde(default, alias = "Conditions", alias = "CONDITIONS")]
pub conditions: Option<HashMap<String, String>>,
#[serde(default, alias = "Vulnerabilities", alias = "VULNERABILITIES")]
pub vulnerabilities: Option<Vec<String>>,
#[serde(
default,
rename = "roles",
alias = "Roles",
alias = "ROLES",
skip_serializing
)]
_roles_helper: Option<HelperRoles>,
#[serde(skip_deserializing)]
pub roles: Option<Roles>,
}
#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
#[serde(untagged)]
pub enum RoleTypes {
Username(String),
Role(Role),
}
#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
#[serde(untagged)]
pub enum HelperRoles {
MixedRoles(HashMap<String, RoleTypes>),
}
impl From<HelperRoles> for Roles {
fn from(helper_role: HelperRoles) -> Self {
match helper_role {
HelperRoles::MixedRoles(mixed_role) => mixed_role
.into_iter()
.map(|(role_name, role_value)| {
let role_value = match role_value {
RoleTypes::Role(role) => role,
RoleTypes::Username(username) => Role {
username,
entities: None,
},
};
(role_name, role_value)
})
.collect::<Roles>(),
}
}
}
impl Connection<Vulnerability> for (&String, &Node) {
fn validate_connections(
&self,
potential_vulnerability_names: &Option<Vec<String>>,
) -> Result<()> {
if let Some(node_vulnerabilities) = &self.1.vulnerabilities {
if let Some(vulnerabilities) = potential_vulnerability_names {
for vulnerability_name in node_vulnerabilities.iter() {
if !vulnerabilities.contains(vulnerability_name) {
return Err(anyhow!(
"Vulnerability \"{vulnerability_name}\" not found under Scenario Vulnerabilities",
));
}
}
} else {
return Err(anyhow!(
"Node \"{node_name}\" has Vulnerabilities but none found under Scenario",
node_name = self.0
));
}
}
Ok(())
}
}
impl Connection<Feature> for (&String, &Node) {
fn validate_connections(&self, potential_feature_names: &Option<Vec<String>>) -> Result<()> {
if let Some(node_features) = &self.1.features {
if let Some(features) = potential_feature_names {
for node_feature in node_features.keys() {
if !features.contains(node_feature) {
return Err(anyhow!(
"Node \"{node_name}\" Feature \"{node_feature}\" not found under Scenario Features",
node_name = &self.0,
));
}
}
} else if !node_features.is_empty() {
return Err(anyhow!(
"Node \"{node_name}\" has Features but none found under Scenario",
node_name = &self.0,
));
}
}
Ok(())
}
}
impl Connection<Condition> for (&String, &Node, &Option<Infrastructure>) {
fn validate_connections(&self, potential_condition_names: &Option<Vec<String>>) -> Result<()> {
let (node_name, node, infrastructure) = self;
if let Some(node_conditions) = &node.conditions {
if let Some(conditions) = potential_condition_names {
for condition_name in node_conditions.keys() {
if !conditions.contains(condition_name) {
return Err(anyhow!(
"Node \"{node_name}\" Condition \"{condition_name}\" not found under Scenario Conditions"
));
}
}
if node_conditions.keys().len() > 0 {
if let Some(infrastructure) = infrastructure {
if let Some(infra_node) = infrastructure.get(node_name.to_owned()) {
if infra_node.count > 1 {
return Err(anyhow!(
"Node \"{node_name}\" can not have count bigger than 1, if it has conditions defined"
));
}
}
}
}
} else if !node_conditions.is_empty() {
return Err(anyhow!(
"Node \"{node_name}\" has Conditions but none found under Scenario"
));
}
}
Ok(())
}
}
impl Connection<Node> for (&String, &Option<Roles>) {
fn validate_connections(&self, potential_role_names: &Option<Vec<String>>) -> Result<()> {
if let Some(role_names) = potential_role_names {
if let Some(roles) = self.1 {
for role_name in role_names {
if !roles.contains_key(role_name) {
return Err(anyhow!(
"Role {role_name} not found under for Node {node_name}'s roles",
node_name = self.0
));
}
}
} else {
return Err(anyhow!(
"Roles list is empty for Node {node_name} but it has Role requirements",
node_name = self.0
));
}
}
Ok(())
}
}
impl Connection<Entity> for (&String, &Option<HashMap<String, Role>>) {
fn validate_connections(&self, potential_entity_names: &Option<Vec<String>>) -> Result<()> {
if let Some(node_roles) = self.1 {
for role in node_roles.values() {
if let Some(role_entities) = &role.entities {
if let Some(entity_names) = potential_entity_names {
for role_entity in role_entities {
if !entity_names.contains(role_entity) {
return Err(anyhow!(
"Role Entity {role_entity} for Node {node_name} not found under Entities",
node_name = self.0
));
}
}
} else {
return Err(anyhow!(
"Entities list under Scenario is empty but Node {node_name} has Role Entities",
node_name = self.0
));
}
}
}
}
Ok(())
}
}
impl Formalize for Node {
fn formalize(&mut self) -> Result<()> {
if self.type_field == NodeType::VM {
if let Some(source_helper) = &self._source_helper {
self.source = Some(source_helper.to_owned().into());
} else {
return Err(anyhow::anyhow!("A Node is missing a source field"));
}
if self.resources.is_none() {
return Err(anyhow::anyhow!(
"Nodes of type VM must have Resources defined"
));
}
} else if self.type_field == NodeType::Switch {
if self._source_helper.is_some() {
return Err(anyhow::anyhow!("Nodes of type Switch can not have Sources"));
}
if self.resources.is_some() {
return Err(anyhow::anyhow!(
"Nodes of type Switch can not have Resources"
));
}
}
if let Some(helper_roles) = &self._roles_helper {
self.roles = Some(helper_roles.to_owned().into());
}
Ok(())
}
}
pub type Nodes = HashMap<String, Node>;
#[cfg(test)]
mod tests {
use crate::parse_sdl;
use super::*;
#[test]
fn vm_source_fields_are_mapped_correctly() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
win-10:
type: VM
resources:
ram: 2 gib
cpu: 2
source: windows10
deb-10:
type: VM
resources:
ram: 2 gib
cpu: 2
source:
name: debian10
version: 1.2.3
"#;
let nodes = parse_sdl(sdl).unwrap().nodes;
insta::with_settings!({sort_maps => true}, {
insta::assert_yaml_snapshot!(nodes);
});
}
#[test]
fn vm_source_longhand_is_parsed() {
let longhand_source = r#"
type: VM
source:
name: package-name
version: 1.2.3
"#;
let node = serde_yaml::from_str::<Node>(longhand_source).unwrap();
insta::assert_debug_snapshot!(node);
}
#[test]
fn vm_source_shorthand_is_parsed() {
let shorthand_source = r#"
type: VM
source: package-name
"#;
let node = serde_yaml::from_str::<Node>(shorthand_source).unwrap();
insta::assert_debug_snapshot!(node);
}
#[test]
fn node_conditions_are_parsed() {
let node_sdl = r#"
type: VM
roles:
admin: "username"
conditions:
condition-1: "admin"
"#;
let node = serde_yaml::from_str::<Node>(node_sdl).unwrap();
insta::assert_debug_snapshot!(node);
}
#[test]
fn switch_source_is_not_required() {
let shorthand_source = r#"
type: Switch
"#;
serde_yaml::from_str::<Node>(shorthand_source)
.unwrap()
.formalize()
.unwrap();
}
#[test]
fn includes_node_requirements_with_switch_type() {
let node_sdl = r#"
type: Switch
description: a network switch
"#;
let node = serde_yaml::from_str::<Node>(node_sdl).unwrap();
insta::assert_debug_snapshot!(node);
}
#[test]
fn includes_nodes_with_defined_features() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
win-10:
type: VM
resources:
ram: 2 gib
cpu: 2
source: windows10
roles:
admin: "username"
moderator: "name"
features:
feature-1: "admin"
feature-2: "moderator"
features:
feature-1:
type: service
source: dl-library
feature-2:
type: artifact
source:
name: my-cool-artifact
version: 1.0.0
"#;
let scenario = parse_sdl(sdl).unwrap();
insta::with_settings!({sort_maps => true}, {
insta::assert_yaml_snapshot!(scenario);
});
}
#[test]
#[should_panic(expected = "Roles list is empty for Node win-10 but it has Role requirements")]
fn roles_missing_when_features_exist() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
win-10:
type: VM
source: windows10
resources:
ram: 4 GiB
cpu: 2
features:
feature-1: "admin"
features:
feature-1:
type: service
source: dl-library
"#;
parse_sdl(sdl).unwrap();
}
#[test]
#[should_panic(expected = "Role admin not found under for Node win-10's roles")]
fn role_under_feature_missing_from_node() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
win-10:
type: VM
resources:
ram: 2 gib
cpu: 2
source: windows10
roles:
moderator: "name"
features:
feature-1: "admin"
features:
feature-1:
type: service
source: dl-library
"#;
parse_sdl(sdl).unwrap();
}
#[test]
fn same_name_for_role_only_saves_one_role() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
win-10:
type: VM
resources:
ram: 2 gib
cpu: 2
source: windows10
roles:
admin: "username"
admin: "username2"
"#;
let scenario = parse_sdl(sdl).unwrap();
insta::with_settings!({sort_maps => true}, {
insta::assert_yaml_snapshot!(scenario);
});
}
#[test]
fn nested_node_role_entity_found_under_entities() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
win-10:
type: VM
resources:
cpu: 2
ram: 32 gib
source: windows10
roles:
admin:
username: "admin"
entities:
- blue-team.bob
entities:
blue-team:
name: The Blue Team
entities:
bob:
name: Blue Bob
"#;
parse_sdl(sdl).unwrap();
}
#[test]
#[should_panic(expected = "Role Entity blue-team.bob for Node win-10 not found under Entities")]
fn entity_missing_for_node_role_entity() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
win-10:
type: VM
resources:
cpu: 2
ram: 32 gib
source: windows10
roles:
admin:
username: "admin"
entities:
- blue-team.bob
entities:
blue-team:
name: The Blue Team
"#;
parse_sdl(sdl).unwrap();
}
#[test]
#[should_panic(
expected = "Entities list under Scenario is empty but Node win-10 has Role Entities"
)]
fn entities_missing_while_node_has_role_entity() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
win-10:
type: VM
resources:
cpu: 2
ram: 32 gib
source: windows10
roles:
admin:
username: "admin"
entities:
- blue-team.bob
"#;
parse_sdl(sdl).unwrap();
}
#[test]
fn can_parse_shorthand_node_roles() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
win-10:
type: VM
resources:
cpu: 2
ram: 32 gib
source: windows10
roles:
admin: admin
"#;
let scenario = parse_sdl(sdl).unwrap();
insta::with_settings!({sort_maps => true}, {
insta::assert_yaml_snapshot!(scenario);
});
}
#[test]
fn can_parse_longhand_node_roles() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
win-10:
type: VM
resources:
cpu: 2
ram: 2 gib
source: windows10
roles:
user:
username: user
"#;
let scenario = parse_sdl(sdl).unwrap();
insta::with_settings!({sort_maps => true}, {
insta::assert_yaml_snapshot!(scenario);
});
}
#[test]
fn can_parse_mixed_short_and_longhand_node_roles() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
win-10:
type: VM
resources:
cpu: 2
ram: 2 gib
source: windows10
roles:
admin: admin
user:
username: user
entities:
- blue-team.bob
entities:
blue-team:
name: The Blue Team
entities:
bob:
name: Blue Bob
"#;
let parsed_sdl = parse_sdl(sdl).unwrap();
insta::with_settings!({sort_maps => true}, {
insta::assert_yaml_snapshot!(parsed_sdl);
});
}
#[test]
#[should_panic(expected = "Nodes of type VM must have Resources defined")]
fn resources_missing_for_vm_node() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
win-10:
type: VM
source: windows10
"#;
parse_sdl(sdl).unwrap();
}
#[test]
#[should_panic(expected = "Nodes of type Switch can not have Sources")]
fn source_defined_for_switch_node() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
switch-1:
type: Switch
source: windows10
"#;
parse_sdl(sdl).unwrap();
}
#[test]
#[should_panic(expected = "Nodes of type Switch can not have Resources")]
fn resources_defined_for_switch_node() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
switch-1:
type: Switch
resources:
cpu: 2
ram: 2 gib
"#;
parse_sdl(sdl).unwrap();
}
#[test]
fn node_type_is_case_insensitive() {
let sdl = r#"
name: test-scenario
description: some-description
nodes:
vm-1:
type: vm
source: debian11
resources:
cpu: 2
ram: 2 gib
switch-1:
type: SWITCH
"#;
parse_sdl(sdl).unwrap();
}
}