forensic_rs/
data.rs

1use std::{borrow::Cow, collections::BTreeMap};
2
3#[cfg(feature="serde")]
4use serde::{Deserialize, Serialize, de::Visitor, Deserializer, ser::SerializeMap};
5
6use crate::{prelude::{Artifact, *}, field::{internal::{InternalField, PreStoredField}, Text, Field, Ip}, context::context};
7
8/// Basic container for all Forensic Data inside an artifact
9#[derive(Debug, Clone)]
10pub struct ForensicData {
11    artifact : Artifact,
12    pub(crate) fields: BTreeMap<Text, InternalField>,
13}
14
15impl Default for ForensicData {
16    fn default() -> Self {
17        let context = context();
18        let mut fields = BTreeMap::new();
19        fields.insert(Text::Borrowed(ARTIFACT_HOST), Field::Text(Text::Owned(context.host)).into());
20        fields.insert(Text::Borrowed(ARTIFACT_TENANT), Field::Text(Text::Owned(context.tenant)).into());
21        fields.insert(Text::Borrowed(ARTIFACT_NAME), Field::Text(Text::Owned(context.artifact.to_string())).into());
22        Self {
23            fields,
24            artifact : context.artifact
25        }
26    }
27}
28
29
30impl<'a> ForensicData {
31    pub fn new(host : &str, artifact : Artifact) -> Self {
32        let mut fields = BTreeMap::new();
33        fields.insert(Text::Borrowed(ARTIFACT_HOST), Field::Text(Text::Owned(host.to_string())).into());
34        fields.insert(Text::Borrowed(ARTIFACT_NAME), Field::Text(Text::Owned(artifact.to_string())).into());
35        Self {
36            fields,
37            artifact
38        }
39    }
40
41    pub fn artifact(&self) -> &Artifact {
42        &self.artifact
43    }
44
45    pub fn host(&'a self) -> &'a str {
46        match self.field(ARTIFACT_HOST) {
47            Some(v) => {
48                match v {
49                    Field::Text(v) => v,
50                    _ => ""
51                }
52            },
53            None => ""
54        }
55    }
56
57    pub fn field(&self, field_name : &str) -> Option<&Field> {
58        Some(&self.fields.get(field_name)?.original)
59    }
60
61    pub fn has_field(&self, field_name : &str) -> bool {
62        self.fields.contains_key(field_name)
63    }
64
65    pub fn field_mut(&'a mut self, field_name: &str) -> Option<&mut Field> {
66        Some(&mut self.fields.get_mut(field_name)?.original)
67    }
68    pub fn add_field(&mut self, field_name: &'static str, field_value: Field) {
69        self.insert(Text::Borrowed(field_name), field_value);
70    }
71    pub fn insert(&mut self, field_name: Text, field_value: Field) {
72        self.fields.insert(field_name, field_value.into());
73    }
74    /// Obtains the casted value of the field into i64 and caches it
75    pub fn i64_field(&'a mut self, field_name: &str) -> Option<i64> {
76        let field = self.fields.get_mut(field_name)?;
77        match field.ni64.as_ref() {
78            PreStoredField::Invalid => return None,
79            PreStoredField::None => {},
80            PreStoredField::Some(v) => return Some(*v)
81        };
82        let i64field : Option<i64> = (&field.original).try_into().ok();
83        let pfield = match i64field {
84            Some(v) => PreStoredField::Some(v),
85            None => PreStoredField::Invalid
86        };
87        field.ni64 = Box::new(pfield);
88        match field.ni64.as_ref() {
89            PreStoredField::Some(v) => Some(*v),
90            _ => None
91        }
92    }
93    /// Obtains the casted value of the field into f64 and caches it
94    pub fn f64_field(&'a mut self, field_name: &str) -> Option<f64> {
95        let field = self.fields.get_mut(field_name)?;
96        match field.nf64.as_ref() {
97            PreStoredField::Invalid => return None,
98            PreStoredField::None => {},
99            PreStoredField::Some(v) => return Some(*v)
100        };
101        let i64field : Option<f64> = (&field.original).try_into().ok();
102        let pfield = match i64field {
103            Some(v) => PreStoredField::Some(v),
104            None => PreStoredField::Invalid
105        };
106        field.nf64 = Box::new(pfield);
107        match field.nf64.as_ref() {
108            PreStoredField::Some(v) => Some(*v),
109            _ => None
110        }
111    }
112    /// Obtains the casted value of the field into u64 and caches it
113    pub fn u64_field(&'a mut self, field_name: &str) -> Option<u64> {
114        let field = self.fields.get_mut(field_name)?;
115        match field.nu64.as_ref() {
116            PreStoredField::Invalid => return None,
117            PreStoredField::None => {},
118            PreStoredField::Some(v) => return Some(*v)
119        };
120        let i64field : Option<u64> = (&field.original).try_into().ok();
121        let pfield = match i64field {
122            Some(v) => PreStoredField::Some(v),
123            None => PreStoredField::Invalid
124        };
125        field.nu64 = Box::new(pfield);
126        match field.nu64.as_ref() {
127            PreStoredField::Some(v) => Some(*v),
128            _ => None
129        }
130    }
131    /// Obtains the casted value of the field into IP and caches it
132    pub fn ip_field(&'a mut self, field_name: &str) -> Option<Ip> {
133        let field = self.fields.get_mut(field_name)?;
134        match field.ip.as_ref() {
135            PreStoredField::Invalid => return None,
136            PreStoredField::None => {},
137            PreStoredField::Some(v) => return Some(*v)
138        };
139        let i64field : Option<Ip> = (&field.original).try_into().ok();
140        let pfield = match i64field {
141            Some(v) => PreStoredField::Some(v),
142            None => PreStoredField::Invalid
143        };
144        field.ip = Box::new(pfield);
145        match field.ip.as_ref() {
146            PreStoredField::Some(v) => Some(*v),
147            _ => None
148        }
149    }
150    /// Obtains the casted value of the field into Text and caches it
151    pub fn txt_field(&'a mut self, field_name: &str) -> Option<&Text> {
152
153        let mut has_value = false;
154
155        let field = self.fields.get_mut(field_name)?;
156        match field.text.as_ref() {
157            PreStoredField::Invalid => return None,
158            PreStoredField::None => {},
159            PreStoredField::Some(_) => {
160                has_value = true;
161            }
162        };
163        if has_value {
164            match field.text.as_ref() {
165                PreStoredField::Some(v) => return Some(v),
166                _ => return None
167            }
168        }
169        let txtfield : Option<Text> = (&field.original).try_into().ok();
170        let pfield = match txtfield {
171            Some(v) => PreStoredField::Some(v),
172            None => PreStoredField::Invalid
173        };
174        field.text = Box::new(pfield);
175        match field.text.as_ref() {
176            PreStoredField::Some(v) => Some(v),
177            _ => None
178        }
179    }
180    /// Obtains the casted value of the field into Vec<Text> and caches it
181    pub fn array_field(&'a mut self, field_name: &str) -> Option<&Vec<Text>> {
182
183        let mut has_value = false;
184
185        let field = self.fields.get_mut(field_name)?;
186        match field.array.as_ref() {
187            PreStoredField::Invalid => return None,
188            PreStoredField::None => {},
189            PreStoredField::Some(_) => {
190                has_value = true;
191            }
192        };
193        if has_value {
194            match field.array.as_ref() {
195                PreStoredField::Some(v) => return Some(v),
196                _ => return None
197            }
198        }
199        let txtfield : Option<Vec<Text>> = (&field.original).try_into().ok();
200        let pfield = match txtfield {
201            Some(v) => PreStoredField::Some(v),
202            None => PreStoredField::Invalid
203        };
204        field.array = Box::new(pfield);
205        match field.array.as_ref() {
206            PreStoredField::Some(v) => Some(v),
207            _ => None
208        }
209    }
210
211    pub fn fields(&self) -> EventIter<'_> {
212        EventIter {
213            children: self.fields.iter(),
214        }
215    }
216    pub fn iter(&self) -> EventIter<'_> {
217        EventIter {
218            children: self.fields.iter(),
219        }
220    }
221    pub fn iter_mut(&mut self) -> EventIterMut<'_> {
222        EventIterMut {
223            children: self.fields.iter_mut(),
224        }
225    }
226
227}
228
229
230pub struct ForensicDataInspector<'a> {
231    iter : std::collections::btree_map::Iter<'a, Cow<'static, str>, String>
232}
233pub struct ForensicDataInspectorMut<'a> {
234    iter : std::collections::btree_map::IterMut<'a, Cow<'static, str>, String>
235}
236
237impl<'a> Iterator for ForensicDataInspector<'a> {
238    type Item = (&'a Cow<'a,str>,&'a String);
239
240    fn next(&mut self) -> Option<Self::Item> {
241        self.iter.next().map(|wrapper| (wrapper.0, wrapper.1))
242    }
243}
244impl<'a> Iterator for ForensicDataInspectorMut<'a> {
245    type Item = (&'a Cow<'a,str>,&'a mut String);
246
247    fn next(&mut self) -> Option<Self::Item> {
248        self.iter.next().map(|wrapper| (wrapper.0, wrapper.1))
249    }
250}
251
252impl std::fmt::Display for ForensicData {
253    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
254        write!(f, "{{artifact:{:?}, fields:{:?}}}", self.artifact, self.fields)
255    }
256}
257
258pub struct EventIter<'a> {
259    children: std::collections::btree_map::Iter<'a, Text, InternalField>,
260}
261pub struct EventFieldIter<'a> {
262    names: std::collections::btree_set::Iter<'a, Text>,
263    fields: &'a BTreeMap<Text, InternalField>,
264}
265
266pub struct EventIterMut<'a> {
267    children: std::collections::btree_map::IterMut<'a, Text, InternalField>,
268}
269
270impl<'a> Iterator for EventIter<'a> {
271    type Item = (&'a Text, &'a Field);
272
273    fn next(&mut self) -> Option<Self::Item> {
274        let evt = self.children.next()?;
275        Some((evt.0, &evt.1.original))
276    }
277}
278impl<'a> Iterator for EventIterMut<'a> {
279    type Item = (&'a Text, &'a mut Field);
280
281    fn next(&mut self) -> Option<Self::Item> {
282        let evt = self.children.next()?;
283        Some((evt.0, &mut evt.1.original))
284    }
285}
286impl<'a> Iterator for EventFieldIter<'a> {
287    type Item = (&'a Text, &'a Field);
288
289    fn next(&mut self) -> Option<Self::Item> {
290        let field = self.names.next()?;
291        let value = self.fields.get(field)?;
292        Some((field, &value.original))
293    }
294}
295#[cfg(feature = "serde")]
296impl<'de> Deserialize<'de> for ForensicData {
297    fn deserialize<D>(deserializer: D) -> Result<ForensicData, D::Error>
298    where
299        D: Deserializer<'de>,
300    {
301        deserializer.deserialize_any(DataVisitor)
302    }
303}
304#[cfg(feature = "serde")]
305impl Serialize for ForensicData {
306    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
307    where
308        S: serde::Serializer {
309        let mut map = serializer.serialize_map(Some(self.fields.len()))?;
310        for (k,v) in &self.fields {
311            map.serialize_entry(k, &v.original)?;
312        }
313        map.end()
314    }
315}
316#[cfg(feature = "serde")]
317struct DataVisitor;
318#[cfg(feature = "serde")]
319impl<'de> Visitor<'de> for DataVisitor {
320    type Value = ForensicData;
321
322    fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
323        formatter.write_str("a valid forensic data")
324    }
325
326    fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
327        where
328            A: serde::de::MapAccess<'de>, {
329        let mut artifact = Artifact::default();
330        let mut fields = BTreeMap::new();
331        while let Some((key, value)) = map.next_entry()? {
332            fields.insert(Cow::Owned(key), InternalField::new(value));
333        }
334        if let Some(artf) = fields.get(ARTIFACT_NAME) {
335            if let Field::Text(artf) = &artf.original {
336                artifact = (&artf[..]).into();
337            }
338        }
339        Ok(ForensicData { artifact, fields })
340    }
341}
342
343#[cfg(test)]
344mod data_tests {
345    use crate::{prelude::RegistryArtifacts, artifact::{Artifact, WindowsArtifacts}};
346
347    use super::ForensicData;
348
349    #[test]
350    fn iterate_fields_test() {
351        let mut data = ForensicData::new("host007", RegistryArtifacts::ShellBags.into());
352        data.insert("field001".into(), "value001".into());
353        data.insert("field002".into(), "value002".into());
354        data.insert("field003".into(), "value003".into());
355
356        let mut count = 0;
357        for (_name, _value) in data.fields() {
358            count += 1;
359        }
360        assert_eq!(5, count);// 3 + 2
361    }
362
363    #[test]
364    fn should_serialize_data() {
365        let mut data = ForensicData::new("host007", RegistryArtifacts::ShellBags.into());
366        data.insert("field001".into(), "value001".into());
367        data.insert("field002".into(), "value002".into());
368        data.insert("field003".into(), "value003".into());
369        data.insert("field004".into(), crate::field::Field::Array(vec!["aaa".into(), "bbb".into()]));
370        let deserialized = serde_json::to_string(&data).unwrap();
371        assert_eq!(r#"{"artifact.host":"host007","artifact.name":"Windows::Registry::ShellBags","field001":"value001","field002":"value002","field003":"value003","field004":["aaa","bbb"]}"#, deserialized);
372        let serialized : ForensicData = serde_json::from_str(&deserialized).unwrap();
373        assert_eq!(Artifact::Windows(WindowsArtifacts::Registry(RegistryArtifacts::ShellBags)), serialized.artifact);
374        let deserialized2 = serde_json::to_string(&serialized).unwrap();
375        assert_eq!(deserialized, deserialized2);
376    }
377}
forensic_rs/data.rs

forensic_rs/
data.rs
