1use crate::{
2 condition::Condition, entity::Entity, feature::Feature, helpers::Connection,
3 infrastructure::Infrastructure, inject::Inject, vulnerability::Vulnerability, Formalize,
4};
5use anyhow::{anyhow, Result};
6use bytesize::ByteSize;
7use serde::{Deserialize, Deserializer, Serialize};
8use serde_aux::prelude::*;
9use std::collections::HashMap;
10
11use crate::common::{HelperSource, Source};
12
13fn parse_bytesize<'de, D>(deserializer: D) -> Result<u64, D::Error>
14where
15 D: Deserializer<'de>,
16{
17 let s = String::deserialize(deserializer)?;
18 Ok(s.parse::<ByteSize>()
19 .map_err(|_| serde::de::Error::custom("Failed"))?
20 .0)
21}
22
23#[allow(clippy::large_enum_variant)]
24#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone)]
25#[serde(tag = "type")]
26pub enum NodeType {
27 #[serde(alias = "SWITCH", alias = "switch", alias = "Switch")]
28 Switch(Switch),
29 #[serde(alias = "VM", alias = "vm", alias = "Vm", alias = "vM")]
30 VM(VM),
31}
32
33#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone)]
34#[serde(deny_unknown_fields)]
35pub struct Switch {}
36
37#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone)]
38pub struct VM {
39 #[serde(default, skip_deserializing)]
40 pub source: Option<Source>,
41 #[serde(
42 default,
43 rename = "source",
44 alias = "Source",
45 alias = "SOURCE",
46 skip_serializing
47 )]
48 _source_helper: Option<HelperSource>,
49 #[serde(
50 alias = "Resources",
51 alias = "RESOURCES",
52 deserialize_with = "deserialize_struct_case_insensitive"
53 )]
54 pub resources: Resources,
55 #[serde(default, alias = "Features", alias = "FEATURES")]
56 pub features: HashMap<String, String>,
57 #[serde(default, alias = "Conditions", alias = "CONDITIONS")]
58 pub conditions: HashMap<String, String>,
59 #[serde(default, alias = "Injects", alias = "INJECTS")]
60 pub injects: HashMap<String, String>,
61 #[serde(default, alias = "Vulnerabilities", alias = "VULNERABILITIES")]
62 pub vulnerabilities: Vec<String>,
63 #[serde(
64 default,
65 rename = "roles",
66 alias = "Roles",
67 alias = "ROLES",
68 skip_serializing
69 )]
70 _roles_helper: Option<HelperRoles>,
71 #[serde(skip_deserializing)]
72 pub roles: Option<Roles>,
73}
74
75#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
76pub struct Resources {
77 #[serde(deserialize_with = "parse_bytesize")]
78 pub ram: u64,
79 pub cpu: u32,
80}
81
82#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
83pub struct Role {
84 #[serde(alias = "Username", alias = "USERNAME")]
85 pub username: String,
86 #[serde(alias = "Entity", alias = "ENTITY")]
87 pub entities: Option<Vec<String>>,
88}
89
90pub type Roles = HashMap<String, Role>;
91
92#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
93pub struct Node {
94 #[serde(flatten)]
95 pub type_field: NodeType,
96 #[serde(alias = "Description", alias = "DESCRIPTION")]
97 pub description: Option<String>,
98}
99
100#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
101#[serde(untagged)]
102pub enum RoleTypes {
103 Username(String),
104 Role(Role),
105}
106#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
107#[serde(untagged)]
108pub enum HelperRoles {
109 MixedRoles(HashMap<String, RoleTypes>),
110}
111
112impl From<HelperRoles> for Roles {
113 fn from(helper_role: HelperRoles) -> Self {
114 match helper_role {
115 HelperRoles::MixedRoles(mixed_role) => mixed_role
116 .into_iter()
117 .map(|(role_name, role_value)| {
118 let role_value = match role_value {
119 RoleTypes::Role(role) => role,
120 RoleTypes::Username(username) => Role {
121 username,
122 entities: None,
123 },
124 };
125 (role_name, role_value)
126 })
127 .collect::<Roles>(),
128 }
129 }
130}
131
132impl Connection<Vulnerability> for (&String, &VM) {
133 fn validate_connections(
134 &self,
135 potential_vulnerability_names: &Option<Vec<String>>,
136 ) -> Result<()> {
137 let node_vulnerabilities = &self.1.vulnerabilities;
138
139 if !node_vulnerabilities.is_empty() {
140 if let Some(vulnerabilities) = potential_vulnerability_names {
141 for vulnerability_name in node_vulnerabilities.iter() {
142 if !vulnerabilities.contains(vulnerability_name) {
143 return Err(anyhow!(
144 "Vulnerability \"{vulnerability_name}\" not found under Scenario Vulnerabilities",
145 ));
146 }
147 }
148 } else {
149 return Err(anyhow!(
150 "Node \"{node_name}\" has Vulnerabilities but none found under Scenario",
151 node_name = self.0
152 ));
153 }
154 }
155
156 Ok(())
157 }
158}
159
160impl Connection<Feature> for (&String, &VM) {
161 fn validate_connections(&self, potential_feature_names: &Option<Vec<String>>) -> Result<()> {
162 let vm_features = &self.1.features;
163
164 if !vm_features.is_empty() {
165 if let Some(features) = potential_feature_names {
166 for node_feature in vm_features.keys() {
167 if !features.contains(node_feature) {
168 return Err(anyhow!(
169 "VM \"{node_name}\" Feature \"{node_feature}\" not found under Scenario Features",
170 node_name = &self.0,
171 ));
172 }
173 }
174 } else if !vm_features.is_empty() {
175 return Err(anyhow!(
176 "VM \"{node_name}\" has Features but none found under Scenario",
177 node_name = &self.0,
178 ));
179 }
180 }
181 Ok(())
182 }
183}
184
185impl Connection<Condition> for (&String, &VM, &Option<Infrastructure>) {
186 fn validate_connections(&self, potential_condition_names: &Option<Vec<String>>) -> Result<()> {
187 let (node_name, node, infrastructure) = self;
188 let vm_conditions = &node.conditions;
189
190 if let Some(conditions) = potential_condition_names {
191 for condition_name in vm_conditions.keys() {
192 if !conditions.contains(condition_name) {
193 return Err(anyhow!(
194 "Node \"{node_name}\" Condition \"{condition_name}\" not found under Scenario Conditions"
195 ));
196 }
197 }
198 if vm_conditions.keys().len() > 0 {
199 if let Some(infrastructure) = infrastructure {
200 if let Some(infra_node) = infrastructure.get(node_name.to_owned()) {
201 if infra_node.count > 1 {
202 return Err(anyhow!(
203 "Node \"{node_name}\" can not have count bigger than 1, if it has conditions defined"
204 ));
205 }
206 }
207 }
208 }
209 } else if !vm_conditions.is_empty() {
210 return Err(anyhow!(
211 "Node \"{node_name}\" has Conditions but none found under Scenario"
212 ));
213 }
214
215 Ok(())
216 }
217}
218
219impl Connection<Inject> for (&String, &VM, &Option<Infrastructure>) {
220 fn validate_connections(&self, potential_inject_names: &Option<Vec<String>>) -> Result<()> {
221 let (node_name, node, infrastructure) = self;
222 let vm_injects = &node.injects;
223
224 if let Some(injects) = potential_inject_names {
225 for inject_name in vm_injects.keys() {
226 if !injects.contains(inject_name) {
227 return Err(anyhow!(
228 "Node \"{node_name}\" Inject \"{inject_name}\" not found under Scenario Injects"
229 ));
230 }
231 }
232 if !vm_injects.is_empty() {
233 if let Some(infrastructure) = infrastructure {
234 if let Some(infra_node) = infrastructure.get(node_name.to_owned()) {
235 if infra_node.count > 1 {
236 return Err(anyhow!(
237 "Node \"{node_name}\" can not have count bigger than 1, if it has injects defined"
238 ));
239 }
240 }
241 }
242 }
243 } else if !vm_injects.is_empty() {
244 return Err(anyhow!(
245 "Node \"{node_name}\" has Injects but none found under Scenario"
246 ));
247 }
248
249 Ok(())
250 }
251}
252
253impl Connection<Node> for (&String, &Option<Roles>) {
254 fn validate_connections(&self, potential_role_names: &Option<Vec<String>>) -> Result<()> {
255 if let Some(role_names) = potential_role_names {
256 if !role_names.is_empty() {
257 if let Some(roles) = self.1 {
258 for role_name in role_names {
259 if !roles.contains_key(role_name) {
260 return Err(anyhow!(
261 "Role {role_name} not found under for Node {node_name}'s roles",
262 node_name = self.0
263 ));
264 }
265 }
266 } else {
267 return Err(anyhow!(
268 "Roles list is empty for Node {node_name} but it has Role requirements",
269 node_name = self.0
270 ));
271 }
272 }
273 }
274
275 Ok(())
276 }
277}
278
279impl Connection<Entity> for (&String, &Option<HashMap<String, Role>>) {
280 fn validate_connections(&self, potential_entity_names: &Option<Vec<String>>) -> Result<()> {
281 if let Some(node_roles) = self.1 {
282 for role in node_roles.values() {
283 if let Some(role_entities) = &role.entities {
284 if let Some(entity_names) = potential_entity_names {
285 for role_entity in role_entities {
286 if !entity_names.contains(role_entity) {
287 return Err(anyhow!(
288 "Role Entity {role_entity} for Node {node_name} not found under Entities",
289 node_name = self.0
290 ));
291 }
292 }
293 } else {
294 return Err(anyhow!(
295 "Entities list under Scenario is empty but Node {node_name} has Role Entities",
296 node_name = self.0
297 ));
298 }
299 }
300 }
301 }
302
303 Ok(())
304 }
305}
306
307impl Formalize for VM {
308 fn formalize(&mut self) -> Result<()> {
309 if let Some(source_helper) = &self._source_helper {
310 self.source = Some(source_helper.to_owned().into());
311 } else {
312 return Err(anyhow::anyhow!("A Node is missing a source field"));
313 }
314
315 if let Some(helper_roles) = &self._roles_helper {
316 self.roles = Some(helper_roles.to_owned().into());
317 }
318
319 Ok(())
320 }
321}
322
323pub type Nodes = HashMap<String, Node>;
324
325#[cfg(test)]
326mod tests {
327 use crate::parse_sdl;
328
329 use super::*;
330
331 #[test]
332 fn vm_source_fields_are_mapped_correctly() {
333 let sdl = r#"
334 name: test-scenario
335 description: some-description
336 nodes:
337 win-10:
338 type: VM
339 resources:
340 ram: 2 gib
341 cpu: 2
342 source: windows10
343 deb-10:
344 type: VM
345 resources:
346 ram: 2 gib
347 cpu: 2
348 source:
349 name: debian10
350 version: 1.2.3
351
352 "#;
353 let nodes = parse_sdl(sdl).unwrap().nodes;
354 insta::with_settings!({sort_maps => true}, {
355 insta::assert_yaml_snapshot!(nodes);
356 });
357 }
358
359 #[test]
360 fn vm_source_longhand_is_parsed() {
361 let longhand_source = r#"
362 type: VM
363 source:
364 name: package-name
365 version: 1.2.3
366 resources:
367 cpu: 2
368 ram: 2GB
369 "#;
370 let node = serde_yaml::from_str::<Node>(longhand_source).unwrap();
371 insta::assert_debug_snapshot!(node);
372 }
373
374 #[test]
375 fn vm_source_shorthand_is_parsed() {
376 let shorthand_source = r#"
377 type: VM
378 source: package-name
379 resources:
380 cpu: 2
381 ram: 2GB
382 "#;
383 let node = serde_yaml::from_str::<Node>(shorthand_source).unwrap();
384 insta::assert_debug_snapshot!(node);
385 }
386
387 #[test]
388 fn node_conditions_are_parsed() {
389 let node_sdl = r#"
390 type: VM
391 roles:
392 admin: "username"
393 conditions:
394 condition-1: "admin"
395 resources:
396 cpu: 2
397 ram: 2GB
398 "#;
399 let node = serde_yaml::from_str::<Node>(node_sdl).unwrap();
400 insta::assert_debug_snapshot!(node);
401 }
402
403 #[test]
404 fn node_injects_are_parsed() {
405 let node_sdl = r#"
406 type: VM
407 roles:
408 admin: "username"
409 injects:
410 inject-1: "admin"
411 resources:
412 cpu: 2
413 ram: 2GB
414 "#;
415 let node = serde_yaml::from_str::<Node>(node_sdl).unwrap();
416 insta::assert_debug_snapshot!(node);
417 }
418
419 #[test]
420 fn switch_source_is_not_required() {
421 let shorthand_source = r#"
422 type: Switch
423 "#;
424 serde_yaml::from_str::<Node>(shorthand_source).unwrap();
425 }
426
427 #[test]
428 fn includes_node_requirements_with_switch_type() {
429 let node_sdl = r#"
430 type: Switch
431 description: a network switch
432
433 "#;
434 let node = serde_yaml::from_str::<Node>(node_sdl).unwrap();
435 insta::assert_debug_snapshot!(node);
436 }
437
438 #[test]
439 fn includes_nodes_with_defined_features() {
440 let sdl = r#"
441 name: test-scenario
442 description: some-description
443 nodes:
444 win-10:
445 type: VM
446 resources:
447 ram: 2 gib
448 cpu: 2
449 source: windows10
450 roles:
451 admin: "username"
452 moderator: "name"
453 features:
454 feature-1: "admin"
455 feature-2: "moderator"
456 features:
457 feature-1:
458 type: service
459 source: dl-library
460 feature-2:
461 type: artifact
462 source:
463 name: my-cool-artifact
464 version: 1.0.0
465
466 "#;
467 let scenario = parse_sdl(sdl).unwrap();
468 insta::with_settings!({sort_maps => true}, {
469 insta::assert_yaml_snapshot!(scenario);
470 });
471 }
472
473 #[test]
474 fn includes_nodes_with_defined_injects() {
475 let sdl = r#"
476 name: test-scenario
477 description: some-description
478 nodes:
479 win-10:
480 type: VM
481 resources:
482 ram: 2 gib
483 cpu: 2
484 source: windows10
485 roles:
486 admin: "username"
487 moderator: "name"
488 injects:
489 inject-1: "admin"
490 inject-2: "moderator"
491 injects:
492 inject-1:
493 source: dl-library
494 inject-2:
495 source: dl-library
496 "#;
497 let scenario = parse_sdl(sdl).unwrap();
498 insta::with_settings!({sort_maps => true}, {
499 insta::assert_yaml_snapshot!(scenario);
500 });
501 }
502
503 #[test]
504 #[should_panic(expected = "Roles list is empty for Node win-10 but it has Role requirements")]
505 fn roles_missing_when_features_exist() {
506 let sdl = r#"
507 name: test-scenario
508 description: some-description
509 nodes:
510 win-10:
511 type: VM
512 source: windows10
513 resources:
514 ram: 4 GiB
515 cpu: 2
516 features:
517 feature-1: "admin"
518 features:
519 feature-1:
520 type: service
521 source: dl-library
522
523 "#;
524 parse_sdl(sdl).unwrap();
525 }
526
527 #[test]
528 #[should_panic(expected = "Role admin not found under for Node win-10's roles")]
529 fn role_under_feature_missing_from_node() {
530 let sdl = r#"
531 name: test-scenario
532 description: some-description
533 nodes:
534 win-10:
535 type: VM
536 resources:
537 ram: 2 gib
538 cpu: 2
539 source: windows10
540 roles:
541 moderator: "name"
542 features:
543 feature-1: "admin"
544 features:
545 feature-1:
546 type: service
547 source: dl-library
548
549 "#;
550 parse_sdl(sdl).unwrap();
551 }
552
553 #[test]
554 fn same_name_for_role_only_saves_one_role() {
555 let sdl = r#"
556 name: test-scenario
557 description: some-description
558 nodes:
559 win-10:
560 type: VM
561 resources:
562 ram: 2 gib
563 cpu: 2
564 source: windows10
565 roles:
566 admin: "username"
567 admin: "username2"
568
569 "#;
570 let scenario = parse_sdl(sdl).unwrap();
571 insta::with_settings!({sort_maps => true}, {
572 insta::assert_yaml_snapshot!(scenario);
573 });
574 }
575
576 #[test]
577 fn nested_node_role_entity_found_under_entities() {
578 let sdl = r#"
579 name: test-scenario
580 description: some-description
581 nodes:
582 win-10:
583 type: VM
584 resources:
585 cpu: 2
586 ram: 32 gib
587 source: windows10
588 roles:
589 admin:
590 username: "admin"
591 entities:
592 - blue-team.bob
593 entities:
594 blue-team:
595 name: The Blue Team
596 entities:
597 bob:
598 name: Blue Bob
599 "#;
600 parse_sdl(sdl).unwrap();
601 }
602
603 #[test]
604 #[should_panic(expected = "Role Entity blue-team.bob for Node win-10 not found under Entities")]
605 fn entity_missing_for_node_role_entity() {
606 let sdl = r#"
607 name: test-scenario
608 description: some-description
609 nodes:
610 win-10:
611 type: VM
612 resources:
613 cpu: 2
614 ram: 32 gib
615 source: windows10
616 roles:
617 admin:
618 username: "admin"
619 entities:
620 - blue-team.bob
621 entities:
622 blue-team:
623 name: The Blue Team
624 "#;
625 parse_sdl(sdl).unwrap();
626 }
627
628 #[test]
629 #[should_panic(
630 expected = "Entities list under Scenario is empty but Node win-10 has Role Entities"
631 )]
632 fn entities_missing_while_node_has_role_entity() {
633 let sdl = r#"
634 name: test-scenario
635 description: some-description
636 nodes:
637 win-10:
638 type: VM
639 resources:
640 cpu: 2
641 ram: 32 gib
642 source: windows10
643 roles:
644 admin:
645 username: "admin"
646 entities:
647 - blue-team.bob
648 "#;
649 parse_sdl(sdl).unwrap();
650 }
651
652 #[test]
653 fn can_parse_shorthand_node_roles() {
654 let sdl = r#"
655 name: test-scenario
656 description: some-description
657 nodes:
658 win-10:
659 type: VM
660 resources:
661 cpu: 2
662 ram: 32 gib
663 source: windows10
664 roles:
665 admin: admin
666 "#;
667 let scenario = parse_sdl(sdl).unwrap();
668 insta::with_settings!({sort_maps => true}, {
669 insta::assert_yaml_snapshot!(scenario);
670 });
671 }
672 #[test]
673 fn can_parse_longhand_node_roles() {
674 let sdl = r#"
675 name: test-scenario
676 description: some-description
677 nodes:
678 win-10:
679 type: VM
680 resources:
681 cpu: 2
682 ram: 2 gib
683 source: windows10
684 roles:
685 user:
686 username: user
687 "#;
688 let scenario = parse_sdl(sdl).unwrap();
689 insta::with_settings!({sort_maps => true}, {
690 insta::assert_yaml_snapshot!(scenario);
691 });
692 }
693 #[test]
694 fn can_parse_mixed_short_and_longhand_node_roles() {
695 let sdl = r#"
696 name: test-scenario
697 description: some-description
698 nodes:
699 win-10:
700 type: VM
701 resources:
702 cpu: 2
703 ram: 2 gib
704 source: windows10
705 roles:
706 admin: admin
707 user:
708 username: user
709 entities:
710 - blue-team.bob
711
712 entities:
713 blue-team:
714 name: The Blue Team
715 entities:
716 bob:
717 name: Blue Bob
718 "#;
719 let parsed_sdl = parse_sdl(sdl).unwrap();
720 insta::with_settings!({sort_maps => true}, {
721 insta::assert_yaml_snapshot!(parsed_sdl);
722 });
723 }
724
725 #[test]
726 #[should_panic(expected = "missing field `resources`")]
727 fn resources_missing_for_vm_node() {
728 let sdl = r#"
729 name: test-scenario
730 description: some-description
731 nodes:
732 win-10:
733 type: VM
734 source: windows10
735 "#;
736 parse_sdl(sdl).unwrap();
737 }
738
739 #[test]
740 #[should_panic(expected = "unknown field `source`")]
741 fn source_defined_for_switch_node() {
742 let sdl = r#"
743 name: test-scenario
744 description: some-description
745 nodes:
746 switch-1:
747 type: Switch
748 source: windows10
749 "#;
750 parse_sdl(sdl).unwrap();
751 }
752
753 #[test]
754 #[should_panic(expected = "unknown field `resources`")]
755 fn resources_defined_for_switch_node() {
756 let sdl = r#"
757 name: test-scenario
758 description: some-description
759 nodes:
760 switch-1:
761 type: Switch
762 resources:
763 cpu: 2
764 ram: 2 gib
765 "#;
766 parse_sdl(sdl).unwrap();
767 }
768
769 #[test]
770 fn node_type_is_case_insensitive() {
771 let sdl = r#"
772 name: test-scenario
773 description: some-description
774 nodes:
775 vm-1:
776 type: vm
777 source: debian11
778 resources:
779 cpu: 2
780 ram: 2 gib
781 switch-1:
782 type: SWITCH
783 "#;
784 parse_sdl(sdl).unwrap();
785 }
786}