Skip to content

Commit f17c6ca

Browse files
codeamtcodeamt
committed
global ignored DS_Store, fixed clippy linting errors
1 parent 85edb7e commit f17c6ca

File tree

4 files changed

+82
-97
lines changed

4 files changed

+82
-97
lines changed

‎.DS_Store‎

-8 KB
Binary file not shown.

‎src/bin/main.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1+
use log::LevelFilter;
2+
use pe_packer::packer::{EncryptionAlgorithm, PEPacker, PackerConfig};
3+
use pe_packer::utils::init_logger;
14
use std::env;
25
use std::path::Path;
3-
use pe_packer::packer::{PEPacker, PackerConfig, EncryptionAlgorithm};
4-
use pe_packer::utils::{init_logger};
5-
use log::LevelFilter;
66

77
fn print_usage(program: &str) {
88
eprintln!("Usage: {} <input.exe> <output.exe> [xor|aes|none]", program);
@@ -12,7 +12,7 @@ fn main() {
1212
init_logger(LevelFilter::Info);
1313

1414
let args: Vec<String> = env::args().collect();
15-
let program = args.get(0).map(String::as_str).unwrap_or("pe-packer");
15+
let program = args.first().map(String::as_str).unwrap_or("pe-packer");
1616

1717
if args.len() < 3 {
1818
print_usage(program);

‎src/packer/core.rs‎

Lines changed: 25 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,15 @@
1-
use std::fs::File;
2-
use std::io::{Read, Write};
3-
use std::path::Path;
4-
use crate::pe::parser::PEParser;
5-
use crate::pe::builder::PEBuilder;
61
use super::encryption::EncryptionAlgorithm;
72
use super::stub_generator::StubGenerator;
3+
use crate::pe::builder::PEBuilder;
4+
use crate::pe::parser::PEParser;
85
use crate::utils::error::PackerError;
9-
use rand::Rng;
106
use goblin::pe::section_table::{
11-
IMAGE_SCN_CNT_INITIALIZED_DATA,
12-
IMAGE_SCN_MEM_READ,
13-
IMAGE_SCN_MEM_WRITE,
14-
IMAGE_SCN_MEM_EXECUTE,
7+
IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE,
158
};
9+
use rand::Rng;
10+
use std::fs::File;
11+
use std::io::{Read, Write};
12+
use std::path::Path;
1613

1714
/// Configuration for the PE packer
1815
#[derive(Debug, Clone)]
@@ -88,7 +85,10 @@ impl PEPacker {
8885

8986
let stub_generator = StubGenerator::new(&config, &key);
9087

91-
PEPacker { config, stub_generator }
88+
PEPacker {
89+
config,
90+
stub_generator,
91+
}
9292
}
9393

9494
/// Pack a PE file from disk
@@ -111,7 +111,8 @@ impl PEPacker {
111111
let mut output = File::create(output_path)
112112
.map_err(|e| PackerError::IoError(format!("Failed to create output file: {}", e)))?;
113113

114-
output.write_all(&packed_data)
114+
output
115+
.write_all(&packed_data)
115116
.map_err(|e| PackerError::IoError(format!("Failed to write output file: {}", e)))?;
116117

117118
Ok(())
@@ -132,7 +133,7 @@ impl PEPacker {
132133
fn process_pe_file(
133134
&self,
134135
parser: &PEParser,
135-
original_data: &[u8]
136+
original_data: &[u8],
136137
) -> Result<Vec<u8>, PackerError> {
137138
// Extract and encrypt sections
138139
let encrypted_sections = self.process_sections(parser, original_data)?;
@@ -142,11 +143,8 @@ impl PEPacker {
142143

143144
// Rebuild PE with new structure
144145
let pe_builder = PEBuilder::new(parser);
145-
let packed_pe = pe_builder.build_packed_pe(
146-
&encrypted_sections,
147-
&stub_data,
148-
&self.config
149-
)?;
146+
let packed_pe =
147+
pe_builder.build_packed_pe(&encrypted_sections, &stub_data, &self.config)?;
150148

151149
Ok(packed_pe)
152150
}
@@ -164,14 +162,13 @@ impl PEPacker {
164162

165163
if (section.characteristics & IMAGE_SCN_MEM_EXECUTE) != 0 {
166164
// Encrypt executable sections
167-
let encrypted_data = self.encrypt_section(&section_data)?;
165+
let encrypted_data = self.encrypt_section(section_data)?;
168166

169167
// Create modified section header
170168
let mut new_header = section.clone();
171169
new_header.name = self.generate_section_name();
172-
new_header.characteristics = IMAGE_SCN_CNT_INITIALIZED_DATA
173-
| IMAGE_SCN_MEM_READ
174-
| IMAGE_SCN_MEM_WRITE;
170+
new_header.characteristics =
171+
IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE;
175172

176173
processed_sections.push((encrypted_data, new_header));
177174
} else {
@@ -187,26 +184,25 @@ impl PEPacker {
187184
fn encrypt_section(&self, data: &[u8]) -> Result<Vec<u8>, PackerError> {
188185
use super::encryption;
189186

190-
let key = self.config.key.as_ref()
191-
.ok_or_else(|| PackerError::EncryptionError("No encryption key provided".to_string()))?;
187+
let key = self.config.key.as_ref().ok_or_else(|| {
188+
PackerError::EncryptionError("No encryption key provided".to_string())
189+
})?;
192190

193191
match self.config.encryption {
194192
EncryptionAlgorithm::Xor => {
195193
let mut encrypted = data.to_vec();
196194
encryption::xor(&mut encrypted, key);
197195
Ok(encrypted)
198-
},
199-
EncryptionAlgorithm::Aes => {
200-
Ok(encryption::aes_encrypt(data, key))
201-
},
196+
}
197+
EncryptionAlgorithm::Aes => Ok(encryption::aes_encrypt(data, key)),
202198
EncryptionAlgorithm::None => Ok(data.to_vec()),
203199
}
204200
}
205201

206202
/// Generate a section name based on configuration
207203
fn generate_section_name(&self) -> [u8; 8] {
208204
if self.config.randomize_sections {
209-
use rand::{Rng, distributions::Alphanumeric};
205+
use rand::{distributions::Alphanumeric, Rng};
210206
let mut rng = rand::thread_rng();
211207
let name: String = (0..8).map(|_| rng.sample(Alphanumeric) as char).collect();
212208
let mut arr = [0u8; 8];

‎src/pe/builder.rs‎

Lines changed: 53 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,12 @@
33
//! This module handles the reconstruction of PE files with modified sections,
44
//! encrypted payloads, and decryption stubs.
55
6-
use crate::utils::error::PackerError;
7-
use crate::pe::parser::PEParser;
86
use crate::packer::PackerConfig;
7+
use crate::pe::parser::PEParser;
8+
use crate::utils::error::PackerError;
99
// use goblin::pe::header; // not currently needed
1010
use goblin::pe::section_table::{
11-
IMAGE_SCN_CNT_CODE,
12-
IMAGE_SCN_MEM_EXECUTE,
13-
IMAGE_SCN_MEM_READ,
14-
IMAGE_SCN_MEM_WRITE,
11+
IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE,
1512
};
1613
// use std::io::Write; // not currently needed
1714

@@ -69,12 +66,8 @@ impl<'a> PEBuilder<'a> {
6966

7067
// Step 4: Build modified section headers
7168
// Update section headers for encrypted content and add stub section
72-
let section_headers = self.build_section_headers(
73-
sections,
74-
stub,
75-
stub_section_index,
76-
config,
77-
)?;
69+
let section_headers =
70+
self.build_section_headers(sections, stub, stub_section_index, config)?;
7871

7972
// Step 5: Write section headers to output
8073
for header in &section_headers {
@@ -94,13 +87,11 @@ impl<'a> PEBuilder<'a> {
9487

9588
// Verify section is written at expected offset
9689
if packed_pe.len() as u32 != header.pointer_to_raw_data {
97-
return Err(PackerError::PeBuildError(
98-
format!(
99-
"Section alignment mismatch: expected offset {}, got {}",
100-
header.pointer_to_raw_data,
101-
packed_pe.len()
102-
)
103-
));
90+
return Err(PackerError::PeBuildError(format!(
91+
"Section alignment mismatch: expected offset {}, got {}",
92+
header.pointer_to_raw_data,
93+
packed_pe.len()
94+
)));
10495
}
10596

10697
// Write the encrypted section data
@@ -114,15 +105,15 @@ impl<'a> PEBuilder<'a> {
114105
// Verify stub section offset
115106
if packed_pe.len() as u32 != section_headers.last().unwrap().pointer_to_raw_data {
116107
return Err(PackerError::PeBuildError(
117-
"Stub section offset mismatch".to_string()
108+
"Stub section offset mismatch".to_string(),
118109
));
119110
}
120111

121112
packed_pe.extend_from_slice(stub);
122113

123114
// Step 9: Update size of image in optional header
124115
// This must reflect the total virtual size of all sections
125-
self.update_image_size(&mut packed_pe, &section_headers)?;
116+
self.update_image_size(&mut packed_pe[..], &section_headers)?;
126117

127118
Ok(packed_pe)
128119
}
@@ -136,15 +127,15 @@ impl<'a> PEBuilder<'a> {
136127
const DOS_HEADER_SIZE: usize = 64;
137128

138129
// Get original DOS header from parser
139-
let dos_data = self.parser.get_dos_header()
140-
.ok_or_else(|| PackerError::PeBuildError(
141-
"Failed to extract DOS header".to_string()
142-
))?;
130+
let dos_data = self
131+
.parser
132+
.get_dos_header()
133+
.ok_or_else(|| PackerError::PeBuildError("Failed to extract DOS header".to_string()))?;
143134

144135
// Ensure we have at least the minimum DOS header
145136
if dos_data.len() < DOS_HEADER_SIZE {
146137
return Err(PackerError::PeBuildError(
147-
"DOS header too small".to_string()
138+
"DOS header too small".to_string(),
148139
));
149140
}
150141

@@ -173,14 +164,13 @@ impl<'a> PEBuilder<'a> {
173164
// COFF header is 20 bytes
174165
const COFF_HEADER_SIZE: usize = 20;
175166

176-
let coff_data = self.parser.get_coff_header()
177-
.ok_or_else(|| PackerError::PeBuildError(
178-
"Failed to extract COFF header".to_string()
179-
))?;
167+
let coff_data = self.parser.get_coff_header().ok_or_else(|| {
168+
PackerError::PeBuildError("Failed to extract COFF header".to_string())
169+
})?;
180170

181171
if coff_data.len() < COFF_HEADER_SIZE {
182172
return Err(PackerError::PeBuildError(
183-
"COFF header too small".to_string()
173+
"COFF header too small".to_string(),
184174
));
185175
}
186176

@@ -201,16 +191,15 @@ impl<'a> PEBuilder<'a> {
201191
/// The optional header contains critical information like entry point and image base.
202192
/// Returns the index where the stub section will be placed.
203193
fn copy_optional_header(&self, output: &mut Vec<u8>) -> Result<usize, PackerError> {
204-
let optional_data = self.parser.get_optional_header()
205-
.ok_or_else(|| PackerError::PeBuildError(
206-
"Failed to extract optional header".to_string()
207-
))?;
194+
let optional_data = self.parser.get_optional_header().ok_or_else(|| {
195+
PackerError::PeBuildError("Failed to extract optional header".to_string())
196+
})?;
208197

209198
// Optional header is typically 96 (32-bit) or 112 (64-bit) bytes
210199
// Validate minimum size
211200
if optional_data.len() < 96 {
212201
return Err(PackerError::PeBuildError(
213-
"Optional header too small".to_string()
202+
"Optional header too small".to_string(),
214203
));
215204
}
216205

@@ -220,7 +209,7 @@ impl<'a> PEBuilder<'a> {
220209

221210
// Calculate stub section index (placed after all other sections)
222211
// This is used later to update the entry point
223-
let section_count = self.parser.section_count() as usize;
212+
let section_count = self.parser.section_count();
224213
Ok(section_count)
225214
}
226215

@@ -267,25 +256,23 @@ impl<'a> PEBuilder<'a> {
267256
if stub_section_index == sections.len() {
268257
current_offset = Self::align_value(current_offset, file_alignment);
269258

270-
let mut stub_header = goblin::pe::section_table::SectionTable::default();
271-
272-
// Set stub section name
273-
stub_header.name = *b".stub\0\0\0";
274-
stub_header.virtual_size = stub.len() as u32;
275-
stub_header.virtual_address = Self::align_value(
276-
sections.last()
277-
.map(|(_, h)| h.virtual_address + h.virtual_size)
278-
.unwrap_or(0x1000),
279-
0x1000,
280-
);
281-
stub_header.size_of_raw_data = stub.len() as u32;
282-
stub_header.pointer_to_raw_data = current_offset;
283-
284-
// Mark as executable code
285-
stub_header.characteristics =
286-
IMAGE_SCN_CNT_CODE |
287-
IMAGE_SCN_MEM_EXECUTE |
288-
IMAGE_SCN_MEM_READ;
259+
let stub_header = goblin::pe::section_table::SectionTable {
260+
name: *b".stub\0\0\0",
261+
virtual_size: stub.len() as u32,
262+
virtual_address: Self::align_value(
263+
sections
264+
.last()
265+
.map(|(_, h)| h.virtual_address + h.virtual_size)
266+
.unwrap_or(0x1000),
267+
0x1000,
268+
),
269+
size_of_raw_data: stub.len() as u32,
270+
pointer_to_raw_data: current_offset,
271+
characteristics: IMAGE_SCN_CNT_CODE
272+
| IMAGE_SCN_MEM_EXECUTE
273+
| IMAGE_SCN_MEM_READ,
274+
..Default::default()
275+
};
289276

290277
headers.push(stub_header);
291278
}
@@ -357,7 +344,7 @@ impl<'a> PEBuilder<'a> {
357344
if alignment == 0 {
358345
return value;
359346
}
360-
((value + alignment - 1) / alignment) * alignment
347+
value.div_ceil(alignment)
361348
}
362349

363350
/// Get the file alignment value (typically 0x200)
@@ -375,9 +362,10 @@ impl<'a> PEBuilder<'a> {
375362

376363
// Ensure we don't exceed maximum sections
377364
if count > 32 {
378-
return Err(PackerError::PeBuildError(
379-
format!("Too many sections: {}", count)
380-
));
365+
return Err(PackerError::PeBuildError(format!(
366+
"Too many sections: {}",
367+
count
368+
)));
381369
}
382370

383371
Ok(count)
@@ -386,11 +374,12 @@ impl<'a> PEBuilder<'a> {
386374
/// Update the size of image field in the optional header
387375
fn update_image_size(
388376
&self,
389-
output: &mut Vec<u8>,
377+
output: &mut [u8],
390378
headers: &[goblin::pe::section_table::SectionTable],
391379
) -> Result<(), PackerError> {
392380
// Calculate total virtual size
393-
let total_size = headers.iter()
381+
let total_size = headers
382+
.iter()
394383
.map(|h| h.virtual_address + h.virtual_size)
395384
.max()
396385
.unwrap_or(0);
@@ -404,12 +393,12 @@ impl<'a> PEBuilder<'a> {
404393

405394
if output.len() < size_of_image_offset + 4 {
406395
return Err(PackerError::PeBuildError(
407-
"Output buffer too small to update size of image".to_string()
396+
"Output buffer too small to update size of image".to_string(),
408397
));
409398
}
410399

411400
// Update the size of image field
412-
output[size_of_image_offset..size_of_image_offset+4]
401+
output[size_of_image_offset..size_of_image_offset + 4]
413402
.copy_from_slice(&aligned_size.to_le_bytes());
414403

415404
Ok(())

0 commit comments

Comments
 (0)