use visitor to traverse resource tree

src/main.rs

@@ -5,6 +5,7 @@ mod pefile;
 #[allow(non_snake_case)]
 #[allow(non_camel_case_types)]
 mod winnt;
+mod utils;
 
 use argparse::{ArgumentParser, Store};
 use std::path::PathBuf;
@@ -24,5 +25,5 @@ fn main() {
         Ok(file)    =>  file,
         Err(why)    =>  {log::error!("{}", why); std::process::exit(1); }
     };
-    let _resources = pefile.get_resources();
+    pefile.print_resources();
 }

src/pefile.rs

@@ -1,16 +1,41 @@
-use std::path::PathBuf;
-use std::fs::File;
-use std::io::{Error, ErrorKind};
-use memmap::MmapOptions;
 use byteorder::{ByteOrder, LittleEndian};
-use num_traits::ToPrimitive;
+use memmap::MmapOptions;
 use num_traits::FromPrimitive;
+use num_traits::ToPrimitive;
+use std::fs::File;
+use std::io::{Error, ErrorKind};
+use std::path::PathBuf;
 use std::str;
 
+use crate::winnt::IMAGE_OPTIONAL_HEADER::*;
+use crate::winnt::*;
 use from_bytes::StructFromBytes;
 use packed_size::*;
