1use crate::constants::{default_node_count, MINIMUM_NODE_COUNT};
2use crate::helpers::Connection;
3use crate::Formalize;
4use anyhow::{anyhow, Result};
5use ipnetwork::IpNetwork;
6use serde::{Deserialize, Deserializer, Serialize};
7use std::collections::{HashMap, HashSet};
8use std::fmt;
9use std::net::IpAddr;
10
11#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone, Default)]
12pub struct InfraNode {
13 #[serde(default, alias = "Name", alias = "NAME")]
14 pub name: Option<String>,
15 #[serde(default = "default_node_count", alias = "Count", alias = "COUNT")]
16 pub count: i32,
17 #[serde(
18 default,
19 alias = "Links",
20 alias = "LINKS",
21 deserialize_with = "deserialize_unique_list"
22 )]
23 pub links: Option<Vec<String>>,
24 #[serde(
25 default,
26 alias = "Dependencies",
27 alias = "DEPENDENCIES",
28 alias = "dependencies",
29 deserialize_with = "deserialize_unique_list"
30 )]
31 pub dependencies: Option<Vec<String>>,
32 #[serde(
33 default,
34 alias = "Properties",
35 alias = "PROPERTIES",
36 alias = "properties",
37 deserialize_with = "deserialize_properties"
38 )]
39 pub properties: Option<Properties>,
40 #[serde(alias = "Description", alias = "DESCRIPTION", alias = "description")]
41 pub description: Option<String>,
42}
43
44impl InfraNode {
45 pub fn new(potential_count: Option<i32>) -> Self {
46 Self {
47 count: match potential_count {
48 Some(count) => count,
49 None => default_node_count(),
50 },
51 ..Default::default()
52 }
53 }
54}
55
56#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
57#[serde(untagged)]
58pub enum Properties {
59 Simple { cidr: IpNetwork, gateway: IpAddr },
60 Complex(Vec<HashMap<String, IpAddr>>),
61}
62
63#[derive(PartialEq, Eq, Debug, Deserialize, Clone)]
64#[serde(untagged)]
65pub enum HelperNode {
66 Empty,
67 Short(i32),
68 Long(InfraNode),
69 Nested {
70 count: Option<i32>,
71 #[serde(default, deserialize_with = "deserialize_unique_list")]
72 links: Option<Vec<String>>,
73 #[serde(default, deserialize_with = "deserialize_unique_list")]
74 dependencies: Option<Vec<String>>,
75 #[serde(default, deserialize_with = "deserialize_properties")]
76 properties: Option<Properties>,
77 },
78}
79
80#[derive(PartialEq, Eq, Debug, Deserialize, Clone)]
81pub struct InfrastructureHelper(pub HashMap<String, HelperNode>);
82
83impl From<HelperNode> for InfraNode {
84 fn from(helper: HelperNode) -> Self {
85 match helper {
86 HelperNode::Empty => InfraNode::default(),
87 HelperNode::Short(count) => InfraNode {
88 count,
89 ..Default::default()
90 },
91 HelperNode::Long(node) => node,
92 HelperNode::Nested {
93 count,
94 links,
95 dependencies,
96 properties,
97 } => InfraNode {
98 count: count.unwrap_or_else(default_node_count),
99 links,
100 dependencies,
101 properties,
102 ..Default::default()
103 },
104 }
105 }
106}
107
108impl From<InfrastructureHelper> for Infrastructure {
109 fn from(helper_infrastructure: InfrastructureHelper) -> Self {
110 helper_infrastructure
111 .0
112 .into_iter()
113 .map(|(node_name, helper_node)| {
114 let mut infra_node: InfraNode = helper_node.into();
115 infra_node.formalize().unwrap();
116 (node_name, infra_node)
117 })
118 .collect()
119 }
120}
121
122pub type Infrastructure = HashMap<String, InfraNode>;
123
124impl Connection<Infrastructure> for String {
125 fn validate_connections(&self, potential_node_names: &Option<Vec<String>>) -> Result<()> {
126 if let Some(node_names) = potential_node_names {
127 if !node_names.contains(self) {
128 return Err(anyhow!(
129 "Infrastructure entry '{}' does not exist under Nodes",
130 self
131 ));
132 }
133 }
134 Ok(())
135 }
136}
137
138impl Formalize for InfraNode {
139 fn formalize(&mut self) -> Result<()> {
140 if self.count < MINIMUM_NODE_COUNT {
141 return Err(anyhow!(
142 "Infrastructure Count field cannot be less than {MINIMUM_NODE_COUNT}"
143 ));
144 }
145 Ok(())
146 }
147}
148
149fn deserialize_unique_list<'de, D>(deserializer: D) -> Result<Option<Vec<String>>, D::Error>
150where
151 D: Deserializer<'de>,
152{
153 struct UniqueListVisitor;
154
155 impl<'de> serde::de::Visitor<'de> for UniqueListVisitor {
156 type Value = Vec<String>;
157
158 fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
159 formatter.write_str("a list of unique strings")
160 }
161
162 fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
163 where
164 A: serde::de::SeqAccess<'de>,
165 {
166 let mut items = HashSet::new();
167 let mut result = Vec::new();
168
169 while let Some(value) = seq.next_element::<String>()? {
170 if items.contains(&value) {
171 return Err(serde::de::Error::custom(format!(
172 "duplicate value found: {}",
173 value
174 )));
175 }
176 items.insert(value.clone());
177 result.push(value);
178 }
179
180 Ok(result)
181 }
182 }
183
184 deserializer.deserialize_seq(UniqueListVisitor).map(Some)
185}
186
187fn deserialize_properties<'de, D>(deserializer: D) -> Result<Option<Properties>, D::Error>
188where
189 D: Deserializer<'de>,
190{
191 #[derive(Deserialize)]
192 struct SimpleProperties {
193 cidr: IpNetwork,
194 gateway: IpAddr,
195 }
196
197 #[derive(Deserialize)]
198 #[serde(untagged)]
199 enum RawProperties {
200 Simple(SimpleProperties),
201 Complex(Vec<HashMap<String, String>>),
202 }
203
204 match RawProperties::deserialize(deserializer)? {
205 RawProperties::Simple(simple) => Ok(Some(Properties::Simple {
206 cidr: simple.cidr,
207 gateway: simple.gateway,
208 })),
209 RawProperties::Complex(complex) => {
210 let list = complex
211 .into_iter()
212 .map(|map| {
213 map.into_iter()
214 .map(|(key, value)| {
215 let ip = value.parse::<IpAddr>().map_err(serde::de::Error::custom)?;
216 Ok((key, ip))
217 })
218 .collect::<Result<HashMap<String, IpAddr>, _>>()
219 })
220 .collect::<Result<Vec<HashMap<String, IpAddr>>, _>>()?;
221 Ok(Some(Properties::Complex(list)))
222 }
223 }
224}
225
226#[cfg(test)]
227mod tests {
228 use super::*;
229 use crate::parse_sdl;
230
231 #[test]
232 fn infranode_count_longhand_is_parsed() {
233 let sdl = r#"
234 count: 23
235 "#;
236 let infra_node = serde_yaml::from_str::<InfraNode>(sdl).unwrap();
237 insta::assert_debug_snapshot!(infra_node);
238 }
239
240 #[test]
241 fn infranode_count_shorthand_is_parsed() {
242 let sdl = r#"
243 23
244 "#;
245 let infra_node: InfraNode = serde_yaml::from_str::<HelperNode>(sdl).unwrap().into();
246 insta::assert_debug_snapshot!(infra_node);
247 }
248
249 #[test]
250 fn infranode_with_links_and_dependencies_is_parsed() {
251 let sdl = r#"
252 count: 25
253 links:
254 - switch-2
255 dependencies:
256 - windows-10
257 - windows-10-vuln-1
258 "#;
259 let infra_node = serde_yaml::from_str::<InfraNode>(sdl).unwrap();
260 insta::assert_debug_snapshot!(infra_node);
261 }
262
263 #[test]
264 fn infranode_with_default_count_is_parsed() {
265 let sdl = r#"
266 links:
267 - switch-1
268 dependencies:
269 - windows-10
270 - windows-10-vuln-1
271 "#;
272 let infra_node = serde_yaml::from_str::<InfraNode>(sdl).unwrap();
273 insta::assert_debug_snapshot!(infra_node);
274 }
275
276 #[test]
277 fn simple_infrastructure_is_parsed() {
278 let sdl = r#"
279 windows-10-vuln-1:
280 count: 10
281 description: "A vulnerable Windows 10 machine"
282 debian-2:
283 count: 4
284 description: "A Debian server"
285 "#;
286 let infrastructure = serde_yaml::from_str::<Infrastructure>(sdl).unwrap();
287 insta::with_settings!({sort_maps => true}, {
288 insta::assert_yaml_snapshot!(infrastructure);
289 });
290 }
291
292 #[test]
293 fn simple_infrastructure_with_shorthand_is_parsed() {
294 let sdl = r#"
295 windows-10-vuln-2:
296 count: 10
297 windows-10-vuln-1: 10
298 ubuntu-10: 5
299 "#;
300 let infrastructure_helper = serde_yaml::from_str::<InfrastructureHelper>(sdl).unwrap();
301 let infrastructure: Infrastructure = infrastructure_helper.into();
302
303 insta::with_settings!({sort_maps => true}, {
304 insta::assert_yaml_snapshot!(infrastructure);
305 });
306 }
307
308 #[test]
309 fn bigger_infrastructure_is_parsed() {
310 let sdl = r#"
311 switch-1: 1
312 windows-10: 3
313 windows-10-vuln-1:
314 count: 1
315 switch-2:
316 count: 2
317 links:
318 - switch-1
319 ubuntu-10:
320 links:
321 - switch-1
322 dependencies:
323 - windows-10
324 - windows-10-vuln-1
325 "#;
326 let infrastructure_helper = serde_yaml::from_str::<InfrastructureHelper>(sdl).unwrap();
327 let infrastructure: Infrastructure = infrastructure_helper.into();
328
329 insta::with_settings!({sort_maps => true}, {
330 insta::assert_yaml_snapshot!(infrastructure);
331 });
332 }
333
334 #[test]
335 fn sdl_keys_are_valid_in_lowercase_uppercase_capitalized() {
336 let sdl = r#"
337 switch-1: 1
338 windows-10: 3
339 windows-10-vuln-1:
340 Count: 1
341 DEPENDENCIES:
342 - windows-10
343 switch-2:
344 COUNT: 2
345 Links:
346 - switch-1
347 ubuntu-10:
348 LINKS:
349 - switch-1
350 Dependencies:
351 - windows-10
352 - windows-10-vuln-1
353 "#;
354 let infrastructure_helper = serde_yaml::from_str::<InfrastructureHelper>(sdl).unwrap();
355 let infrastructure: Infrastructure = infrastructure_helper.into();
356
357 insta::with_settings!({sort_maps => true}, {
358 insta::assert_yaml_snapshot!(infrastructure);
359 });
360 }
361
362 #[test]
363 fn infrastructure_with_links_and_properties() {
364 let sdl = r#"
365switch-1:
366 count: 1
367 properties:
368 cidr: 10.10.10.0/24
369 gateway: 10.10.10.1
370windows-10: 3
371windows-10-vuln-1:
372 count: 1
373 links:
374 - switch-1
375 properties:
376 - switch-1: 10.10.10.10
377switch-2:
378 count: 2
379 links:
380 - switch-1
381ubuntu-10:
382 links:
383 - switch-1
384 dependencies:
385 - windows-10
386 - windows-10-vuln-1
387
388 "#;
389
390 let infrastructure_helper = serde_yaml::from_str::<InfrastructureHelper>(sdl).unwrap();
391 let infrastructure: Infrastructure = infrastructure_helper.into();
392
393 insta::with_settings!({sort_maps => true}, {
394 insta::assert_yaml_snapshot!(infrastructure);
395 });
396 }
397
398 #[test]
399 fn empty_count_is_allowed() {
400 let sdl = r#"
401 switch-1:
402 "#;
403 serde_yaml::from_str::<InfrastructureHelper>(sdl).unwrap();
404 }
405
406 #[should_panic(expected = "Infrastructure Count field cannot be less than 1")]
407 #[test]
408 fn infranode_with_negative_count_is_rejected() {
409 let sdl = r#"
410 name: test-scenario
411 description: some-description
412 nodes:
413 win-10:
414 type: VM
415 resources:
416 ram: 2 gib
417 cpu: 2
418 source: windows10
419 infrastructure:
420 win-10: -1
421 "#;
422 parse_sdl(sdl).unwrap();
423 }
424
425 #[should_panic(expected = "Infrastructure entry \"debian\" does not exist under Nodes")]
426 #[test]
427 fn infranode_with_unknown_name_is_rejected() {
428 let sdl = r#"
429 name: test-scenario
430 description: some-description
431 nodes:
432 win-10:
433 type: VM
434 resources:
435 ram: 2 gib
436 cpu: 2
437 source: windows10
438 infrastructure:
439 debian: 1
440 "#;
441 parse_sdl(sdl).unwrap();
442 }
443
444 #[should_panic(
445 expected = "Infrastructure entry \"main-switch\" does not exist under Infrastructure even though it is a dependency for \"win-10\""
446 )]
447 #[test]
448 fn error_on_missing_infrastructure_link() {
449 let sdl = r#"
450 name: test-scenario
451 description: some-description
452 nodes:
453 win-10:
454 type: VM
455 resources:
456 ram: 2 gib
457 cpu: 2
458 source: windows10
459 main-switch:
460 type: Switch
461 infrastructure:
462 win-10:
463 count: 1
464 links:
465 - main-switch
466 "#;
467
468 parse_sdl(sdl).unwrap();
469 }
470
471 #[should_panic(
472 expected = "Infrastructure entry \"router\" does not exist under Infrastructure even though it is a dependency for \"win-10\""
473 )]
474 #[test]
475 fn error_on_missing_infrastructure_dependency() {
476 let sdl = r#"
477 name: test-scenario
478 description: some-description
479 nodes:
480 win-10:
481 type: VM
482 resources:
483 ram: 2 gib
484 cpu: 2
485 source: windows10
486 router:
487 type: VM
488 resources:
489 ram: 2 gib
490 cpu: 2
491 source: debian11
492 infrastructure:
493 win-10:
494 count: 1
495 dependencies:
496 - router
497 "#;
498
499 parse_sdl(sdl).unwrap();
500 }
501
502 #[should_panic(
503 expected = "IP address '10.20.10.10' for 'main-switch' in properties of node 'win-10' is not within the CIDR '10.10.10.0/24' of the linked node 'main-switch'"
504 )]
505 #[test]
506 fn error_on_wrong_ip_in_infranode_properties() {
507 let sdl = r#"
508 name: test-scenario
509 description: some-description
510 nodes:
511 win-10:
512 type: VM
513 resources:
514 ram: 2 gib
515 cpu: 2
516 source: windows10
517 main-switch:
518 type: Switch
519 infrastructure:
520 main-switch:
521 count: 1
522 properties:
523 cidr: 10.10.10.0/24
524 gateway: 10.10.10.1
525 win-10:
526 count: 1
527 links:
528 - main-switch
529 properties:
530 - main-switch: 10.20.10.10
531 "#;
532
533 parse_sdl(sdl).unwrap();
534 }
535}