pe_assembler/helpers/builder/
mod.rs

1use crate::{
2    formats::exe::writer::ExeWriter,
3    helpers::pe_writer::PeWriter,
4    types::{
5        tables::{ImportEntry, ImportTable},
6        CoffHeader, DataDirectory, DosHeader, NtHeader, OptionalHeader, PeHeader, PeProgram, PeSection, SubsystemType,
7    },
8};
9use gaia_types::{helpers::Architecture, GaiaError};
10use std::io::Cursor;
11/// PE 汇编器构建器
12#[derive(Debug)]
13pub struct PeBuilder {
14    architecture: Option<Architecture>,
15    subsystem: Option<SubsystemType>,
16    entry_point: Option<u32>,
17    image_base: Option<u64>,
18    imports: Vec<(String, Vec<String>)>, // (dll_name, functions)
19    code: Option<Vec<u8>>,
20    data: Option<Vec<u8>>,
21    sections: Vec<PeSection>, // Add this field
22}
23
24impl PeBuilder {
25    /// 创建新的 PE 汇编器构建器
26    pub fn new() -> Self {
27        Self {
28            architecture: None,
29            subsystem: None,
30            entry_point: None,
31            image_base: None,
32            imports: Vec::new(),
33            code: None,
34            data: None,
35            sections: Vec::new(), // Initialize the new field
36        }
37    }
38
39    /// 设置目标架构
40    pub fn architecture(mut self, arch: Architecture) -> Self {
41        self.architecture = Some(arch);
42        self
43    }
44
45    /// 设置子系统类型
46    pub fn subsystem(mut self, subsystem: SubsystemType) -> Self {
47        self.subsystem = Some(subsystem);
48        self
49    }
50
51    /// 设置入口点地址
52    pub fn entry_point(mut self, entry_point: u32) -> Self {
53        self.entry_point = Some(entry_point);
54        self
55    }
56
57    /// 设置映像基地址
58    pub fn image_base(mut self, image_base: u64) -> Self {
59        self.image_base = Some(image_base);
60        self
61    }
62
63    /// 导入单个函数
64    pub fn import_function(mut self, dll_name: &str, function_name: &str) -> Self {
65        // 查找是否已存在该 DLL
66        if let Some(entry) = self.imports.iter_mut().find(|(name, _)| name == dll_name) {
67            entry.1.push(function_name.to_string());
68        }
69        else {
70            self.imports.push((dll_name.to_string(), vec![function_name.to_string()]));
71        }
72        self
73    }
74
75    /// 导入多个函数
76    pub fn import_functions(mut self, dll_name: &str, function_names: &[&str]) -> Self {
77        let functions: Vec<String> = function_names.iter().map(|&s| s.to_string()).collect();
78
79        // 查找是否已存在该 DLL
80        if let Some(entry) = self.imports.iter_mut().find(|(name, _)| name == dll_name) {
81            entry.1.extend(functions);
82        }
83        else {
84            self.imports.push((dll_name.to_string(), functions));
85        }
86        self
87    }
88
89    /// 设置代码数据
90    pub fn code(mut self, code: Vec<u8>) -> Self {
91        self.code = Some(code);
92        self
93    }
94
95    /// 设置数据
96    pub fn data(mut self, data: Vec<u8>) -> Self {
97        self.data = Some(data);
98        self
99    }
100
101    /// 获取导入信息（用于创建导入表）
102    pub fn get_imports(&self) -> &Vec<(String, Vec<String>)> {
103        &self.imports
104    }
105
106    /// 生成 PE 头部信息
107    pub fn build_header(&self) -> Result<PeHeader, GaiaError> {
108        let architecture = self
109            .architecture
110            .as_ref()
111            .ok_or_else(|| GaiaError::syntax_error("Architecture is required", gaia_types::SourceLocation::default()))?;
112        let pointer_size: u32 = if *architecture == Architecture::X86_64 { 8 } else { 4 };
113        let subsystem = self
114            .subsystem
115            .ok_or_else(|| GaiaError::syntax_error("Subsystem is required", gaia_types::SourceLocation::default()))?;
116        let entry_point = self.entry_point.unwrap_or(0x1000);
117        let image_base = self.image_base.unwrap_or(match architecture {
118            Architecture::X86 => 0x400000,
119            Architecture::X86_64 => 0x140000000,
120            _ => 0x400000,
121        });
122
123        // 创建 DOS 头
124        let dos_header = DosHeader::new(0x80);
125
126        // 创建 NT 头
127        let nt_header = NtHeader {
128            signature: 0x00004550, // "PE\0\0"
129        };
130
131        // 创建 COFF 头
132        let machine = match architecture {
133            Architecture::X86 => 0x014C,
134            Architecture::X86_64 => 0x8664,
135            _ => 0x014C,
136        };
137
138        let mut section_count = 0;
139        if self.code.is_some() {
140            section_count += 1;
141        }
142        if self.data.is_some() {
143            section_count += 1;
144        }
145        if !self.imports.is_empty() {
146            section_count += 1;
147        }
148
149        let optional_header_size = match architecture {
150            Architecture::X86_64 => 240,
151            _ => 224,
152        };
153
154        // 根据架构设置 COFF 特征位：
155        // - x86: 可执行映像 | 32 位机器
156        // - x64: 可执行映像 | 大地址感知（不设置 32 位机器位）
157        let characteristics = match architecture {
158            Architecture::X86 => 0x0102,    // IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_32BIT_MACHINE
159            Architecture::X86_64 => 0x0022, // IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_LARGE_ADDRESS_AWARE
160            _ => 0x0102,
161        };
162
163        let coff_header = CoffHeader::new(machine, section_count)
164            .with_timestamp(0)
165            .with_symbol_table(0, 0)
166            .with_optional_header_size(optional_header_size)
167            .with_characteristics(characteristics);
168
169        // 规范化 size_of_code 和 size_of_initialized_data，使用节的对齐后大小
170        let size_of_code = if let Some(code) = &self.code { ((code.len() + 0x1FF) / 0x200 * 0x200) as u32 } else { 0 };
171        let mut size_of_initialized_data = 0;
172        if let Some(data) = &self.data {
173            size_of_initialized_data += ((data.len() + 0x1FF) / 0x200 * 0x200) as u32;
174            // .data
175        }
176        if !self.imports.is_empty() {
177            // 这里我们无法准确预测 .idata 大小，但 0x200 应该足够大多数情况
178            // 或者我们可以先构建节来获取准确大小
179            size_of_initialized_data += 0x200; // .idata
180        }
181
182        // 计算 size_of_image：从 0x1000 开始，每个实际存在的节增加 0x1000
183        let mut size_of_image = 0x1000; // DOS/Headers 占一个对齐页
184        if self.code.is_some() {
185            size_of_image += 0x1000;
186        }
187        if self.data.is_some() {
188            size_of_image += 0x1000;
189        }
190        if !self.imports.is_empty() {
191            size_of_image += 0x1000;
192        }
193
194        let mut optional_header = OptionalHeader::new_for_architecture(
195            architecture,
196            entry_point,
197            image_base,
198            size_of_code,
199            0x200,         // size_of_headers
200            size_of_image, // 动态计算的镜像大小
201            subsystem,
202        );
203        optional_header.size_of_initialized_data = size_of_initialized_data;
204
205        // 关闭 ASLR（DYNAMIC_BASE），否则我们修补的绝对地址会因随机基址而失效
206        // DYNAMIC_BASE 位值为 0x0040
207        optional_header.dll_characteristics &= !0x0040;
208
209        // 设置导入表数据目录（动态计算 RVA）
210        // 兼容模式（x64）：IAT 与 INT 初始都指向 IMAGE_IMPORT_BY_NAME（Hint+Name）的 RVA。
211        // - x64：INT=名称指针数组，IAT=名称RVA（加载器解析后覆盖为真实地址）
212        // - x86：OFT/INT=0，IAT=名称RVA（常见兼容布局）
213        if !self.imports.is_empty() {
214            // 查找 .idata 节的实际 RVA
215            let idata_section = self
216                .sections
217                .iter()
218                .find(|s| s.name == ".idata")
219                .ok_or_else(|| GaiaError::syntax_error("Missing .idata section", gaia_types::SourceLocation::default()))?;
220            let import_rva_base = idata_section.virtual_address;
221
222            // 后续计算保持不变，但使用 import_rva_base 代替硬编码值
223            let mut current_rva = import_rva_base + ((self.imports.len() + 1) as u32) * 20;
224            for (dll_name, _) in &self.imports {
225                current_rva += (dll_name.len() as u32) + 1;
226            }
227            if current_rva % 2 != 0 {
228                current_rva += 1;
229            }
230            // 函数 Hint/Name
231            for (_, functions) in &self.imports {
232                for func in functions {
233                    // 名称按 2 字节对齐
234                    if current_rva % 2 != 0 {
235                        current_rva += 1;
236                    }
237                    current_rva += 2 + (func.len() as u32) + 1;
238                }
239            }
240            // 移除这里多余的对齐，改为与 write_import_table 保持一致
241            // INT
242            if current_rva % pointer_size != 0 {
243                current_rva = (current_rva + pointer_size - 1) & !(pointer_size - 1);
244            }
245            for (_, functions) in &self.imports {
246                current_rva += ((functions.len() as u32) + 1) * pointer_size;
247            }
248            // IAT（到此处即为 IAT 起点）
249            if current_rva % pointer_size != 0 {
250                current_rva = (current_rva + pointer_size - 1) & !(pointer_size - 1);
251            }
252            let iat_rva_start = current_rva; // 记录 IAT 起点
253                                             // IAT 长度
254            let mut end_rva = current_rva;
255            for (_, functions) in &self.imports {
256                end_rva += ((functions.len() as u32) + 1) * pointer_size;
257            }
258            optional_header.data_directories[1] =
259                DataDirectory { virtual_address: import_rva_base, size: end_rva - import_rva_base };
260            // 同时填写 IAT Directory（索引 12），让加载器知道 IAT 范围
261            let mut iat_rva_end = iat_rva_start;
262            for (_, functions) in &self.imports {
263                iat_rva_end += ((functions.len() as u32) + 1) * pointer_size;
264            }
265            optional_header.data_directories[12] =
266                DataDirectory { virtual_address: iat_rva_start, size: iat_rva_end - iat_rva_start };
267        }
268
269        Ok(PeHeader { dos_header, nt_header, coff_header, optional_header })
270    }
271
272    /// 生成节列表
273    pub fn build_sections(&mut self) -> Vec<PeSection> {
274        let mut sections = Vec::new();
275        let mut next_virtual_address = 0x1000;
276        let mut next_raw_data_offset = 0x200;
277
278        // 预先计算节的 RVA，以便 fix_code_relocations 使用
279        // 注意：这里的逻辑必须与下面的节添加顺序完全一致
280        let mut code_rva = None;
281        let mut data_rva = None;
282
283        if self.code.is_some() {
284            code_rva = Some(next_virtual_address);
285            next_virtual_address += 0x1000;
286        }
287        if self.data.is_some() {
288            data_rva = Some(next_virtual_address);
289            next_virtual_address += 0x1000;
290        }
291
292        // 重置以进行实际构建
293        next_virtual_address = 0x1000;
294
295        // 添加代码节
296        if let Some(code) = &self.code {
297            let mut code_data = code.clone();
298
299            // 修复重定位
300            self.fix_code_relocations_with_rvas(&mut code_data, code_rva.unwrap_or(0x1000), data_rva);
301
302            // 对齐到 512 字节
303            let raw_size = ((code_data.len() + 0x1FF) / 0x200 * 0x200) as u32;
304            while code_data.len() < raw_size as usize {
305                code_data.push(0);
306            }
307
308            let text_section = PeSection {
309                name: ".text".to_string(),
310                virtual_size: 0x1000,
311                virtual_address: next_virtual_address,
312                size_of_raw_data: raw_size,
313                pointer_to_raw_data: next_raw_data_offset,
314                pointer_to_relocations: 0,
315                pointer_to_line_numbers: 0,
316                number_of_relocations: 0,
317                number_of_line_numbers: 0,
318                characteristics: 0x60000020,
319                data: code_data,
320            };
321            sections.push(text_section);
322            next_virtual_address += 0x1000;
323            next_raw_data_offset += raw_size;
324        }
325
326        // 添加数据节
327        if let Some(data) = &self.data {
328            let mut data_bytes = data.clone();
329            // 对齐到 512 字节
330            let raw_size = ((data_bytes.len() + 0x1FF) / 0x200 * 0x200) as u32;
331            while data_bytes.len() < raw_size as usize {
332                data_bytes.push(0);
333            }
334
335            let data_section = PeSection {
336                name: ".data".to_string(),
337                virtual_size: 0x1000,
338                virtual_address: next_virtual_address,
339                size_of_raw_data: raw_size,
340                pointer_to_raw_data: next_raw_data_offset,
341                pointer_to_relocations: 0,
342                pointer_to_line_numbers: 0,
343                number_of_relocations: 0,
344                number_of_line_numbers: 0,
345                characteristics: 0xC0000040,
346                data: data_bytes,
347            };
348            sections.push(data_section);
349            next_virtual_address += 0x1000;
350            next_raw_data_offset += raw_size;
351        }
352
353        // 添加导入表节（如果有导入）
354        if !self.imports.is_empty() {
355            let mut idata_section = self.build_import_section();
356            idata_section.virtual_address = next_virtual_address;
357            idata_section.pointer_to_raw_data = next_raw_data_offset;
358            sections.push(idata_section);
359        }
360
361        sections
362    }
363
364    /// 构建导入表节
365    fn build_import_section(&self) -> PeSection {
366        // 不在这里填充数据，让 write_import_table 方法来处理。
367        // 注意：write_import_table 使用“兼容模式”（见上文），在 x64 下 IAT 初始填入 Hint/Name 的 RVA。
368        PeSection {
369            name: ".idata".to_string(),
370            virtual_size: 0x1000,
371            virtual_address: 0x3000, // 这个值会在 build_sections 中被覆盖
372            size_of_raw_data: 0x200,
373            pointer_to_raw_data: 0x600, // 这个值会在 build_sections 中被覆盖
374            pointer_to_relocations: 0,
375            pointer_to_line_numbers: 0,
376            number_of_relocations: 0,
377            number_of_line_numbers: 0,
378            characteristics: 0xC0000040, // IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE
379            data: Vec::new(),            // 空数据，由 write_import_table 方法填充
380        }
381    }
382
383    /// 修复代码中的重定位
384    fn fix_code_relocations_with_rvas(&self, code: &mut Vec<u8>, code_section_rva: u32, data_section_rva: Option<u32>) {
385        // 查找 CALL 指令的占位符并替换为正确的地址
386        let mut i = 0;
387
388        let arch = self.architecture.as_ref().unwrap_or(&Architecture::X86).clone();
389        let pointer_size: usize = if arch == Architecture::X86_64 { 8 } else { 4 };
390        let image_base: u64 = self.image_base.as_ref().copied().unwrap_or(match arch {
391            Architecture::X86 => 0x400000,
392            Architecture::X86_64 => 0x140000000,
393            _ => 0x400000,
394        });
395
396        // 计算 .idata 节的起始 RVA（与 build_sections 一致）
397        let mut import_rva_base: u64 = 0x1000; // 起始 RVA
398        if self.code.is_some() {
399            import_rva_base += 0x1000;
400        }
401        if self.data.is_some() {
402            import_rva_base += 0x1000;
403        }
404
405        // 计算 IAT 的实际位置（与 write_import_table 中的计算完全一致）
406        let mut current_rva: u64 = import_rva_base;
407
408        // 计算 IAT 的偏移量（跟 write_import_table 中的计算保持一致）
409        if !self.imports.is_empty() {
410            // 导入描述符表大小：(导入DLL数量 + 1) * 20字节
411            current_rva += ((self.imports.len() + 1) * 20) as u64;
412
413            // DLL 名称大小
414            for (dll_name, _) in &self.imports {
415                current_rva += (dll_name.len() + 1) as u64; // 包括空终止符
416            }
417
418            // 对齐到 2 字节，与 write_import_table 保持一致
419            if current_rva % 2 != 0 {
420                current_rva += 1;
421            }
422
423            // 函数名称（Hint(2) + Name + '\0'），逐个累加
424            for (_, functions) in &self.imports {
425                for function in functions {
426                    // 名称按 2 字节对齐
427                    if current_rva % 2 != 0 {
428                        current_rva += 1;
429                    }
430                    current_rva += 2 + (function.len() + 1) as u64;
431                }
432            }
433
434            // INT（OriginalFirstThunk）按 pointer_size 字节对齐并分配
435            if current_rva % pointer_size as u64 != 0 {
436                current_rva = (current_rva + pointer_size as u64 - 1) & !(pointer_size as u64 - 1);
437            }
438            for (_, functions) in &self.imports {
439                current_rva += ((functions.len() as u64) + 1) * pointer_size as u64;
440                // 包含终止符
441            }
442
443            // IAT（FirstThunk）按 pointer_size 字节对齐并分配 -> 这里就是 IAT 起始 RVA
444            if current_rva % pointer_size as u64 != 0 {
445                current_rva = (current_rva + pointer_size as u64 - 1) & !(pointer_size as u64 - 1);
446            }
447        }
448
449        let iat_start_rva = current_rva;
450
451        while i < code.len() {
452            // x64: 修补多种指令的 [rip+disp32] 寻址
453            if arch == Architecture::X86_64 && i + 2 < code.len() {
454                let mut pos = i;
455                let mut _rex_prefix = None;
456                // 检查是否为 REX 前缀 (0x40 - 0x4F)
457                if code[pos] >= 0x40 && code[pos] <= 0x4F {
458                    _rex_prefix = Some(code[pos]);
459                    pos += 1;
460                }
461
462                if pos + 1 < code.len() {
463                    let mut opcode = code[pos] as u32;
464                    let mut modrm_pos = pos + 1;
465
466                    // 处理 2 字节指令 (0x0F 开头)
467                    if opcode == 0x0F && pos + 2 < code.len() {
468                        opcode = (opcode << 8) | (code[pos + 1] as u32);
469                        modrm_pos = pos + 2;
470                    }
471
472                    if modrm_pos < code.len() {
473                        let modrm = code[modrm_pos];
474
475                        // 检查 ModR/M 是否为 [RIP + disp32] (mod=00, rm=101)
476                        if (modrm & 0xC7) == 0x05 && modrm_pos + 4 < code.len() {
477                            let disp_offset = modrm_pos + 1;
478                            let next_instr_offset = disp_offset + 4;
479
480                            // 分支 1: 间接调用/跳转 (CALL/JMP [RIP+disp32]) -> 指向导入表 (IAT)
481                            // CALL: FF /2, JMP: FF /4
482                            if opcode == 0xFF && ((modrm >> 3) & 7 == 2 || (modrm >> 3) & 7 == 4) {
483                                // 读取当前的 displacement 作为导入索引
484                                let import_index = u32::from_le_bytes([
485                                    code[disp_offset],
486                                    code[disp_offset + 1],
487                                    code[disp_offset + 2],
488                                    code[disp_offset + 3],
489                                ]) as usize;
490
491                                // 计算实际的 IAT RVA
492                                let mut current_iat_offset = 0;
493                                let mut found_rva = None;
494                                let mut flat_idx = 0;
495
496                                'import_search: for (_, functions) in &self.imports {
497                                    for _ in functions {
498                                        if flat_idx == import_index {
499                                            found_rva = Some(iat_start_rva + current_iat_offset);
500                                            break 'import_search;
501                                        }
502                                        flat_idx += 1;
503                                        current_iat_offset += pointer_size as u64;
504                                    }
505                                    current_iat_offset += pointer_size as u64; // 跳过 DLL 终止符
506                                }
507
508                                let target_rva =
509                                    found_rva.unwrap_or(iat_start_rva + (import_index as u64 * pointer_size as u64));
510                                let rip_rva = (code_section_rva as u64) + (next_instr_offset as u64);
511                                let disp_i32 = (target_rva as i64 - rip_rva as i64) as i32;
512
513                                code[disp_offset..disp_offset + 4].copy_from_slice(&disp_i32.to_le_bytes());
514
515                                // 调试输出
516                                tracing::trace!(
517                                    "IMPORT 修补(RVA): i={}, rip_rva={:08X}, target_rva={:08X}, disp={:08X}, import_index={}",
518                                    i,
519                                    rip_rva as u32,
520                                    target_rva as u32,
521                                    disp_i32 as u32,
522                                    import_index
523                                );
524
525                                i = next_instr_offset;
526                                continue;
527                            }
528                            // 分支 2: 数据访问指令 -> 指向数据段 (.data)
529                            // 只要是 [RIP + disp32] 且不是上面处理过的导入调用，都视为数据访问
530                            else {
531                                // 读取当前 displacement 作为数据段内的偏移
532                                let current_disp = u32::from_le_bytes([
533                                    code[disp_offset],
534                                    code[disp_offset + 1],
535                                    code[disp_offset + 2],
536                                    code[disp_offset + 3],
537                                ]);
538
539                                let target_rva: u64 = (data_section_rva.unwrap_or(0) as u64) + (current_disp as u64);
540                                let rip_rva: u64 = (code_section_rva as u64) + (next_instr_offset as u64);
541                                let disp_i32 = (target_rva as i64 - rip_rva as i64) as i32;
542
543                                code[disp_offset..disp_offset + 4].copy_from_slice(&disp_i32.to_le_bytes());
544
545                                // 调试输出
546                                tracing::trace!(
547                                    "DATA 修补(RVA): i={}, opcode={:04X}, rip_rva={:08X}, target_rva={:08X}, disp={:08X}",
548                                    i,
549                                    opcode,
550                                    rip_rva as u32,
551                                    target_rva as u32,
552                                    disp_i32 as u32
553                                );
554
555                                i = next_instr_offset;
556                                continue;
557                            }
558                        }
559                    }
560                }
561            }
562
563            // x86 (32-bit) 的处理逻辑保持不变
564            if arch == Architecture::X86 && i + 1 < code.len() && code[i] == 0xFF && code[i + 1] == 0x15 && i + 5 < code.len() {
565                // 读取当前的 displacement 作为导入索引
566                let import_index = u32::from_le_bytes([code[i + 2], code[i + 3], code[i + 4], code[i + 5]]) as usize;
567
568                // 计算实际的 IAT RVA
569                let mut current_iat_offset = 0;
570                let mut found_rva = None;
571                let mut flat_idx = 0;
572
573                'import_search_x86: for (_, functions) in &self.imports {
574                    for _ in functions {
575                        if flat_idx == import_index {
576                            found_rva = Some(iat_start_rva + current_iat_offset);
577                            break 'import_search_x86;
578                        }
579                        flat_idx += 1;
580                        current_iat_offset += pointer_size as u64;
581                    }
582                    current_iat_offset += pointer_size as u64;
583                }
584
585                let target_rva = found_rva.unwrap_or(iat_start_rva + (import_index as u64 * pointer_size as u64));
586                let target_va = image_base + target_rva;
587
588                let disp: u32 = target_va as u32;
589                let address_bytes = disp.to_le_bytes();
590                code[i + 2..i + 6].copy_from_slice(&address_bytes);
591                i += 6;
592            }
593            // 查找直接 CALL 指令 (0xE8) - 保留用于内部函数调用
594            else if code[i] == 0xE8 && i + 4 < code.len() {
595                // 检查是否是占位符地址 (0x00000000)
596                let placeholder = u32::from_le_bytes([code[i + 1], code[i + 2], code[i + 3], code[i + 4]]);
597
598                if placeholder == 0x00000000 {
599                    // 对于直接调用，计算相对偏移（用于内部函数）
600                    // 这里可以根据需要实现内部函数的重定位
601                    // 暂时跳过
602                }
603
604                i += 5; // 跳过整个直接 CALL 指令
605            }
606            // 查找 PUSH imm32 指令 (0x68) - 用于 push_label 或数据地址占位符
607            else if code[i] == 0x68 && i + 4 < code.len() {
608                // 读取立即数
609                let imm = u32::from_le_bytes([code[i + 1], code[i + 2], code[i + 3], code[i + 4]]);
610
611                if arch == Architecture::X86 {
612                    // 如果这是一个 0 占位符，检查是否应该修补为 .data 开始地址
613                    if imm == 0 {
614                        // 检查前面是否有 push_imm8 + push_label 的模式
615                        // push_imm(msg_len) 编译为 6a <len>（2字节）
616                        // push_label("msg") 编译为 68 00 00 00 00（5字节）
617                        let mut should_patch = false;
618
619                        // 向前查找，看是否有 push imm8 <msg_len> 紧接着当前的 push 0 的模式
620                        if i >= 2 {
621                            // 确保有足够的空间向前查找
622                            // 查找前面的 push imm8 指令 (6a)
623                            let prev_push_pos = i - 2;
624                            if prev_push_pos < code.len() && code[prev_push_pos] == 0x6a {
625                                let prev_imm8 = code[prev_push_pos + 1] as u32;
626
627                                // 检查这个立即数是否等于消息长度
628                                let msg_len = if let Some(data) = &self.data {
629                                    data.iter().position(|&b| b == 0).map(|p| p as u32).unwrap_or(0)
630                                }
631                                else {
632                                    0
633                                };
634
635                                if prev_imm8 == msg_len {
636                                    should_patch = true;
637                                }
638                            }
639                        }
640
641                        if should_patch {
642                            // 数据段起始 VA
643                            let data_section_va: u64 = image_base + (data_section_rva.unwrap_or(0) as u64);
644                            let addr_u32 = data_section_va as u32;
645                            code[i + 1..i + 5].copy_from_slice(&addr_u32.to_le_bytes());
646                        }
647                    }
648                }
649
650                i += 5; // 跳过整个 PUSH 指令
651            }
652            else {
653                i += 1;
654            }
655        }
656    }
657
658    /// 生成 PE 文件字节数组
659    pub fn generate(&mut self) -> Result<Vec<u8>, GaiaError> {
660        // 构建节
661        self.sections = self.build_sections(); // Populate sections first
662
663        // 构建头部
664        let header = self.build_header()?;
665
666        // 构建导入表
667        let mut import_table = ImportTable::new();
668        for (dll_name, functions) in &self.imports {
669            let entry = ImportEntry { dll_name: dll_name.clone(), functions: functions.clone() };
670            import_table.entries.push(entry);
671        }
672
673        // 创建 PE 程序
674        let program = PeProgram {
675            header,
676            sections: self.sections.clone(),
677            imports: import_table,
678            exports: crate::types::tables::ExportTable::new(),
679        };
680
681        // 写入到字节数组
682        let mut buffer = Vec::new();
683        let cursor = Cursor::new(&mut buffer);
684        let mut writer = ExeWriter::new(cursor);
685        writer.write_program(&program)?;
686
687        Ok(buffer)
688    }
689}
690
691impl Default for PeBuilder {
692    fn default() -> Self {
693        Self::new()
694    }
695}
pe_assembler/helpers/builder/mod.rs

pe_assembler/helpers/builder/
mod.rs