-use crate::winnt::*;
-use crate::winnt::IMAGE_OPTIONAL_HEADER::*;
+
+pub trait ResourceDirectoryVisitor {
+    fn init(&mut self, _pefile: &PEFile) {}
+    fn finalize(&mut self, _pefile: &PEFile) {}
+
+    fn enter_resource_directory(
+        &mut self,
+        pefile: &PEFile,
+        dir: &IMAGE_RESOURCE_DIRECTORY,
+        identifier: &EntryIdentifier,
+    );
+    fn leave_resource_directory(
+        &mut self,
+        pefile: &PEFile,
+        dir: &IMAGE_RESOURCE_DIRECTORY,
+        identifier: &EntryIdentifier,
+    );
+
+    fn visit_resource_data_entry(
+        &mut self,
+        pefile: &PEFile,
+        entry: &IMAGE_RESOURCE_DATA_ENTRY,
+        identifier: &EntryIdentifier,
+    );
+}
 
 #[allow(dead_code)]
 pub struct PEFile {
@@ -20,34 +45,43 @@ pub struct PEFile {
     image_file_header: IMAGE_FILE_HEADER,
     image_optional_header: Option<IMAGE_OPTIONAL_HEADER>,
     directories: Vec<Option<IMAGE_DATA_DIRECTORY>>,
-    sections: Vec<IMAGE_SECTION_HEADER>
+    sections: Vec<IMAGE_SECTION_HEADER>,
 }
 
 impl PEFile {
     pub fn new(filename: PathBuf) -> std::io::Result<PEFile> {
         let mut offset = 0;
-        let file=File::open(&filename)?;
+        let file = File::open(&filename)?;
         let mmap = unsafe { MmapOptions::new().map(&file)? };
 
         let image_dos_header = IMAGE_DOS_HEADER::from_bytes(&mmap, offset)?;
 
         if image_dos_header.e_magic != LittleEndian::read_u16(b"MZ") {
-            return Err(Error::new(ErrorKind::InvalidData, format!("illegal DOS magic: {:?}", &mmap[0..2])));
+            return Err(Error::new(
+                ErrorKind::InvalidData,
+                format!("illegal DOS magic: {:?}", &mmap[0..2]),
+            ));
         } else {
             log::debug!("DOS magic is ok");
         }
 
         let nt_magic_offset = image_dos_header.e_lfanew as usize;
-        let nt_magic = &mmap[nt_magic_offset .. nt_magic_offset+4];
+        let nt_magic = &mmap[nt_magic_offset..nt_magic_offset + 4];
         if nt_magic != [b'P', b'E', 0, 0] {
-            return Err(Error::new(ErrorKind::InvalidData, format!("illegal NT magic: {:?}", nt_magic)));
+            return Err(Error::new(
+                ErrorKind::InvalidData,
+                format!("illegal NT magic: {:?}", nt_magic),
+            ));
         } else {
             log::debug!("NT magic is ok");
         }
-
         offset = (image_dos_header.e_lfanew + 4) as usize;
         let nt_header_size = IMAGE_FILE_HEADER::packed_size();
-        log::debug!("searching extended header at 0x{:08x}, size = {}", offset, nt_header_size);
+        log::debug!(
+            "searching extended header at 0x{:08x}, size = {}",
+            offset,
+            nt_header_size
+        );
         let image_file_header = IMAGE_FILE_HEADER::from_bytes(&mmap, offset)?;
         offset += nt_header_size;
 
@@ -57,7 +91,7 @@ impl PEFile {
         let image_optional_header = if optional_header_size == 0 {
             None
         } else {
-            let header_magic = &mmap[offset..offset+2];
+            let header_magic = &mmap[offset..offset + 2];
             match FromPrimitive::from_u16(LittleEndian::read_u16(header_magic)) {
                 Some(IMAGE_NT_OPTIONAL_HEADER::IMAGE_NT_OPTIONAL_HDR32_MAGIC) => {
                     let header = IMAGE_OPTIONAL_HEADER32::from_bytes(&mmap, offset)?;
@@ -69,8 +103,11 @@ impl PEFile {
                     offset += IMAGE_OPTIONAL_HEADER64::packed_size();
                     Some(AMD64(*header))
                 }
-                _  => {
-                    return Err(Error::new(ErrorKind::InvalidData, format!("illegal optional header magic: {:?}", header_magic)));
+                _ => {
+                    return Err(Error::new(
+                        ErrorKind::InvalidData,
+                        format!("illegal optional header magic: {:?}", header_magic),
+                    ));
                 }
             }
         };
@@ -86,7 +123,12 @@ impl PEFile {
                 let entry = IMAGE_DATA_DIRECTORY::from_bytes(&mmap, offset + (entry_size * idx))?;
 
                 if entry.VirtualAddress != 0 {
-                    log::debug!("DATA DIRECTORY {:02}: address = 0x{:08x}, size = {}", idx, entry.VirtualAddress, entry.Size);
+                    log::debug!(
+                        "DATA DIRECTORY {:02}: address = 0x{:08x}, size = {}",
+                        idx,
+                        entry.VirtualAddress,
+                        entry.Size
+                    );
                     directories.push(Some(*entry));
                 } else {
                     log::debug!("DATA DIRECTORY {:02}: <EMPTY>", idx);
@@ -100,17 +142,28 @@ impl PEFile {
         // load section headers
         let mut sections = Vec::new();
         let entry_size = IMAGE_SECTION_HEADER::packed_size();
-        for idx in 0 .. image_file_header.NumberOfSections {
-            let entry = IMAGE_SECTION_HEADER::from_bytes(&mmap, offset + (entry_size * idx as usize))?;
+        for idx in 0..image_file_header.NumberOfSections {
+            let entry =
+                IMAGE_SECTION_HEADER::from_bytes(&mmap, offset + (entry_size * idx as usize))?;
 
             let section_name = str::from_utf8(&entry.Name[..]).unwrap();
-            let virt_size  = entry.Misc;
-            let virt_addr  = entry.VirtualAddress;
+            let virt_size = entry.Misc;
+            let virt_addr = entry.VirtualAddress;
             let raw_offset = entry.PointerToRawData;
-            let raw_size   = entry.SizeOfRawData;
+            let raw_size = entry.SizeOfRawData;
 
-            log::debug!("{:02} {}  VirtAddr: {:08x}      VirtSize: {:08x}", idx, section_name, virt_addr, virt_size);
-            log::debug!("  raw data offs: {:08x} raw data size: {:08x} ",raw_offset, raw_size);
+            log::debug!(
+                "{:02} {}  VirtAddr: {:08x}      VirtSize: {:08x}",
+                idx,
+                section_name,
+                virt_addr,
+                virt_size
+            );
+            log::debug!(
+                "  raw data offs: {:08x} raw data size: {:08x} ",
+                raw_offset,
+                raw_size
+            );
 
             sections.push(*entry);
         }
@@ -126,7 +179,6 @@ impl PEFile {
         };
         return Ok(me);
     }
-
     #[allow(dead_code)]
     pub fn info(&self) -> String {
         let mut lines = Vec::new();
@@ -138,33 +190,158 @@ impl PEFile {
     }
 
     pub fn get_raw_address(&self, rva: usize) -> Option<usize> {
-        match self.sections.iter()
-                .find(|&x| (x.VirtualAddress as usize .. (x.VirtualAddress + x.Misc) as usize)
-                .contains(&rva)) {
-            None        => None,
-            Some(sect)  => {
-                log::debug!("found rva {:08x} in section {}", rva, str::from_utf8(&sect.Name[..]).unwrap());
-                let raw_address = rva - sect.VirtualAddress as usize + sect.PointerToRawData as usize;
-                log::debug!("raw address is {:08x} ({:x} - {:x} + {:x})", raw_address, rva, sect.VirtualAddress, sect.PointerToRawData);
+        match self.sections.iter().find(|&x| {
+            (x.VirtualAddress as usize..(x.VirtualAddress + x.Misc) as usize).contains(&rva)
+        }) {
+            None => None,
+            Some(sect) => {
+                log::debug!(
+                    "found rva {:08x} in section {}",
+                    rva,
+                    str::from_utf8(&sect.Name[..]).unwrap()
+                );
+                let raw_address =
+                    rva - sect.VirtualAddress as usize + sect.PointerToRawData as usize;
+                log::debug!(
+                    "raw address is {:08x} ({:x} - {:x} + {:x})",
+                    raw_address,
+                    rva,
+                    sect.VirtualAddress,
+                    sect.PointerToRawData
+                );
                 Some(raw_address)
             }
         }
     }
-
-    pub fn get_resources(&self)  -> std::io::Result<Option<ImageResourceDirectory>> {
-
-        let idx_resources = ToPrimitive::to_usize(&IMAGE_DIRECTORY_ENTRY::IMAGE_DIRECTORY_ENTRY_RESOURCE).unwrap();
+    
+    pub fn get_resources_section(&self) -> Option<&[u8]> {
+        let idx_resources =
+            ToPrimitive::to_usize(&IMAGE_DIRECTORY_ENTRY::IMAGE_DIRECTORY_ENTRY_RESOURCE).unwrap();
         if let Some(entry) = &self.directories[idx_resources] {
             if let Some(offset) = self.get_raw_address(entry.VirtualAddress as usize) {
-                log::debug!("loading resources of size {} at 0x{:08x}", entry.Size, offset);
+                log::debug!(
+                    "loading resources of size {} at 0x{:08x}",
+                    entry.Size,
+                    offset
+                );
 
                 // create slice to enforce bounds checking
-                let resources = &self.mmap[offset .. offset+entry.Size as usize];
-
-                let root = ImageResourceDirectory::from_bytes(resources, 0)?;
-                return Ok(Some(root));
+                return Some(&self.mmap[offset..offset + entry.Size as usize]);
             }
         }
-        Ok(None)
+        None
+    }
+
+    pub fn print_resources(&self) {
+        let mut visitor = ConsoleVisitor::new();
+        self.visit_resource_tree(&mut visitor).unwrap();
     }
-}
+
+    pub fn visit_resource_tree<V: ResourceDirectoryVisitor>(
+        &self,
+        visitor: &mut V,
+    ) -> std::io::Result<()> {
+        visitor.init(self);
+
+        if let Some(resources) = self.get_resources_section() {
+            self.visit_directory(resources, visitor, 0, EntryIdentifier::NoIdentifier)?;
+        }
+
+        visitor.finalize(self);
+        Ok(())
+    }
+
+    pub fn visit_directory<V: ResourceDirectoryVisitor>(
+        &self,
+        resources: &[u8],
+        visitor: &mut V,
+        offset: usize,
+        identifier: EntryIdentifier,
+    ) -> std::io::Result<()> {
+        let dir = IMAGE_RESOURCE_DIRECTORY::from_bytes(resources, offset)?;
+        visitor.enter_resource_directory(self, &dir, &identifier);
+
+        let offset = offset + IMAGE_RESOURCE_DIRECTORY::packed_size();
+        let entry_size = IMAGE_RESOURCE_DIRECTORY_ENTRY::packed_size();
+        let count = (dir.NumberOfNamedEntries + dir.NumberOfIdEntries) as usize;
+        for idx in 0..count {
+            let entry_offset = offset + idx * entry_size;
+            self.visit_directory_entry(resources, visitor, entry_offset)?;
+        }
+        visitor.leave_resource_directory(self, &dir, &identifier);
+        Ok(())
+    }
+
+    pub fn visit_directory_entry<V: ResourceDirectoryVisitor>(
+        &self,
+        resources: &[u8],
+        visitor: &mut V,
+        offset: usize,
+    ) -> std::io::Result<()> {
+        let raw_entry = IMAGE_RESOURCE_DIRECTORY_ENTRY::from_bytes(resources, offset)?;
+        let identifier = raw_entry.parse_identifier(resources);
+        if (raw_entry.OffsetToData & 0x80000000) == 0x80000000 {
+            let entry_offset = raw_entry.OffsetToData & 0x7fffffff;
+            self.visit_directory(resources, visitor, entry_offset as usize, identifier)?;
+        } else {
+            let entry_offset = raw_entry.OffsetToData as usize;
+            let raw_entry = IMAGE_RESOURCE_DATA_ENTRY::from_bytes(resources, entry_offset)?;
+            visitor.visit_resource_data_entry(self, &raw_entry, &identifier);
+        }
+
+        Ok(())
+    }
+}
+
+pub struct ConsoleVisitor {
+    level: u32,
+}
+impl ConsoleVisitor {
+    pub fn new() -> Self {
+        Self { level: 0 }
+    }
+    fn indent(&self) -> String {
+        let mut res = String::with_capacity(self.level as usize * 2);
+        for _ in 0..self.level {
+            res.push_str("  ");
+        }
+        res
+    }
+    fn enter(&mut self) {
+        self.level += 1;
+    }
+    fn leave(&mut self) {
+        if self.level == 0 {
+            panic!("stack underflow");
+        }
+        self.level -= 1;
+    }
+}
+impl ResourceDirectoryVisitor for ConsoleVisitor {
+    fn enter_resource_directory(
+        &mut self,
+        _pefile: &PEFile,
+        _dir: &IMAGE_RESOURCE_DIRECTORY,
+        identifier: &EntryIdentifier,
+    ) {
+        self.enter();
+        println!("{}{:?}", self.indent(), identifier);
+    }
+    fn leave_resource_directory(
+        &mut self,
+        _pefile: &PEFile,
+        _dir: &IMAGE_RESOURCE_DIRECTORY,
+        _identifier: &EntryIdentifier,
+    ) {
+        self.leave();
+    }
+
+    fn visit_resource_data_entry(
+        &mut self,
+        _pefile: &PEFile,
+        _entry: &IMAGE_RESOURCE_DATA_ENTRY,
+        identifier: &EntryIdentifier,
+    ) {
+        println!("{} -> {:?}", self.indent(), identifier);
+    }
+}

src/utils.rs

@@ -0,0 +1,8 @@
+pub fn utf16_from_slice(slice: &[u8], mut offset: usize, length: usize) -> String {
+    let mut name_chars = Vec::new();
+    for _ in 0..length {
+        name_chars.push(slice[offset] as u16 | ((slice[offset+1] as u16)<<8 as u8));
+        offset += 2;
+    }
+    String::from_utf16_lossy(&name_chars[..])
+}