diff --git a/auto_obfuscate/src/main.rs b/auto_obfuscate/src/main.rs index 05e6184..8e3caef 100644 --- a/auto_obfuscate/src/main.rs +++ b/auto_obfuscate/src/main.rs @@ -28,7 +28,11 @@ fn main() { .help("disable macro and modify source directly for flow obfuscation") ) .arg(Arg::with_name("var").long("var").help("Enable variable renaming")) - .arg(Arg::with_name("p").long("p").help("set upper bound for string literal encryption")) + .arg(Arg::with_name("p") + .short('p') + .long("percent_strings_to_encrypt") + .help("set upper bound for string literal encryption") + .value_name("PERCENTAGE")) .get_matches(); let path = matches.value_of("path").unwrap(); @@ -53,7 +57,14 @@ fn main() { } //set upper bound for string literal encryption if let Some(percentage) = matches.value_of("p") { - config.string_config.percentage = percentage.parse().unwrap_or(100); + config.string_config.percentage = match percentage.parse() { + Ok(n) if n <= 100 => n, + _ => { + eprintln!("-p: expected integer between 0 and 100, got: `{}`", percentage); + eprintln!("defaulting to 100%"); + 100 + } + }; } process_path(&path, &config); diff --git a/auto_obfuscate/src/string.rs b/auto_obfuscate/src/string.rs index f81248d..7850a50 100644 --- a/auto_obfuscate/src/string.rs +++ b/auto_obfuscate/src/string.rs @@ -37,6 +37,7 @@ pub struct StringObfuscator { percentage: u8, encrypted_count: usize, strings_to_encrypt: usize, + num_strings_encrypted: usize, } impl StringObfuscator { @@ -46,6 +47,7 @@ impl StringObfuscator { percentage: config.percentage, encrypted_count: 0, strings_to_encrypt: 0, + num_strings_encrypted: 0, } } #[allow(dead_code)] @@ -104,6 +106,12 @@ impl VisitMut for StringObfuscator { //replace all string literals with call to obfuscation macro fn visit_local_mut(&mut self, local: &mut Local) { if let Some(local_init) = &mut local.init { + if self.num_strings_encrypted >= self.strings_to_encrypt { + return; + } + self.num_strings_encrypted += 1; + + //match on local variables that contain string literal assignments if let Expr::Lit(ExprLit { lit: Lit::Str(lit_str), .. }) = &*local_init.expr { let encrypted = quote! { cryptify::encrypt_string!(#lit_str) }; diff --git a/auto_obfuscate/src/string/string_tests.rs b/auto_obfuscate/src/string/string_tests.rs index b4c178b..d6cf9cb 100644 --- a/auto_obfuscate/src/string/string_tests.rs +++ b/auto_obfuscate/src/string/string_tests.rs @@ -34,3 +34,40 @@ fn test_no_macro() { let parse_result = syn::parse_file(&obfuscated_code); assert!(parse_result.is_ok(), "Modified code is not valid Rust code"); } + +#[test] +fn test_percentage() { + let code = + r#" + fn main() { + let a = "a"; + let b = "b"; + let c = "c"; + let d = "d"; + let e = "e"; + let f = "f"; + let g = "g"; + let h = "h"; + let i = "i"; + let j = "j"; + println!("Hello"); + println!("Hello"); + println!("Hello"); + println!("Hello"); + println!("Hello"); + println!("Hello"); + } +"#; + + let mut string_config = StringConfig::default(); + string_config.percentage = 80; + let mut string_obfuscator = StringObfuscator::new(string_config); + let obfuscated_code = string_obfuscator.obfuscate_strings(code); + assert_ne!(code, obfuscated_code); + assert!(obfuscated_code.contains("encrypt_string ! (\"h\")")); + assert!(obfuscated_code.contains("let i = \"i\"")); + println!("{}", obfuscated_code); + + let parse_result = syn::parse_file(&obfuscated_code); + assert!(parse_result.is_ok(), "Modified code is not valid Rust code"); +}