Fixed more linting errors for rust src

codeamt · codeamt · commit ea52c824fa23 · 2025-10-18T03:35:17.000-04:00
diff --git a/.gitignore b/.gitignore
@@ -54,4 +54,5 @@ target/wheels/
 .ssh/ssh-agent
 .ssh/ssh-agent.pid
 .ssh/ssh-agent.env
-.DS_Store
+.DS_Store
+**/.DS_Store
diff --git a/src/lib.rs b/src/lib.rs
@@ -6,4 +6,3 @@ pub mod pe;
 #[cfg(feature = "python")]
 pub mod python;
 pub mod utils;
-
diff --git a/src/packer/core.rs b/src/packer/core.rs
@@ -236,4 +236,4 @@ mod tests {
 
         assert!(packer.config().key.is_some());
     }
-}
+}
diff --git a/src/packer/encryption.rs b/src/packer/encryption.rs
@@ -1,6 +1,6 @@
 use aes::Aes128;
-use cbc::{Encryptor as Aes128CbcEnc, Decryptor as Aes128CbcDec};
-use cbc::cipher::{KeyIvInit, block_padding::Pkcs7, BlockEncryptMut, BlockDecryptMut};
+use cbc::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
+use cbc::{Decryptor as Aes128CbcDec, Encryptor as Aes128CbcEnc};
 use rand::Rng;
 
 /// Encryption algorithms supported by the packer
@@ -170,7 +170,7 @@ mod tests {
 
         assert!(high_entropy > low_entropy);
         assert!(high_entropy > 7.0); // Random data should have ~8 bits of entropy
-        assert!(low_entropy < 1.0);  // Repeated data should have very low entropy
+        assert!(low_entropy < 1.0); // Repeated data should have very low entropy
     }
 
     #[test]
@@ -184,4 +184,4 @@ mod tests {
         let none_key = generate_key(EncryptionAlgorithm::None);
         assert!(none_key.is_empty());
     }
-}
+}
diff --git a/src/packer/mod.rs b/src/packer/mod.rs
@@ -9,5 +9,5 @@ pub use encryption::EncryptionAlgorithm;
 pub use stub_generator::StubGenerator;
 
 // Re-export for convenience
+pub use crate::pe::builder::PEBuilder;
 pub use crate::pe::parser::PEParser;
-pub use crate::pe::builder::PEBuilder;
diff --git a/src/packer/stub_generator.rs b/src/packer/stub_generator.rs
@@ -1,5 +1,5 @@
-use rand::Rng;
 use super::encryption::EncryptionAlgorithm;
+use rand::Rng;
 
 /// Generates decryption stubs for packed executables
 pub struct StubGenerator {
@@ -120,14 +120,14 @@ impl StubGenerator {
                 stub.extend_from_slice(b"XOR_DECRYPTION_ROUTINE");
                 // In a real implementation, this would contain actual x86/x64 assembly
                 // for XOR decryption
-            },
+            }
             EncryptionAlgorithm::Aes => {
                 stub.extend_from_slice(b"AES_DECRYPTION_ROUTINE");
                 // AES decryption assembly would go here
-            },
+            }
             EncryptionAlgorithm::None => {
                 stub.extend_from_slice(b"NO_DECRYPTION_NEEDED");
-            },
+            }
         }
 
         stub.extend_from_slice(b"DECRYPT_LOGIC_END");
