merge arguements into main

add arguments handling with clap

Cargo.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+clap = "4.5.17"
 rand = { version = "0.8.5", default-features = false, features = ["rand_chacha", "small_rng", "std", "std_rng"] }
 tree-sitter = "0.22"
 tree-sitter-python = "0.21"

src/main.rs

@@ -1,28 +1,65 @@
 mod obfuscator;
-use std::env;
+use clap::{arg, value_parser, ArgAction, ArgMatches, Command, ValueHint};
 
 use obfuscator::Obfuscator;
+const AVAILABLE_OBFUSCATION_METHODS: [&str; 5] = ["int", "string", "fn", "bools", "dead"];
 
-fn run_obfuscation(mut obfuscator: Obfuscator) -> obfuscator::Result<()> {
-    obfuscator.insert_dead_branches()?;
-    obfuscator.obfuscate_booleans()?;
-    obfuscator.obfuscate_strings()?;
-    obfuscator.obfuctate_functions()?;
-    obfuscator.obfuscate_integers()?;
-    obfuscator.print_tree();
+fn run_obfuscation(mut obfuscator: Obfuscator, matches: ArgMatches) -> obfuscator::Result<()> {
+    let run_all = !matches.contains_id("set");
+    let set_options = matches
+        .get_many::<String>("set")
+        .map(|values| values.collect::<Vec<_>>())
+        .unwrap_or_default(); // Default to empty if --set not provided
+
+    if run_all || set_options.contains(&&"dead".to_string()) {
+        obfuscator.insert_dead_branches()?;
+    }
+    if run_all || set_options.contains(&&"bools".to_string()) {
+        obfuscator.obfuscate_booleans()?;
+    }
+    if run_all || set_options.contains(&&"string".to_string()) {
+        obfuscator.obfuscate_strings()?;
+    }
+    if run_all || set_options.contains(&&"fn".to_string()) {
+        obfuscator.obfuscate_functions()?;
+    }
+    if run_all || set_options.contains(&&"int".to_string()) {
+        obfuscator.obfuscate_integers()?;
+    }
+
+    if matches.get_flag("debug") {
+        obfuscator.print_tree();
+    }
     println!("{}", obfuscator);
     Ok(())
 }
 
 fn main() -> obfuscator::Result<()> {
-    let args = env::args().collect::<Vec<String>>();
+    let matches = Command::new("Obfuscator")
+        .version("0.0")
+        .about("static obfuscation of python code")
+        .arg(
+            arg!(<script> "Script to obfuscate")
+                .required(true)
+                .value_parser(value_parser!(std::path::PathBuf))
+                .value_hint(ValueHint::FilePath),
+        )
+        .arg(
+            arg!(-d --debug "dumps more debug stuff, for now simply prints the syntax tree")
+                .required(false)
+                .action(ArgAction::SetTrue),
+        )
+        .arg(
+            arg!(-s --set [list] "comma separated list of options")
+                .required(false)
+                .num_args(0..=AVAILABLE_OBFUSCATION_METHODS.len())
+                .value_parser(AVAILABLE_OBFUSCATION_METHODS),
+        )
+        .get_matches();
 
-    if args.len() < 2 {
-        eprintln!("Usage: {} <file>", args[0]);
-        std::process::exit(1);
-    } else {
-        run_obfuscation(Obfuscator::new(
-            std::fs::read_to_string(&args[1]).expect("error reading file"),
-        )?)
-    }
+    let file_name = matches.get_one::<std::path::PathBuf>("script").unwrap();
+    run_obfuscation(
+        Obfuscator::new(std::fs::read_to_string(file_name).expect("Must be a real file"))?,
+        matches,
+    )
 }

src/obfuscator/functions.rs

@@ -65,7 +65,7 @@ fn replace_fn(tree: &mut Tree, code: &str, replace: &str, replacement: &str) ->
 }
 
 impl Obfuscator {
-    pub fn obfuctate_functions(&mut self) -> Result<()> {
+    pub fn obfuscate_functions(&mut self) -> Result<()> {
         let _ = self.reparse(ObfuscatorError::Booleans);
         let fns = get_fn(&self.tree, &self.code);
         for e in &fns {