Merge pull request #2 from lucasscharenbroch/addStringPercentage

Add string percentage

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);

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) };

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");
+}