@@ -149,7 +149,7 @@ impl StubGenerator {
 
     /// Generate a unique stub ID for tracking
     pub fn generate_stub_id(&self) -> String {
-        use sha2::{Sha256, Digest};
+        use sha2::{Digest, Sha256};
 
         let mut hasher = Sha256::new();
         hasher.update(&self.key);
@@ -163,7 +163,11 @@ impl StubGenerator {
     pub fn estimate_stub_size(&self) -> usize {
         let base_size = 64; // Header, flags, etc.
         let anti_debug_size = if self.config.add_anti_debug { 128 } else { 0 };
-        let metadata_size = if self.config.add_benign_metadata { 256 } else { 0 };
+        let metadata_size = if self.config.add_benign_metadata {
+            256
+        } else {
+            0
+        };
         let padding_size = 40; // Average padding
 
         base_size + anti_debug_size + metadata_size + padding_size
@@ -233,4 +237,4 @@ mod tests {
         // Should contain AES marker
         assert!(stub_str.contains("AES"));
     }
-}
+}
diff --git a/src/pe/builder.rs b/src/pe/builder.rs
@@ -268,9 +268,7 @@ impl<'a> PEBuilder<'a> {
                 ),
                 size_of_raw_data: stub.len() as u32,
                 pointer_to_raw_data: current_offset,
-                characteristics: IMAGE_SCN_CNT_CODE
-                    | IMAGE_SCN_MEM_EXECUTE
-                    | IMAGE_SCN_MEM_READ,
+                characteristics: IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ,
                 ..Default::default()
             };
 
@@ -330,10 +328,8 @@ impl<'a> PEBuilder<'a> {
     fn pad_to_alignment(&self, output: &mut Vec<u8>, alignment: u32) -> Result<(), PackerError> {
         let current_size = output.len() as u32;
         let aligned_size = Self::align_value(current_size, alignment);
-        let padding = aligned_size - current_size;
-
-        if padding > 0 {
-            output.extend(vec![0u8; padding as usize]);
+        if aligned_size > current_size {
+            output.resize(aligned_size as usize, 0u8);
         }
 
         Ok(())
@@ -421,4 +417,4 @@ mod tests {
     fn test_file_alignment() {
         assert_eq!(PEBuilder::get_file_alignment(), 0x200);
     }
-}
+}
diff --git a/src/pe/mod.rs b/src/pe/mod.rs
@@ -1,9 +1,9 @@
 //! PE file parsing and building functionality
 
-pub mod parser;
 pub mod builder;
+pub mod parser;
 mod structures;
 
-pub use parser::PEParser;
 pub use builder::PEBuilder;
-pub use structures::*;
+pub use parser::PEParser;
+pub use structures::*;
diff --git a/src/pe/parser.rs b/src/pe/parser.rs
@@ -1,6 +1,6 @@
-use goblin::pe::PE;
 use crate::utils::error::PackerError;
 use goblin::pe::characteristic::IMAGE_FILE_DLL;
+use goblin::pe::PE;
 
 pub struct PEParser<'a> {
     pe: PE<'a>,
@@ -20,19 +20,22 @@ impl<'a> PEParser<'a> {
     pub fn extract_section_data<'b>(
         &self,
         section: &goblin::pe::section_table::SectionTable,
-        buffer: &'b [u8]
+        buffer: &'b [u8],
     ) -> &'b [u8] {
         let start = section.pointer_to_raw_data as usize;
         let size = section.size_of_raw_data as usize;
-        if start >= buffer.len() { return &buffer[0..0]; }
+        if start >= buffer.len() { 
+            return &buffer[0..0];
+        }
         let end = start.saturating_add(size).min(buffer.len());
         &buffer[start..end]
     }
 
     pub fn entry_point(&self) -> u32 {
         // goblin exposes optional_header as Option<OptionalHeader>
         // address_of_entry_point is in standard_fields
-        (self.pe
+        (self
+            .pe
             .header
             .optional_header
             .as_ref()
@@ -46,15 +49,23 @@ impl<'a> PEParser<'a> {
 
     // Helper: raw DOS header bytes (first 64 bytes)
     pub fn get_dos_header(&self) -> Option<&[u8]> {
-        if self.data.len() >= 64 { Some(&self.data[..64]) } else { None }
+        if self.data.len() >= 64 { 
+            Some(&self.data[..64]) 
+        } else { 
+            None 
+        }
     }
 
     // Helper: raw COFF header bytes (20 bytes at e_lfanew + 4)
     pub fn get_coff_header(&self) -> Option<&[u8]> {
         let pe_off = self.read_u32_le(0x3C)? as usize;
         let start = pe_off.checked_add(4)?;
         let end = start.checked_add(20)?;
-        if end <= self.data.len() { Some(&self.data[start..end]) } else { None }
+        if end <= self.data.len() { 
+            Some(&self.data[start..end]) 
+        } else { 
+            None 
+        }
     }
 
     // Helper: raw Optional header bytes (size given in COFF header at offset +16, u16)
@@ -66,7 +77,11 @@ impl<'a> PEParser<'a> {
         } else { return None };
         let opt_start = coff_start.checked_add(20)?;
         let opt_end = opt_start.checked_add(size_opt)?;
-        if opt_end <= self.data.len() { Some(&self.data[opt_start..opt_end]) } else { None }
+        if opt_end <= self.data.len() { 
+            Some(&self.data[opt_start..opt_end]) 
+        } else { 
+            None 
+        }
     }
 
     pub fn section_count(&self) -> usize {
@@ -85,4 +100,4 @@ impl<'a> PEParser<'a> {
             None
         }
     }
-}
+}
diff --git a/src/pe/structures.rs b/src/pe/structures.rs
@@ -1,4 +1,4 @@
-use serde::{Serialize, Deserialize};
+use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct DOSHeader {
@@ -44,4 +44,4 @@ pub struct PEFile {
     pub nt_headers: NTHeaders,
     pub sections: Vec<SectionHeader>,
     pub data: Vec<u8>,
-}
+}
diff --git a/src/python/bindings.rs b/src/python/bindings.rs
@@ -2,8 +2,10 @@ use pyo3::prelude::*;
 use pyo3::types::PyDict;
 use std::path::Path;
 
-use crate::packer::{PEPacker as RustPEPacker, PackerConfig as RustPackerConfig, EncryptionAlgorithm};
 use crate::packer::encryption;
+use crate::packer::{
+    EncryptionAlgorithm, PEPacker as RustPEPacker, PackerConfig as RustPackerConfig,
+};
 
 /// Python wrapper for EncryptionAlgorithm enum
 #[pyclass]
@@ -35,7 +37,9 @@ pub fn generate_xor_key_py(length: usize) -> PyResult<Vec<u8>> {
     use crate::packer::encryption::EncryptionAlgorithm as RustAlg;
     let mut key = encryption::generate_key(RustAlg::Xor);
     if length > 0 {
-        if key.len() < length { key.extend(std::iter::repeat(0u8).take(length - key.len())); }
+        if key.len() < length {
+            key.extend(std::iter::repeat(0u8).take(length - key.len()));
+        }
         key.truncate(length);
     }
     Ok(key)
@@ -173,17 +177,15 @@ impl PEPacker {
         let input = Path::new(input_path);
         let output = Path::new(output_path);
 
-        self.inner.pack(input, output)
-            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
-                format!("Packing failed: {}", e)
-            ))
+        self.inner.pack(input, output).map_err(|e| {
+            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Packing failed: {}", e))
+        })
     }
 
     pub fn pack_bytes(&self, data: Vec<u8>) -> PyResult<Vec<u8>> {
-        self.inner.pack_bytes(&data)
-            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
-                format!("Packing failed: {}", e)
-            ))
+        self.inner.pack_bytes(&data).map_err(|e| {
+            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Packing failed: {}", e))
+        })
     }
 
     pub fn get_config(&self) -> PackerConfig {
@@ -208,29 +210,38 @@ pub fn generate_training_dataset(
 
     if !input_path.is_dir() {
         return Err(PyErr::new::<pyo3::exceptions::PyFileNotFoundError, _>(
-            format!("Input directory not found: {}", input_dir)
+            format!("Input directory not found: {}", input_dir),
         ));
     }
 
-    fs::create_dir_all(output_path)
-        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
-            format!("Failed to create output directory: {}", e)
-        ))?;
+    fs::create_dir_all(output_path).map_err(|e| {
+        PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
+            "Failed to create output directory: {}",
+            e
+        ))
+    })?;
 
     let mut generated_files = Vec::new();
 
     // Process each PE file in input directory
