fix: 修复load_exe函数中的问题,现在可以正常工作了

2024-12-30 00:02:57 +08:00 · 2024-12-30 00:02:57 +08:00 · 3d1ce2d58a
commit 3d1ce2d58a
parent 1761d87cad
13 changed files with 159 additions and 56 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@ -23,6 +23,12 @@ dependencies = [
 "cfg-if",
 ]

+[[package]]
+name = "equivalent"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5"
+
 [[package]]
 name = "flate2"
 version = "1.0.35"
@ -33,15 +39,50 @@ dependencies = [
 "miniz_oxide",
 ]

+[[package]]
+name = "foldhash"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a0d2fde1f7b3d48b8395d5f2de76c18a528bd6a9cdde438df747bfcba3e05d6f"
+
+[[package]]
+name = "hashbrown"
+version = "0.15.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289"
+dependencies = [
+ "foldhash",
+]
+
+[[package]]
+name = "indexmap"
+version = "2.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "62f822373a4fe84d4bb149bf54e584a7f4abec90e072ed49cda0edea5b95471f"
+dependencies = [
+ "equivalent",
+ "hashbrown",
+]
+
 [[package]]
 name = "loadpe-rs"
 version = "0.1.0"
 dependencies = [
+ "lz4_flex",
 "object",
 "thiserror",
 "windows",
 ]

+[[package]]
+name = "lz4_flex"
+version = "0.11.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "75761162ae2b0e580d7e7c390558127e5f01b4194debd6221fd8c207fc80e3f5"
+dependencies = [
+ "twox-hash",
+]
+
 [[package]]
 name = "memchr"
 version = "2.7.4"
@ -63,7 +104,10 @@ version = "0.36.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87"
 dependencies = [
+ "crc32fast",
 "flate2",
+ "hashbrown",
+ "indexmap",
 "memchr",
 "ruzstd",
 ]
--- a/Cargo.toml
+++ b/Cargo.toml
@ -4,7 +4,8 @@ version = "0.1.0"
 edition = "2021"

 [dependencies]
-object = "0.36.7"
+lz4_flex = "0.11.3"
+object = {version= "0.36.7", features=["default", "write"]}
 thiserror = "2.0.9"

 [dependencies.windows]
--- a/Readme.md
+++ b/Readme.md
@ -5,3 +5,9 @@
 ```bash
 git remote add origin ssh://git@81.70.52.148:17022/asahi/loadpe-rs.git
 ```
+
+如果使用rust进行壳的开发:
+- 压缩程序并生成PE
+- 编写shellcode,shellcode依赖:
+    - 解压缩API
+    - kernel32.dll中的API
--- a/src/bin/main.rs
+++ b/src/bin/main.rs
@ -4,7 +4,7 @@ use loadpe_rs::load_exe;
 const EXE_PATH: &str = r#"E:\work\rust\loadpe-rs\tests\Testx64.exe"#;

 #[cfg(target_arch = "x86")]
