Update binary patching implementation

src/main.rs

@@ -133,6 +133,7 @@ fn main() {
 
                 let mut program_table_offset = file_header.phoff;
                 let mut program_table_count: i32 = 0;
+                let mut pt_note_index: usize = 0;
                 let mut pt_note_offset: usize = 0;
 
                 // Iterate through number of Program Headers
@@ -148,8 +149,9 @@ fn main() {
                     // Parse the program name using the program type
                     let program_name: String = util::parse_program_segment_type(program_header.program_type);
 
-                    if (program_name == "PT_NOTE") && (pt_note_offset == 0) {
+                    if (program_name == "PT_NOTE") && (pt_note_index == 0) {
                         pt_note_offset = program_table_offset as usize;
+                        pt_note_index = i as usize;
                     }
 
                     util::pp_program_header(&program_header, program_table_count, &program_name);
@@ -162,7 +164,7 @@ fn main() {
                 let note_segment: elf::ProgramHeader = util::build_program_header(
                     bytes, 
                     pt_note_offset,
-                    0,
+                    pt_note_index as u16,
                     file_header.is_x86_64
                 );
 

src/patcher.rs

@@ -29,7 +29,7 @@ pub fn patch_binary(
 
 
     // Apply patch to end of binary
-    println!("Patch data read successfully, injecting at end of binary...");
+    print!("Patch data read successfully, injecting at end of binary...");
 
     let injection_offset: usize = program_data.len();
     let injection_size: usize = patch_data.len();
@@ -37,22 +37,22 @@ pub fn patch_binary(
 
     program_data.extend_from_slice(patch_data);
 
-    print!("Done!");
+    println!("Done!");
 
 
     // Locate a note segment
-    println!("Pulling .note.ABI-tag segment data...");
+    print!("Pulling .note.ABI-tag segment data...");
     let note_section: &elf::SectionHeader = section_header_map.get(".note.ABI-tag")
         .expect("[Error] Failed to pull ABI-tag section from binary!");
-    print!("Done!\n");
+    println!("Done!\n");
     
     println!("Note section address: {:#04x}", note_section.addr);
     println!("Note section offset: {:#04x}", note_section.offset);
-    println!("Note section size: {}", note_section.size);
+    println!("Note section size: {}\n", note_section.size);
-    println!("");
+    
     println!("Injected address: {:#04x}", injection_addr);
     println!("Injected section offset: {:#04x}", injection_offset);
-    println!("Injected section size: {}", injection_size);
+    println!("Injected section size: {}\n", injection_size);
 
 
     // Rewrite the section header
@@ -62,6 +62,8 @@ pub fn patch_binary(
     injected_section.addr = injection_addr as u64;
     injected_section.offset = injection_offset as u64;
     injected_section.size = injection_size as u64;
+    injected_section.addralign = 16;
+    injected_section.flags = 6;
     
     util::overwrite_section_header(
         &mut program_data,
@@ -73,8 +75,7 @@ pub fn patch_binary(
     );
 
 
-    // 
+    // Rewrite the program segment
-    // Rewrite the section header
     let mut injected_segment: elf::ProgramHeader = elf::ProgramHeader::from(note_segment.clone());
 
     injected_segment.program_type = 1;
@@ -88,16 +89,15 @@ pub fn patch_binary(
 
     util::overwrite_segment_header(
         &mut program_data,
-        file_header.shoff as usize,
+        file_header.phoff as usize,
-        file_header.shentsize as usize,
+        file_header.phentsize as usize,
-        injected_section.id as usize,
+        injected_segment.id as usize,
-        &injected_section,
+        &injected_segment,
         file_header.is_x86_64
     );
 
 
-
+    // Rewrite the program entrypoint
-
     util::overwrite_entrypoint(&mut program_data, injection_offset);
 
 

src/util.rs

@@ -66,11 +66,12 @@ pub fn build_program_header(data: &Vec<u8>, phoffset: usize, id: u16, is_x86_64:
 }
 
 
-pub fn overwrite_segment_header(program_data: &mut Vec<u8>,
+pub fn overwrite_segment_header(
-    stoffset: usize,
+    program_data: &mut Vec<u8>,
-    shentsize: usize,
+    phoffset: usize,
-    shentidx: usize,
+    phentsize: usize,
-    new_section: &elf::SectionHeader,
+    phentidx: usize,
+    new_segment: &elf::ProgramHeader,
     is_x86_64: bool) {
 
     // Cast the supplied is_x86_64 bool to an array offset
@@ -78,32 +79,51 @@ pub fn overwrite_segment_header(program_data: &mut Vec<u8>,
     // 1 : x64
     let arch: usize = is_x86_64.into();
 
-println!("stoffset {}", stoffset);
+    println!("phoffset {}", phoffset);
-println!("shentsize {}", shentsize);
+    println!("phentsize {}", phentsize);
-println!("shentidx {}", shentidx);
+    println!("phentidx {}", phentidx);
 
-let section_addr_offset: usize = stoffset + (shentsize * shentidx) + elf::SH_ADDR_OFFSET[arch] as usize;
+    let segment_type_offset: usize = phoffset + (phentsize * phentidx) + elf::PH_TYPE_OFFSET as usize;
-let section_offset_offset: usize = stoffset + (shentsize * shentidx) + elf::SH_OFFSET_OFFSET[arch] as usize;
+    let segment_offset_offset: usize = phoffset + (phentsize * phentidx) + elf::PH_OFFSET_OFFSET[arch] as usize;
-let section_size_offset: usize = stoffset + (shentsize * shentidx) + elf::SH_SIZE_OFFSET[arch] as usize;
+    let segment_vaddr_offset: usize = phoffset + (phentsize * phentidx) + elf::PH_VADDR_OFFSET[arch] as usize;
-let section_type_offset: usize = stoffset + (shentsize * shentidx) + elf::SH_TYPE_OFFSET as usize;
+    let segment_paddr_offset: usize = phoffset + (phentsize * phentidx) + elf::PH_PADDR_OFFSET[arch] as usize;
+    let segment_filesz_offset: usize = phoffset + (phentsize * phentidx) + elf::PH_FILESZ_OFFSET[arch] as usize;
+    let segment_memsz_offset: usize = phoffset + (phentsize * phentidx) + elf::PH_MEMSZ_OFFSET[arch] as usize;
+    let segment_flags_offset: usize = phoffset + (phentsize * phentidx) + elf::PH_FLAGS_OFFSET[arch] as usize;
+    let segment_align_offset: usize = phoffset + (phentsize * phentidx) + elf::PH_ALIGN_OFFSET[arch] as usize;
 
-program_data[section_addr_offset..section_addr_offset+8].copy_from_slice(
+    program_data[segment_type_offset..segment_type_offset+4].copy_from_slice(
-&new_section.addr.to_ne_bytes().to_vec());
+            &new_segment.program_type.to_ne_bytes().to_vec());
-println!("Overwriting section addr with {:#04x}", new_section.addr);
+    println!("Overwriting segment type with {:#04x}", new_segment.program_type);
 
-program_data[section_offset_offset..section_offset_offset+8].copy_from_slice(
+    program_data[segment_offset_offset..segment_offset_offset+8].copy_from_slice(
-&new_section.offset.to_ne_bytes().to_vec());
+        &new_segment.offset.to_ne_bytes().to_vec());
-println!("Overwriting section offset with {:#04x}", new_section.offset as usize);
+    println!("Overwriting segment offset with {:#04x}", new_segment.offset);
 
-program_data[section_size_offset..section_size_offset+8].copy_from_slice(
+    program_data[segment_vaddr_offset..segment_vaddr_offset+8].copy_from_slice(
-&new_section.size.to_ne_bytes().to_vec());
+        &new_segment.vaddr.to_ne_bytes().to_vec());
-println!("Overwriting section size with {:#04x}", new_section.size as usize);
+    println!("Overwriting segment vaddr with {:#04x}", new_segment.vaddr);
 
-program_data[section_type_offset..section_type_offset+4].copy_from_slice(
+    program_data[segment_paddr_offset..segment_paddr_offset+8].copy_from_slice(
-&new_section.section_type.to_ne_bytes().to_vec());
+        &new_segment.paddr.to_ne_bytes().to_vec());
-println!("Overwriting section type with {:#04x}", new_section.section_type as usize);
+    println!("Overwriting segment paddr with {:#04x}", new_segment.paddr);
+
+    program_data[segment_filesz_offset..segment_filesz_offset+8].copy_from_slice(
+        &new_segment.filesz.to_ne_bytes().to_vec());
+    println!("Overwriting segment filesz with {:#04x}", new_segment.filesz);
+
+    program_data[segment_memsz_offset..segment_memsz_offset+8].copy_from_slice(
+        &new_segment.memsz.to_ne_bytes().to_vec());
+    println!("Overwriting segment memsz with {:#04x}", new_segment.memsz);
+
+    program_data[segment_flags_offset..segment_flags_offset+4].copy_from_slice(
+        &new_segment.flags.to_ne_bytes().to_vec());
+    println!("Overwriting segment flag with {:#04x}", new_segment.flags);
+
+    program_data[segment_align_offset..segment_align_offset+8].copy_from_slice(
+        &new_segment.align.to_ne_bytes().to_vec());
+    println!("Overwriting segment alignment with {:#04x}\n", new_segment.align);
 
-// return section_header;
 }
 
 
@@ -152,6 +172,8 @@ pub fn overwrite_section_header(program_data: &mut Vec<u8>,
     let section_offset_offset: usize = stoffset + (shentsize * shentidx) + elf::SH_OFFSET_OFFSET[arch] as usize;
     let section_size_offset: usize = stoffset + (shentsize * shentidx) + elf::SH_SIZE_OFFSET[arch] as usize;
     let section_type_offset: usize = stoffset + (shentsize * shentidx) + elf::SH_TYPE_OFFSET as usize;
+    let section_flag_offset: usize = stoffset + (shentsize * shentidx) + elf::SH_FLAGS_OFFSET as usize;
+    let section_align_offset: usize = stoffset + (shentsize * shentidx) + elf::SH_ADDRALIGN_OFFSET[arch] as usize;
 
     program_data[section_addr_offset..section_addr_offset+8].copy_from_slice(
                 &new_section.addr.to_ne_bytes().to_vec());
@@ -169,6 +191,14 @@ pub fn overwrite_section_header(program_data: &mut Vec<u8>,
         &new_section.section_type.to_ne_bytes().to_vec());
     println!("Overwriting section type with {:#04x}", new_section.section_type as usize);
 
+    program_data[section_flag_offset..section_flag_offset+8].copy_from_slice(
+        &new_section.flags.to_ne_bytes().to_vec());
+    println!("Overwriting section flags with {:#04x}", new_section.flags as usize);
+
+    program_data[section_align_offset..section_align_offset+8].copy_from_slice(
+        &new_section.addralign.to_ne_bytes().to_vec());
+    println!("Overwriting section address alignment with {:#04x}\n", new_section.addralign as usize);
+
     // return section_header;
 }
 
@@ -180,6 +210,7 @@ pub fn overwrite_entrypoint(program_data: &mut Vec<u8>,
     program_data[offset..offset+8].copy_from_slice(
         &new_entry_point.to_ne_bytes().to_vec()
     );
+    println!("Overwriting program entrypoint with {:#04x}\n", new_entry_point as usize);
 }
 
 
@@ -403,6 +434,7 @@ pub fn pp_section_header(header: &elf::SectionHeader, number: i32, name: &String
     println!("[{}] {}", number, name);
     println!("\t- Type: {}", parse_section_type(header.section_type));
     println!("\t- Flags: {}", parse_section_flags(header.flags));
+    println!("\t- Flags (Value): {}", header.flags);
     println!("\t- Address: {:#04x}", header.section_type);
     println!("\t- Offset: {:#04x}", header.section_type);
     println!("\t- Link Index: {}", header.link);