-    for entry in fs::read_dir(input_path)
-        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
-            format!("Failed to read input directory: {}", e)
-        ))?
-    {
-        let entry = entry.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
-            format!("Error reading directory entry: {}", e)
-        ))?;
+    for entry in fs::read_dir(input_path).map_err(|e| {
+        PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
+            "Failed to read input directory: {}",
+            e
+        ))
+    })? {
+        let entry = entry.map_err(|e| {
+            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
+                "Error reading directory entry: {}",
+                e
+            ))
+        })?;
 
         let path = entry.path();
-        if path.extension().map_or(false, |ext| ext == "exe" || ext == "dll") {
+        if path
+            .extension()
+            .map_or(false, |ext| ext == "exe" || ext == "dll")
+        {
             // Generate variants with different configurations
             for variant in 0..num_variants {
                 let algo = match variant % 3 {
@@ -247,21 +258,26 @@ pub fn generate_training_dataset(
 
                 let packer = RustPEPacker::new(config);
 
-                let filename = path.file_stem()
+                let filename = path
+                    .file_stem()
                     .map(|s| s.to_string_lossy().to_string())
                     .unwrap_or_else(|| "sample".to_string());
 
-                let output_file = output_path.join(
-                    format!("{}_{}.exe", filename, variant)
-                );
+                let output_file = output_path.join(format!("{}_{}.exe", filename, variant));
 
                 if let Ok(data) = std::fs::read(&path) {
                     if let Ok(packed) = packer.pack_bytes(&data) {
                         if let Ok(mut file) = std::fs::File::create(&output_file) {
-                            let _ = file.write_all(&packed);
-                            generated_files.push(
-                                output_file.to_string_lossy().to_string()
-                            );
+                            if let Err(e) = file.write_all(&packed) {
+                                return Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
+                                    format!(
+                                        "Failed to write output file {}: {}",
+                                        output_file.display(),
+                                        e
+                                    ),
+                                ));
+                            }
+                            generated_files.push(output_file.to_string_lossy().to_string());
                         }
                     }
                 }
@@ -285,4 +301,4 @@ fn _native(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(xor_bytes, m)?)?;
 
     Ok(())
-}
+}
diff --git a/src/python/helpers.rs b/src/python/helpers.rs
diff --git a/src/python/mod.rs b/src/python/mod.rs
diff --git a/src/utils/error.rs b/src/utils/error.rs
diff --git a/src/utils/logger.rs b/src/utils/logger.rs
diff --git a/src/utils/mod.rs b/src/utils/mod.rs

Original file line number	Diff line number	Diff line change
`@@ -236,4 +236,4 @@ mod tests {`
`236`	`236`
`237`	`237`	`assert!(packer.config().key.is_some());`
`238`	`238`	`}`
`239`		`-}`
	`239`	`+}`