-const EXE_PATH: &str = r#"E:\work\rust\loadpe-rs\tests\Testx32.exe"#;
+const EXE_PATH: &str = r#"E:\work\rust\loadpe-rs\tests\Testx32R.exe"#;

 fn main() {
    match load_exe(EXE_PATH, None) {
--- a/src/errors.rs
+++ b/src/errors.rs
@ -2,6 +2,7 @@ use std::alloc::LayoutError;

 use thiserror::Error;
 use object::Error as ObjectError;
+use object::write::Error as ObjectWriteError;
 use windows::core::Error as WindowsError;
 #[derive(Error, Debug)]
 pub enum LoadPEError {
@ -28,4 +29,6 @@ pub enum LoadPEError {
    LayoutError(#[from] LayoutError),
    #[error("更改内存属性失败！")]
    ChangeMemoryProtectFailed,
+    #[error("写入PE文件失败！")]
+    WritePEError(#[from] ObjectWriteError),
 }
--- a/src/lib.rs
+++ b/src/lib.rs
@ -1,4 +1,6 @@
 use errors::LoadPEError;
+
+use lz4_flex::block::compress_prepend_size;
 use object::{
    endian,
    pe::{
@ -6,11 +8,14 @@ use object::{
        IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE,
    },
    read::pe::{self, ImageOptionalHeader, PeFile32},
-    LittleEndian, Object,
+    write::pe::Writer,
+    LittleEndian, Object, ObjectSection, SectionIndex, U16Bytes,
 };
 use std::{
-    alloc::{alloc, Layout}, arch::asm, fs
+    alloc::{alloc, Layout},
+    fs,
 };
+use types::CompressPeInfo;
 use windows::{
    core::PCSTR,
    Win32::System::{
@ -23,12 +28,9 @@ use windows::{
 };
 pub mod errors;
 mod help;
+mod types;

 /// 加载并启动一个exe文件，需要支持命令行参数
-///
-///
-///
-///
 pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEError> {
    // 1. 检查文件是否存在
    if !fs::metadata(exe_path).is_ok() {
@ -49,13 +51,13 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
        .get(endian::LittleEndian);
    // 3. 为镜像分配内存
    // 分配的内存，后续需要修改内存属性
-    //
-    let layout = Layout::from_size_align(image_size as usize, 1)?;
+
+    // 对齐值一定要写0x1000，否则会出现问题
+    let layout = Layout::from_size_align(image_size as usize, 0x1000)?;
    let buf = unsafe { alloc(layout) };
    if buf.is_null() {
        return Err(LoadPEError::MemoryAllocFailed);
    }
-
    dbg!("分配的内存地址: {:?}", buf);

    // 进行数据的拷贝。
@ -100,15 +102,8 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
            let section_data = &data
                [pointer_to_raw_data as usize..(pointer_to_raw_data + size_of_raw_data) as usize];

-            dbg!(
-               "拷贝节区数据: va: {:?}, size_of_raw_data: {:?}, pointer_to_raw_data: {:?}, section_data.len: {:?}",
-                va, size_of_raw_data, pointer_to_raw_data, section_data.len()
-            );
-
            unsafe {
                let dst = (buf as usize + va as usize) as *mut u8;
-                dbg!("dst: {:p}", dst);
-
                std::ptr::copy(section_data.as_ptr(), dst, size_of_raw_data as usize);
            };
        }
@ -140,7 +135,6 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
                    LoadLibraryA(PCSTR::from_raw(p_module_name))
                        .map_err(|err| LoadPEError::LoadLibraryError(err))?
                };
-                dbg!(h_module);
                if h_module.is_invalid() {
                    // 这个逻辑应该走不到
                    unimplemented!("LoadLibraryA failed");
@ -230,7 +224,6 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
                    // 修复IAT表
                    // 将函数地址写入IAT表 如果是x86_64，需要写入8字节
                    // 如果是x86，需要写入4字节
-                    // TODO: 验证一下是否修复了
                    unsafe {
                        #[cfg(target_arch = "x86")]
                        {
@ -251,7 +244,6 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
            // 1. 获取重定向表的数据目录
            let relocation_table_directory = file.data_directory(IMAGE_DIRECTORY_ENTRY_BASERELOC);
            if relocation_table_directory.is_none() {
-                // TODO: 如果没有重定位表,应该要把这个错误返回出去
                return Err(LoadPEError::NoRelocationTable);
            }
            let base_relocation_table_rva = relocation_table_directory
@ -260,6 +252,11 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
                .get(LittleEndian);

            let mut relocation_offset = 0;
+            let image_base = file
+                .nt_headers()
+                .optional_header
+                .image_base
+                .get(LittleEndian);

            loop {
                // 读取一个IMAGE_BASE_RELOCATION
@ -276,6 +273,7 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
                let size_of_block = relocation.size_of_block.get(LittleEndian);
                let base_block = base_relocation_table_rva as usize + relocation_offset + 8;
                let block_num = (size_of_block - 8) / 2;
+
                for i in 0..block_num {
                    // 读取一个TypeOffset
                    let type_offset = unsafe {
@ -291,21 +289,12 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
                    // 需要修复的地址偏移
                    let offset_va = virtual_address + r_offset as u32;
                    // 若要应用基址重定位，会计算首选基址与在其中实际加载映像的基址之间的差异。 如果在首选基址处加载映像，则差异为零，因此无需应用基址重定位。
-                    let reloc_offset = buf as usize
-                        - file
-                            .nt_headers()
-                            .optional_header
-                            .image_base
-                            .get(LittleEndian) as usize;
+                    // 修复后的值: 按类型
+                    let reloc_offset = buf as usize - image_base as usize;
+                    // 被修复的地址:
                    let p_address = (buf as usize + offset_va as usize) as *mut u32;
                    let address_value = unsafe { *p_address };
-                    dbg!(
-                        "r_type: {:?}, r_offset: {:?}, offset_va: {:?}, virtual_address: {:?}",
-                        r_type,
-                        r_offset,
-                        offset_va,
-                        virtual_address
-                    );
+
                    // TODO: 验证一下是否修复了
                    match r_type {
                        0 => {}
@ -317,7 +306,6 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
                            unsafe {
                                *p_address = new_address;
                            }
-                            dbg!("修复一个16位地址: {:?} -> {:?}", address_value, new_address);
                        }
                        2 => {
                            // IMAGE_REL_BASED_LOW 基准重定位将差值的低 16 位加到偏移的 16 位字段。 16 位字段代表 32 位字的低半部分。
@ -327,16 +315,14 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
                            unsafe {
                                *p_address = new_address;
                            }
-                            dbg!("修复一个16位地址: {:?} -> {:?}", address_value, new_address);
                        }
                        3 => {
                            // 如果r_type是3，表示这个TypeOffset是一个IMAGE_REL_BASED_HIGHLOW
                            // 这个时候需要修复一个32位的地址
-                            let new_address = address_value + buf as u32;
+                            let new_address = address_value + reloc_offset as u32;
                            unsafe {
                                *p_address = new_address;
                            }
-                            dbg!("修复一个32位地址: {:?} -> {:?}", address_value, new_address);
                        }
                        10 => {
                            // IMAGE_REL_BASED_DIR64 基址重定位将差值添加到偏移的 64 位字段。
@ -344,7 +330,6 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
                            unsafe {
                                *(p_address as *mut u64) = new_address as u64;
                            }
-                            dbg!("修复一个64位地址: {:?} -> {:?}", address_value, new_address);
                        }
                        _ => {
                            // 其他的类型，暂时不支持
@ -368,13 +353,13 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
            // 遍历节区，设置节区属性
            for section in file.sections() {
                let va = section.pe_section().virtual_address.get(LittleEndian);
-                let size_of_raw_data = section.pe_section().size_of_raw_data.get(LittleEndian);
+                // let size_of_raw_data = section.pe_section().size_of_raw_data.get(LittleEndian);
+                let section_name = section.name().unwrap();
+                let virtual_size = section.pe_section().virtual_size.get(LittleEndian);
                let characteristics = section.pe_section().characteristics.get(LittleEndian);
                let section_va = buf as usize + va as usize;
-                let section_size = align_to!(size_of_raw_data, section_alignment);
-                if section_size == 0 {
-                    continue;
-                }
+                let section_size = align_to!(virtual_size, section_alignment);
+
                // 设置节区属性
                let protect;
                if characteristics & IMAGE_SCN_MEM_EXECUTE != 0
@ -405,16 +390,9 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
                } else {
                    return Err(LoadPEError::ChangeMemoryProtectFailed);
                }
-
                // 设置属性
                let mut old_protect = PAGE_PROTECTION_FLAGS(0);
-
-                dbg!(
-                    "设置节区属性: va: {:?}, size: {:?}, protect: {:?}",
-                    section_va,
-                    section_size,
-                    protect
-                );
+                println!("{section_name} va: {:x}, section_size: {:x} address: {:x} characteristics: {:x} protect: {:x}", va, section_size, section_va, characteristics, protect);

                let result = unsafe {
                    VirtualProtect(
@ -425,7 +403,6 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
                    )
                };
                if result.is_err() {
-                    dbg!("设置节区属性失败: {:?}", result.unwrap_err());
                    return Err(LoadPEError::ChangeMemoryProtectFailed);
                }
            }
@ -439,9 +416,67 @@ pub fn load_exe(exe_path: &str, _args: Option<Vec<&str>>) -> Result<(), LoadPEEr
            let entry_point = buf as usize + entry_point as usize;
            let entry_point: extern "system" fn() = unsafe { std::mem::transmute(entry_point) };
            dbg!("entry_point: {:p}", entry_point);
-
            entry_point();
        }
    }
    Ok(())
 }
+
+/// 压缩PE并保存到指定路径
+/// 需要注入解压缩的位置无关代码
+pub fn compress_pe(pe_data: &[u8], output_path: &str) -> Result<(), LoadPEError> {
+    // 1. 解析PE文件
+    #[cfg(target_arch = "x86")]
+    let origin_pe: PeFile32 = pe::PeFile::parse(pe_data)?;
+    #[cfg(target_arch = "x86_64")]
+    let origin_pe: PeFile64 = pe::PeFile::parse(pe_data)?;
+    // 2. 获取头部数据
+    let origin_pe_header_data = {
+        let head_size = origin_pe
+            .nt_headers()
+            .optional_header
+            .size_of_headers
+            .get(endian::LittleEndian);
+        &pe_data[..head_size as usize]
+    };
+    // 压缩头部数据
+    let compressed_header_data = compress_prepend_size(origin_pe_header_data);
+
+    // 3. 获取节区数据
+    let section_data = {
+        let offset = origin_pe
+            .section_by_index(SectionIndex(0))?
+            .pe_section()
+            .pointer_to_raw_data
+            .get(LittleEndian);
+        &pe_data[offset as usize..]
+    };
+    // 压缩节区数据
+    let compressed_section_data = compress_prepend_size(section_data);
+
+    // 4. 构造压缩后的PE文件
+    {
+        let section_alignment = origin_pe.nt_headers().optional_header.section_alignment();
+        let file_alignment = origin_pe.nt_headers().optional_header.file_alignment();
+        let compress_info = CompressPeInfo {
+            header_origin_size: origin_pe_header_data.len() as u32,
+            header_compress_size: compressed_header_data.len() as u32,
+            code_origin_size: section_data.len() as u32,
+            code_compress_size: compressed_section_data.len() as u32,
+        };
+        let mut buf = Vec::with_capacity(1024 * 1024);
+        let is_64 = origin_pe.is_64();
+        let mut new_pe_writer = Writer::new(is_64, section_alignment, file_alignment, &mut buf);
+
+        // 写入DOS头和dos sub
+        new_pe_writer.write_dos_header_and_stub()?;
+        // 写入PE头 拷贝一份
+        let mut new_pe_header = origin_pe.nt_headers().clone();
+        // 重写节区
+        new_pe_header.file_header.number_of_sections = U16Bytes::new(LittleEndian, 2);
+
+        // TODO: 数据目录中,我还是想要kernel32.dll中的LoadLibraryA和GetProcAddress地址
+    }
+
+    todo!()
+}
--- a/src/types.rs
+++ b/src/types.rs
@ -0,0 +1,14 @@
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy)]
+/// 保存压缩前的原始PE文件信息
+pub struct CompressPeInfo {
+    // 原始PE文件头大小
+    pub header_origin_size: u32,
+    // 压缩后的PE文件头大小
+    pub header_compress_size: u32,
+    // 原始PE文件代码段大小
+    pub code_origin_size: u32,
+    // 压缩后的PE文件代码段大小
+    pub code_compress_size: u32,
+}
--- a/tests/Test.exe
+++ b/tests/Test.exe
--- a/tests/Testx32.exe
+++ b/tests/Testx32.exe
--- a/tests/Testx32D.exe
+++ b/tests/Testx32D.exe
--- a/tests/Testx32R.exe
+++ b/tests/Testx32R.exe
--- a/tests/Testx64.exe
+++ b/tests/Testx64.exe
--- a/tests/Testx64D.exe
+++ b/tests/Testx64D.exe