Write support for PE binaries

Cargo.toml

@@ -54,3 +54,6 @@ archive = ["alloc"]
 [badges.travis-ci]
 branch = "master"
 repository = "m4b/goblin"
+
+[dev-dependencies]
+stderrlog = "0.5.4"

examples/dotnet_pe_analysis.rs

@@ -67,7 +67,8 @@ fn main() {
         .expect("No CLI header");
     let sections = &pe.sections;
 
-    let cli_header_value: CliHeader = get_data(file, sections, cli_header, file_alignment).unwrap();
+    let cli_header_value: CliHeader =
+        get_data(file, sections, *cli_header, file_alignment).unwrap();
     println!("{:#?}", cli_header_value);
     let metadata_root: MetadataRoot =
         get_data(file, sections, cli_header_value.metadata, file_alignment).unwrap();

examples/rewrite_pe.rs

@@ -0,0 +1,42 @@
+use goblin::pe::PE;
+use scroll::Pwrite;
+
+fn main() {
+    stderrlog::new().verbosity(1).init().unwrap();
+    let args: Vec<String> = std::env::args().collect();
+
+    let file = std::fs::read(&args[1]).unwrap();
+    let file = &file[..];
+    let pe = PE::parse(file).unwrap();
+    println!("read {}", &args[1]);
+
+    println!(
+        "file alignment: {:?}",
+        pe.header
+            .optional_header
+            .unwrap()
+            .windows_fields
+            .file_alignment
+    );
+
+    let mut new_pe = vec![0u8; file.len() + 8192];
+    let new_len = new_pe.pwrite(pe, 0).unwrap();
+    let pe = PE::parse(file).unwrap();
+
+    let out = &new_pe[..new_len];
+    std::fs::write(&args[2], &out).unwrap();
+    println!("written as {}", &args[2]);
+    println!(
+        "original PE size: {} bytes, new PE size: {} bytes, delta (new - original): {} bytes",
+        file.len(),
+        out.len(),
+        out.len() as isize - file.len() as isize
+    );
+
+    let new_pe = PE::parse(&new_pe).unwrap();
+    println!(
+        "original signatures: {}, new signatures: {}",
+        pe.certificates.len(),
+        new_pe.certificates.len()
+    );
+}

src/error.rs

@@ -3,6 +3,7 @@
 
 use alloc::string::String;
 use core::fmt;
+use core::num::TryFromIntError;
 use core::result;
 #[cfg(feature = "std")]
 use std::{error, io};
@@ -42,6 +43,12 @@ impl From<io::Error> for Error {
     }
 }
 
+impl From<TryFromIntError> for Error {
+    fn from(err: TryFromIntError) -> Error {
+        Error::Malformed(format!("Integer do not fit: {err}"))
+    }
+}
+
 impl From<scroll::Error> for Error {
     fn from(err: scroll::Error) -> Error {
         Error::Scroll(err)

src/pe/certificate_table.rs

@@ -3,11 +3,13 @@
 /// https://learn.microsoft.com/en-us/windows/win32/debug/pe-format#the-attribute-certificate-table-image-only
 /// https://learn.microsoft.com/en-us/windows/win32/api/wintrust/ns-wintrust-win_certificate
 use crate::error;
-use scroll::Pread;
+use scroll::{ctx, Pread, Pwrite};
 
 use alloc::string::ToString;
 use alloc::vec::Vec;
 
+use super::utils::pad;
+
 #[repr(u16)]
 #[non_exhaustive]
 #[derive(Debug, PartialEq, Copy, Clone)]
@@ -39,7 +41,7 @@ impl TryFrom<u16> for AttributeCertificateRevision {
 }
 
 #[repr(u16)]
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Copy, Clone)]
 pub enum AttributeCertificateType {
     /// WIN_CERT_TYPE_X509
     X509 = 0x0001,
@@ -128,7 +130,28 @@ impl<'a> AttributeCertificate<'a> {
     }
 }
 
