MFT Attributes structure restructure and serialization restructure

3 files changed, 92 insertions, 52 deletions

diff --git a/src/attribute.rs b/src/attribute.rs
index 8d662db..368bd0e 100644
--- a/src/attribute.rs
+++ b/src/attribute.rs
@@ -54,9 +54,12 @@ pub fn serialize_attr_data_flags<S>(&item: &AttributeDataFlags, serializer: S)
 #[derive(Serialize, Clone, Debug)]
 pub struct AttributeHeader {
     pub attribute_type: u32,
+    #[serde(skip_serializing)]
     pub attribute_size: u32,
     pub resident_flag: u8, // 0 -> resident; 1 -> non-resident
+    #[serde(skip_serializing)]
     pub name_size: u8,
+    #[serde(skip_serializing)]
     pub name_offset: u16,
     #[serde(serialize_with = "serialize_attr_data_flags")]
     pub data_flags: AttributeDataFlags,
@@ -142,9 +145,12 @@ pub enum ResidentialHeader{
 
 #[derive(Serialize, Clone, Debug)]
 pub struct ResidentHeader{
+    #[serde(skip_serializing)]
     pub data_size: u32,
+    #[serde(skip_serializing)]
     pub data_offset: u16,
     pub index_flag: u8,
+    #[serde(skip_serializing)]
     pub padding: u8,
 }
 impl ResidentHeader {
@@ -168,8 +174,10 @@ pub struct NonResidentHeader{
     pub vnc_first: u64,
     #[serde(serialize_with = "serialize_u64")]
     pub vnc_last: u64,
+    #[serde(skip_serializing)]
     pub datarun_offset: u16,
     pub unit_compression_size: u16,
+    #[serde(skip_serializing)]
     pub padding: u32,
     #[serde(serialize_with = "serialize_u64")]
     pub size_allocated: u64,
diff --git a/src/entry.rs b/src/entry.rs
index b2f3dac..2af49d5 100644
--- a/src/entry.rs
+++ b/src/entry.rs
@@ -1,3 +1,4 @@
+use std::collections::BTreeMap;
 use errors::{MftError};
 use enumerator::{PathMapping};
 use mft::{MftHandler};
@@ -34,27 +35,35 @@ pub fn serialize_entry_flags<S>(&item: &EntryFlags, serializer: S) -> Result<S::
 #[derive(Serialize, Debug)]
 pub struct EntryHeader{
     pub signature: u32,
+    #[serde(skip_serializing)]
     pub usa_offset: u16,
+    #[serde(skip_serializing)]
     pub usa_size: u16,
     #[serde(serialize_with = "serialize_u64")]
     pub logfile_sequence_number: u64,
+    #[serde(skip_serializing)]
     pub sequence: u16,
     pub hard_link_count: u16,
+    #[serde(skip_serializing)]
     pub fst_attr_offset: u16,
     #[serde(serialize_with = "serialize_entry_flags")]
     pub flags: EntryFlags,
+    #[serde(skip_serializing)]
     pub entry_size_real: u32,
+    #[serde(skip_serializing)]
     pub entry_size_allocated: u32,
     pub base_reference: MftReference,
+    #[serde(skip_serializing)]
     pub next_attribute_id: u16,
     #[serde(skip_serializing)]
     pub padding: Option<u16>,
-    pub record_number: Option<u32>,
+    #[serde(skip_serializing)]
+    pub record_number: u64,
     pub update_sequence_value: u32,
-    pub entry_reference: Option<MftReference>
+    pub entry_reference: MftReference
 }
 impl EntryHeader{
-    pub fn new<R: Read>(mut reader: R, entry: Option<u64>) -> Result<EntryHeader,MftError> {
+    pub fn new<R: Read>(mut reader: R, entry: u64) -> Result<EntryHeader,MftError> {
         let mut entry_header: EntryHeader = unsafe {
             mem::zeroed()
         };
@@ -84,22 +93,19 @@ impl EntryHeader{
 
         if entry_header.usa_offset == 48 {
             entry_header.padding = Some(reader.read_u16::<LittleEndian>()?);
-            entry_header.record_number = Some(reader.read_u32::<LittleEndian>()?);
+            entry_header.record_number = reader.read_u32::<LittleEndian>()? as u64;
         } else {
+            entry_header.record_number = entry as u64;
             panic!(
                 "Unhandled update sequence array offset: {}",
                 entry_header.usa_offset
             )
         }
 
-        if entry_header.record_number.is_some(){
-            entry_header.entry_reference = Some(
-                MftReference::get_from_entry_and_seq(
-                    entry_header.record_number.unwrap() as u64,
-                    entry_header.sequence
-                )
-            );
-        }
+        entry_header.entry_reference = MftReference::get_from_entry_and_seq(
+            entry_header.record_number as u64,
+            entry_header.sequence
+        );
 
         Ok(entry_header)
     }
@@ -108,20 +114,21 @@ impl EntryHeader{
 #[derive(Serialize, Debug)]
 pub struct MftEntry{
     pub header: EntryHeader,
-    pub attributes: Vec<attribute::MftAttribute>
+    pub attributes: BTreeMap<String,Vec<attribute::MftAttribute>>
 }
 impl MftEntry{
-    pub fn new(mut buffer: Vec<u8>, entry: Option<u64>) -> Result<MftEntry,MftError> {
-        let mut mft_entry: MftEntry = unsafe {
-            mem::zeroed()
-        };
-
+    pub fn new(mut buffer: Vec<u8>, entry: u64) -> Result<MftEntry,MftError> {
         // Get Header
-        mft_entry.header = EntryHeader::new(
+        let entry_header = EntryHeader::new(
             buffer.as_slice(),
             entry
         )?;
 
+        let mut mft_entry = MftEntry{
+            header: entry_header,
+            attributes: BTreeMap::new()
+        };
+
         // Fixup buffer
         mft_entry.buffer_fixup(
             &mut buffer
@@ -151,22 +158,25 @@ impl MftEntry{
     }
 
     pub fn get_pathmap(&self) -> Option<PathMapping> {
-        for attribute in self.attributes.iter() {
-            if attribute.header.attribute_type == 0x30 {
-                match attribute.content {
-                    attribute::AttributeContent::AttrX30(ref attrib) => {
-                        if attrib.namespace != 2 {
-                            return Some(
-                                PathMapping {
-                                    name: attrib.name.clone(),
-                                    parent: MftReference(attrib.parent.0)
-                                }
-                            );
+        match self.attributes.get("0x0030"){
+            Some(fn_attr_list) => {
+                for attribute in fn_attr_list {
+                    match attribute.content {
+                        attribute::AttributeContent::AttrX30(ref attrib) => {
+                            if attrib.namespace != 2 {
+                                return Some(
+                                    PathMapping {
+                                        name: attrib.name.clone(),
+                                        parent: MftReference(attrib.parent.0)
+                                    }
+                                );
+                            }
                         }
+                        _ => {}
                     }
-                    _ => {}
                 }
             }
+            None => {}
         }
 
         None
@@ -245,8 +255,7 @@ impl MftEntry{
                         }
                     }
 
-                    // push attribute into attributes
-                    self.attributes.push(
+                    self.set_attribute(
                         attribute::MftAttribute {
                             header: attribute_header.clone(),
                             content: attr_content
@@ -255,7 +264,7 @@ impl MftEntry{
                 },
                 attribute::ResidentialHeader::NonResident(_) => {
                     // No content, so push header into attributes
-                    self.attributes.push(
+                    self.set_attribute(
                         attribute::MftAttribute {
                             header: attribute_header.clone(),
                             content: attribute::AttributeContent::None
@@ -275,24 +284,47 @@ impl MftEntry{
         Ok(attr_count)
     }
 
+    pub fn set_attribute(&mut self, attribute: attribute::MftAttribute){
+        // This could maybe use some refactoring??
+        // Check if attribute type is already in mapping
+        let attr_type = format!("0x{:04X}",attribute.header.attribute_type);
+        if self.attributes.contains_key(&attr_type) {
+            // We already have a list for this attribute, get mut ref
+            if let Some(attr_list) = self.attributes.get_mut(&attr_type) {
+                // append the attribute to list
+                attr_list.push(attribute);
+            }
+        } else {
+            // Create new attribute list
+            let mut attr_list: Vec<attribute::MftAttribute> = Vec::new();
+            // Insert value into list
+            attr_list.push(attribute);
+            // Set attribute list to key of attrib type
+            self.attributes.insert(
+                attr_type,
+                attr_list
+            );
+        }
+    }
+
     pub fn set_fullnames(&mut self, mft_handler: &mut MftHandler){
-        // Iterate through each MFT attribute with mutable reference
-        for attribute in self.attributes.iter_mut() {
-            // If attribute is type 0x30
-            if attribute.header.attribute_type == 0x30 {
-                // Check if resident content
-                match attribute.content {
-                    attribute::AttributeContent::AttrX30(ref mut attrib) => {
-                        // Get fullpath
-                        let fullpath = mft_handler.get_fullpath(
-                            attrib.parent
-                        );
-                        // Set fullname
-                        let fullname = fullpath + "/" + attrib.name.as_str();
-                        // Set attribute to fullname
-                        attrib.fullname = Some(fullname);
+        if self.attributes.contains_key("0x0030") {
+            if let Some(attr_list) = self.attributes.get_mut("0x0030") {
+                for attribute in attr_list.iter_mut() {
+                    // Check if resident content
+                    match attribute.content {
+                        attribute::AttributeContent::AttrX30(ref mut attrib) => {
+                            // Get fullpath
+                            let fullpath = mft_handler.get_fullpath(
+                                attrib.parent
+                            );
+                            // Set fullname
+                            let fullname = fullpath + "/" + attrib.name.as_str();
+                            // Set attribute to fullname
+                            attrib.fullname = Some(fullname);
+                        }
+                        _ => {}
                     }
-                    _ => {}
                 }
             }
         }
diff --git a/src/mft.rs b/src/mft.rs
index 639d221..5a182bb 100644
--- a/src/mft.rs
+++ b/src/mft.rs
@@ -74,7 +74,7 @@ impl MftHandler{
             Some(mapping) => {
                 if mft_entry.is_dir() {
                     &self.path_enumerator.set_mapping(
-                        mft_entry.header.entry_reference.clone().unwrap(),
+                        mft_entry.header.entry_reference.clone(),
                         mapping.clone()
                     );
                 }
@@ -96,7 +96,7 @@ impl MftHandler{
     pub fn entry_from_buffer(&mut self, buffer: Vec<u8>, entry: u64) -> Result<MftEntry,MftError> {
         let mft_entry = MftEntry::new(
             buffer,
-            Some(entry)
+            entry
         )?;
 
         Ok(mft_entry)