1use anyhow::{anyhow, Result};
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4
5use crate::{
6 common::{HelperSource, Source},
7 entity::Entity,
8 helpers::Connection,
9 training_learning_objective::TrainingLearningObjective,
10 Formalize,
11};
12
13#[derive(PartialEq, Eq, Debug, Serialize, Deserialize, Clone)]
14pub struct Inject {
15 #[serde(default, alias = "Name", alias = "NAME")]
16 pub name: Option<String>,
17 #[serde(
18 default,
19 rename = "source",
20 alias = "Source",
21 alias = "SOURCE",
22 skip_serializing
23 )]
24 source_helper: Option<HelperSource>,
25 #[serde(default, skip_deserializing)]
26 pub source: Option<Source>,
27 #[serde(rename = "from-entity", alias = "From-entity", alias = "FROM-ENTITY")]
28 pub from_entity: Option<String>,
29 #[serde(rename = "to-entities", alias = "To-entities", alias = "TO-ENTITIES")]
30 pub to_entities: Option<Vec<String>>,
31 #[serde(alias = "Tlos", alias = "TLOS")]
32 pub tlos: Option<Vec<String>>,
33 #[serde(alias = "Description", alias = "DESCRIPTION")]
34 pub description: Option<String>,
35 #[serde(alias = "Environment", alias = "ENVIRONMENT")]
36 pub environment: Option<Vec<String>>,
37}
38
39pub type Injects = HashMap<String, Inject>;
40
41impl Formalize for Inject {
42 fn formalize(&mut self) -> Result<()> {
43 if self.from_entity.is_some() && self.to_entities.is_none() {
44 return Err(anyhow!(
45 "Inject must have `to-entities` declared if `from-entity` is declared"
46 ));
47 } else if self.from_entity.is_none() && self.to_entities.is_some() {
48 return Err(anyhow!(
49 "Inject must have `from-entity` declared if `to-entities` is declared"
50 ));
51 } else if let Some(source_helper) = &self.source_helper {
52 self.source = Some(source_helper.to_owned().into());
53 }
54 Ok(())
55 }
56}
57
58impl Connection<Entity> for (&String, &Inject) {
59 fn validate_connections(&self, potential_entity_names: &Option<Vec<String>>) -> Result<()> {
60 if self.1.to_entities.is_some() && potential_entity_names.is_none()
61 || self.1.from_entity.is_some() && potential_entity_names.is_none()
62 {
63 return Err(anyhow!(
64 "Inject \"{inject_name}\" has Entities but none found under Scenario",
65 inject_name = self.0
66 ));
67 }
68
69 let mut required_entities: Vec<String> = vec![];
70
71 if let Some(from_entity) = self.1.clone().from_entity {
72 required_entities.push(from_entity);
73 if let Some(to_entities) = self.1.clone().to_entities {
74 required_entities.extend_from_slice(to_entities.as_slice());
75 }
76 }
77 for inject_entity_name in required_entities.iter() {
78 if let Some(scenario_entities) = potential_entity_names {
79 if !scenario_entities.contains(inject_entity_name) {
80 return Err(anyhow!(
81 "Inject \"{inject_name}\" Entity \"{inject_entity_name}\" not found under Scenario Entities",
82 inject_name = self.0
83 ));
84 }
85 }
86 }
87
88 Ok(())
89 }
90}
91
92impl Connection<TrainingLearningObjective> for (&String, &Inject) {
93 fn validate_connections(&self, potential_tlo_names: &Option<Vec<String>>) -> Result<()> {
94 if self.1.tlos.is_some() && potential_tlo_names.is_none() {
95 return Err(anyhow!(
96 "Inject \"{inject_name}\" has TLOs but none found under Scenario",
97 inject_name = self.0
98 ));
99 }
100
101 if let Some(required_tlos) = &self.1.tlos {
102 if let Some(tlo_names) = potential_tlo_names {
103 for tlo_name in required_tlos {
104 if !tlo_names.contains(tlo_name) {
105 return Err(anyhow!("Inject \"{inject_name}\" TLO \"{tlo_name}\" not found under Scenario TLOs",
106 inject_name = self.0
107 ));
108 }
109 }
110 }
111 }
112
113 Ok(())
114 }
115}
116
117#[cfg(test)]
118mod tests {
119 use super::*;
120 use crate::parse_sdl;
121
122 #[test]
123 fn parses_sdl_with_injects() {
124 let sdl = r#"
125 name: test-scenario
126 description: some description
127 start: 2022-01-20T13:00:00Z
128 end: 2022-01-20T23:00:00Z
129 conditions:
130 condition-1:
131 command: executable/path.sh
132 interval: 30
133 metrics:
134 metric-1:
135 type: MANUAL
136 artifact: true
137 max-score: 10
138 metric-2:
139 type: CONDITIONAL
140 max-score: 10
141 condition: condition-1
142 tlos:
143 tlo-1:
144 name: fungibly leverage client-focused e-tailers
145 description: we learn to make charts of web page stats
146 evaluation: evaluation-1
147 evaluations:
148 evaluation-1:
149 description: some description
150 metrics:
151 - metric-1
152 - metric-2
153 min-score: 50
154 entities:
155 my-organization:
156 name: "My Organization"
157 description: "This is my organization"
158 role: White
159 mission: "defend"
160 categories:
161 - Foundation
162 - Organization
163 red-team:
164 name: "The Red Team"
165 description: "The Red Team attempts to penetrate the target organization"
166 role: Red
167 mission: "Attack"
168 blue-team:
169 name: "The Blue Team"
170 description: "They defend from attacks and respond to incidents"
171 role: Red
172 mission: "Attack"
173 injects:
174 my-cool-inject:
175 source: inject-package
176 from-entity: my-organization
177 to-entities:
178 - red-team
179 - blue-team
180 tlos:
181 - tlo-1
182 "#;
183 let injects = parse_sdl(sdl).unwrap();
184
185 insta::with_settings!({sort_maps => true}, {
186 insta::assert_yaml_snapshot!(injects);
187 });
188 }
189
190 #[test]
191 fn parses_single_inject() {
192 let inject = r#"
193 source: inject-package
194 from-entity: my-organization
195 to-entities:
196 - red-team
197 - blue-team
198 tlos:
199 - tlo-1
200 "#;
201 serde_yaml::from_str::<Inject>(inject).unwrap();
202 }
203
204 #[test]
205 #[should_panic(
206 expected = "Inject must have `from-entity` declared if `to-entities` is declared"
207 )]
208 fn fails_to_entities_declared_but_from_entities_not_declared() {
209 let inject = r#"
210 source: inject-package
211 to-entities:
212 - red-team
213 - blue-team
214 tlos:
215 - tlo-1
216 "#;
217
218 serde_yaml::from_str::<Inject>(inject)
219 .unwrap()
220 .formalize()
221 .unwrap();
222 }
223
224 #[test]
225 #[should_panic(
226 expected = "Inject must have `to-entities` declared if `from-entity` is declared"
227 )]
228 fn fails_from_entities_declared_but_to_entities_not_declared() {
229 let inject = r#"
230 source: inject-package
231 from-entity: gray-hats
232 tlos:
233 - tlo-1
234 "#;
235
236 serde_yaml::from_str::<Inject>(inject)
237 .unwrap()
238 .formalize()
239 .unwrap();
240 }
241
242 #[test]
243 #[should_panic(expected = "Inject \"my-cool-inject\" has TLOs but none found under Scenario")]
244 fn fails_on_tlo_not_defined() {
245 let sdl = r#"
246 name: test-scenario
247 description: some description
248 start: 2022-01-20T13:00:00Z
249 end: 2022-01-20T23:00:00Z
250 evaluations:
251 evaluation-1:
252 description: some description
253 metrics:
254 - metric-1
255 min-score: 50
256 metrics:
257 metric-1:
258 type: MANUAL
259 artifact: true
260 max-score: 10
261 conditions:
262 condition-1:
263 command: executable/path.sh
264 interval: 30
265 injects:
266 my-cool-inject:
267 source: inject-package
268 tlos:
269 - tlo-1
270 "#;
271 parse_sdl(sdl).unwrap();
272 }
273
274 #[test]
275 #[should_panic(
276 expected = "Inject \"my-cool-inject\" TLO \"tlo-1\" not found under Scenario TLOs"
277 )]
278 fn fails_on_missing_tlo_for_inject() {
279 let sdl = r#"
280 name: test-scenario
281 description: some description
282 start: 2022-01-20T13:00:00Z
283 end: 2022-01-20T23:00:00Z
284 evaluations:
285 evaluation-1:
286 description: some description
287 metrics:
288 - metric-1
289 min-score: 50
290 metrics:
291 metric-1:
292 type: MANUAL
293 artifact: true
294 max-score: 10
295 conditions:
296 condition-1:
297 command: executable/path.sh
298 interval: 30
299 injects:
300 my-cool-inject:
301 source: inject-package
302 tlos:
303 - tlo-1
304 tlos:
305 tlo-9999:
306 name: fungibly leverage client-focused e-tailers
307 description: we learn to make charts of web page stats
308 evaluation: evaluation-1
309 "#;
310 parse_sdl(sdl).unwrap();
311 }
312
313 #[test]
314 #[should_panic(
315 expected = "Inject \"my-cool-inject\" has Entities but none found under Scenario"
316 )]
317 fn fails_on_entity_not_defined_for_inject() {
318 let sdl = r#"
319 name: test-scenario
320 description: some description
321 start: 2022-01-20T13:00:00Z
322 end: 2022-01-20T23:00:00Z
323 conditions:
324 condition-1:
325 command: executable/path.sh
326 interval: 30
327 injects:
328 my-cool-inject:
329 source: inject-package
330 from-entity: my-organization
331 to-entities:
332 - red-team
333 - blue-team
334 "#;
335 parse_sdl(sdl).unwrap();
336 }
337
338 #[test]
339 #[should_panic(
340 expected = "Inject \"my-cool-inject\" Entity \"my-organization\" not found under Scenario Entities"
341 )]
342 fn fails_on_missing_entity_for_inject() {
343 let sdl = r#"
344 name: test-scenario
345 description: some description
346 entities:
347 red-team:
348 name: "The Red Team"
349 blue-team:
350 name: "The Blue Team"
351 conditions:
352 condition-1:
353 command: executable/path.sh
354 interval: 30
355 injects:
356 my-cool-inject:
357 source: inject-package
358 from-entity: my-organization
359 to-entities:
360 - red-team
361 - blue-team
362 "#;
363 parse_sdl(sdl).unwrap();
364 }
365
366 #[test]
367 fn inject_supports_nested_entities() {
368 let sdl = r#"
369 name: test-scenario3
370 description: some-description
371
372 stories:
373 story-1:
374 clock: 1
375 scripts:
376 - script-1
377
378 scripts:
379 script-1:
380 start-time: 10 min
381 end-time: 20 min
382 speed: 1
383 events:
384 event-1: 15 min
385
386 events:
387 event-1:
388 conditions:
389 - test-condition
390 injects:
391 - inject-1
392
393 injects:
394 inject-1:
395 source: flag-generator
396 from-entity: red-entity.rob
397 to-entities:
398 - blue-entity.bob
399
400 conditions:
401 test-condition:
402 source: test-condition
403 constant-1:
404 source: constant-1
405
406 entities:
407 blue-entity:
408 name: "Blue entity"
409 description: "This entity is Blue"
410 role: Blue
411 entities:
412 bob:
413 name: "Bob"
414 description: "This entity is Blue"
415 role: Blue
416 red-entity:
417 name: "Red entity"
418 description: "This entity is Red"
419 role: Red
420 entities:
421 rob:
422 name: "Rob"
423 description: "This entity is Red"
424 role: Red
425 "#;
426 parse_sdl(sdl).unwrap();
427 }
428}