+impl<'a> ctx::TryIntoCtx<scroll::Endian> for &AttributeCertificate<'a> {
+    type Error = error::Error;
+
+    /// Writes an aligned attribute certificate in the buffer.
+    fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result<usize, Self::Error> {
+        let offset = &mut 0;
+        bytes.gwrite_with(self.length, offset, ctx)?;
+        bytes.gwrite_with(self.revision as u16, offset, ctx)?;
+        bytes.gwrite_with(self.certificate_type as u16, offset, ctx)?;
+        // Extend by zero the buffer until it is aligned on a quadword (16 bytes).
+        let maybe_certificate_padding = pad(self.certificate.len(), Some(16usize));
+        bytes.gwrite(self.certificate, offset)?;
+        if let Some(cert_padding) = maybe_certificate_padding {
+            bytes.gwrite(&cert_padding[..], offset)?;
+        }
+
+        Ok(*offset)
+    }
+}
+
 pub type CertificateDirectoryTable<'a> = Vec<AttributeCertificate<'a>>;
+
 pub(crate) fn enumerate_certificates(
     bytes: &[u8],
     table_virtual_address: u32,

src/pe/data_directories.rs

@@ -1,5 +1,8 @@
 use crate::error;
-use scroll::{Pread, Pwrite, SizeWith};
+use scroll::{
+    ctx::{self},
+    Pread, Pwrite, SizeWith,
+};
 
 #[repr(C)]
 #[derive(Debug, PartialEq, Copy, Clone, Default, Pread, Pwrite, SizeWith)]
@@ -13,14 +16,86 @@ const NUM_DATA_DIRECTORIES: usize = 16;
 
 impl DataDirectory {
     pub fn parse(bytes: &[u8], offset: &mut usize) -> error::Result<Self> {
-        let dd = bytes.gread_with(offset, scroll::LE)?;
-        Ok(dd)
+        Ok(bytes.gread_with(offset, scroll::LE)?)
+    }
+}
+
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum DataDirectoryType {
+    ExportTable,
+    ImportTable,
+    ResourceTable,
+    ExceptionTable,
+    CertificateTable,
+    BaseRelocationTable,
+    DebugTable,
+    Architecture,
+    GlobalPtr,
+    TlsTable,
+    LoadConfigTable,
+    BoundImportTable,
+    ImportAddressTable,
+    DelayImportDescriptor,
+    ClrRuntimeHeader,
+}
+
+impl TryFrom<usize> for DataDirectoryType {
+    type Error = error::Error;
+    fn try_from(value: usize) -> Result<Self, Self::Error> {
+        Ok(match value {
+            0 => Self::ExportTable,
+            1 => Self::ImportTable,
+            2 => Self::ResourceTable,
+            3 => Self::ExceptionTable,
+            4 => Self::CertificateTable,
+            5 => Self::BaseRelocationTable,
+            6 => Self::DebugTable,
+            7 => Self::Architecture,
+            8 => Self::GlobalPtr,
+            9 => Self::TlsTable,
+            10 => Self::LoadConfigTable,
+            11 => Self::BoundImportTable,
+            12 => Self::ImportAddressTable,
+            13 => Self::DelayImportDescriptor,
+            14 => Self::ClrRuntimeHeader,
+            _ => {
+                return Err(error::Error::Malformed(
+                    "Wrong data directory index number".into(),
+                ))
+            }
+        })
     }
 }
 
 #[derive(Debug, PartialEq, Copy, Clone, Default)]
 pub struct DataDirectories {
-    pub data_directories: [Option<DataDirectory>; NUM_DATA_DIRECTORIES],
+    pub data_directories: [Option<(usize, DataDirectory)>; NUM_DATA_DIRECTORIES],
+}
+
+impl ctx::TryIntoCtx<scroll::Endian> for DataDirectories {
+    type Error = error::Error;
+
+    fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result<usize, Self::Error> {
+        let offset = &mut 0;
+        for opt_dd in self.data_directories {
+            if let Some((dd_offset, dd)) = opt_dd {
+                bytes.pwrite_with(dd, dd_offset, ctx)?;
+                *offset += dd_offset;
+            } else {
+                bytes.gwrite(&[0; SIZEOF_DATA_DIRECTORY][..], offset)?;
+            }
+        }
+        Ok(NUM_DATA_DIRECTORIES * SIZEOF_DATA_DIRECTORY)
+    }
+}
+
+macro_rules! build_dd_getter {
+    ($dd_name:tt, $index:tt) => {
+        pub fn $dd_name(&self) -> Option<&DataDirectory> {
+            let idx = $index;
+            self.data_directories[idx].as_ref().map(|(_, dd)| dd)
+        }
+    };
 }
 
 impl DataDirectories {
@@ -37,70 +112,42 @@ impl DataDirectories {
             let dd = if dd.virtual_address == 0 && dd.size == 0 {
                 None
             } else {
-                Some(dd)
+                Some((*offset, dd))
             };
             *dir = dd;
         }
         Ok(DataDirectories { data_directories })
     }
-    pub fn get_export_table(&self) -> &Option<DataDirectory> {
-        let idx = 0;
-        &self.data_directories[idx]
-    }
-    pub fn get_import_table(&self) -> &Option<DataDirectory> {
-        let idx = 1;
-        &self.data_directories[idx]
-    }
-    pub fn get_resource_table(&self) -> &Option<DataDirectory> {
-        let idx = 2;
-        &self.data_directories[idx]
-    }
-    pub fn get_exception_table(&self) -> &Option<DataDirectory> {
-        let idx = 3;
-        &self.data_directories[idx]
-    }
-    pub fn get_certificate_table(&self) -> &Option<DataDirectory> {
-        let idx = 4;
-        &self.data_directories[idx]
-    }
-    pub fn get_base_relocation_table(&self) -> &Option<DataDirectory> {
-        let idx = 5;
-        &self.data_directories[idx]
-    }
-    pub fn get_debug_table(&self) -> &Option<DataDirectory> {
-        let idx = 6;
-        &self.data_directories[idx]
-    }
-    pub fn get_architecture(&self) -> &Option<DataDirectory> {
-        let idx = 7;
-        &self.data_directories[idx]
-    }
-    pub fn get_global_ptr(&self) -> &Option<DataDirectory> {
-        let idx = 8;
-        &self.data_directories[idx]
-    }
-    pub fn get_tls_table(&self) -> &Option<DataDirectory> {
-        let idx = 9;
-        &self.data_directories[idx]
-    }
-    pub fn get_load_config_table(&self) -> &Option<DataDirectory> {
-        let idx = 10;
-        &self.data_directories[idx]
-    }
-    pub fn get_bound_import_table(&self) -> &Option<DataDirectory> {
-        let idx = 11;
-        &self.data_directories[idx]
-    }
-    pub fn get_import_address_table(&self) -> &Option<DataDirectory> {
-        let idx = 12;
-        &self.data_directories[idx]
-    }
-    pub fn get_delay_import_descriptor(&self) -> &Option<DataDirectory> {
-        let idx = 13;
-        &self.data_directories[idx]
-    }
-    pub fn get_clr_runtime_header(&self) -> &Option<DataDirectory> {
-        let idx = 14;
-        &self.data_directories[idx]
+
+    build_dd_getter!(get_export_table, 0);
+    build_dd_getter!(get_import_table, 1);
+    build_dd_getter!(get_resource_table, 2);
+    build_dd_getter!(get_exception_table, 3);
+    build_dd_getter!(get_certificate_table, 4);
+    build_dd_getter!(get_base_relocation_table, 5);
+    build_dd_getter!(get_debug_table, 6);
+    build_dd_getter!(get_architecture, 7);
+    build_dd_getter!(get_global_ptr, 8);
+    build_dd_getter!(get_tls_table, 9);
+    build_dd_getter!(get_load_config_table, 10);
+    build_dd_getter!(get_bound_import_table, 11);
+    build_dd_getter!(get_import_address_table, 12);
+    build_dd_getter!(get_delay_import_descriptor, 13);
+    build_dd_getter!(get_clr_runtime_header, 14);
+
+    pub fn dirs(&self) -> impl Iterator<Item = (DataDirectoryType, DataDirectory)> {
+        self.data_directories
+            .into_iter()
+            .enumerate()
+            // (Index, Option<DD>) -> Option<(Index, DD)> -> (DDT, DD)
+            .filter_map(|(i, o)|
+                // We should not have invalid indexes.
+                // Indeed: `data_directories: &[_; N]` where N is the number
+                // of data directories.
+                // The `TryFrom` trait for integers to DataDirectoryType
+                // takes into account the N possible data directories.
+                // Therefore, the unwrap can never fail as long as Rust guarantees
+                // on types are honored.
+                o.map(|(_, v)| (i.try_into().unwrap(), v)))
     }
 }