Merge pull request #24 from PoCInnovation/syntax_check

added real slightly better syntax checks and better error display

src/main.rs

@@ -1,16 +1,21 @@
 mod obfuscator;
+use std::process::ExitCode;
+
 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, matches: ArgMatches) -> obfuscator::Result<()> {
+fn run_obfuscator(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 !obfuscator.is_syntax_ok()? {
+        Err(obfuscator::error::ObfuscatorError::InvalidCode)?;
+    }
     if run_all || set_options.contains(&&"dead".to_string()) {
         obfuscator.insert_dead_branches()?;
     }
@@ -34,7 +39,24 @@ fn run_obfuscation(mut obfuscator: Obfuscator, matches: ArgMatches) -> obfuscato
     Ok(())
 }
 
-fn main() -> obfuscator::Result<()> {
+fn run_obfuscation(code: String, matches: ArgMatches) -> ExitCode {
+    let obfuscator = match Obfuscator::new(code) {
+        Ok(ob) => ob,
+        Err(err) => {
+            println!("{err}");
+            return ExitCode::SUCCESS;
+        }
+    };
+
+    if let Err(err) = run_obfuscator(obfuscator, matches) {
+        println!("{err}");
+        ExitCode::FAILURE
+    } else {
+        ExitCode::SUCCESS
+    }
+}
+
+fn main() -> ExitCode {
     let matches = Command::new("Obfuscator")
         .version("0.0")
         .about("static obfuscation of python code")
@@ -58,8 +80,6 @@ fn main() -> obfuscator::Result<()> {
         .get_matches();
 
     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,
-    )
+    let code = std::fs::read_to_string(file_name).expect("The Path given was wrong");
+    run_obfuscation(code, matches)
 }

src/obfuscator/dead_code_entry_points.rs

@@ -36,48 +36,58 @@ fn figure_out_indentation(line: &str) -> usize {
     indent
 }
 
+fn insert_dead_branches(code: &str, attempt: usize) -> String {
+    let lines = code.lines().count();
+    let mut rng = rand::thread_rng();
+    let iterations = rng.gen_range(1..(lines / 3) - attempt);
+    let mut new_code: String = code.to_string();
+
+    for _ in 0..iterations {
+        let line = rng.gen_range(init::OBFUSCATOR_HELPER_FUNCTIONS.lines().count()..lines);
+
+        new_code = new_code.lines().enumerate().map(|(i, l)| {
+            if i == line {
+                let indent = figure_out_indentation(l);
+                let mut line = DEAD_CODE_ENTRY_POINT[rng.gen_range(0..DEAD_CODE_ENTRY_POINT.len())]
+                    .to_string();
+                line.insert_str(0, " ".repeat(indent).as_str());
+                line.push('\n');
+                line.push_str(" ".repeat(indent + 4).as_str());
+                line.push_str(RANDOM_USELESS_CODE[rng.gen_range(0..RANDOM_USELESS_CODE.len())]);
+                line.push('\n');
+                line.push_str(l);
+                line.push('\n');
+                line
+            } else {
+                let mut l = l.to_string();
+                l.push('\n');
+                l
+            }
+        })
+        .collect::<String>();
+    }
+    new_code
+}
+
 impl Obfuscator {
     pub fn insert_dead_branches(&mut self) -> Result<()> {
-        let lines = self.code.lines().count();
-        let mut rng = rand::thread_rng();
-        let iterations = rng.gen_range(1..lines / 3);
-
-        for _ in 0..iterations {
-            let line = rng.gen_range(init::OBFUSCATOR_HELPER_FUNCTIONS.lines().count()..lines);
-            let code = self.code.clone();
+        let attempts :usize = 5;
+        let old = self.code.clone();
 
-            self.code = code
-                .lines()
-                .enumerate()
-                .map(|(i, l)| {
-                    if i == line {
-                        let indent = figure_out_indentation(l);
-                        let mut line = DEAD_CODE_ENTRY_POINT
-                            [rng.gen_range(0..DEAD_CODE_ENTRY_POINT.len())]
-                        .to_string();
-                        line.insert_str(0, " ".repeat(indent).as_str());
-                        line.push('\n');
-                        line.push_str(" ".repeat(indent + 4).as_str());
-                        line.push_str(
-                            RANDOM_USELESS_CODE[rng.gen_range(0..RANDOM_USELESS_CODE.len())],
-                        );
-                        line.push('\n');
-                        line.push_str(l);
-                        line.push('\n');
-                        line
-                    } else {
-                        let mut l = l.to_string();
-                        l.push('\n');
-                        l
-                    }
-                })
-                .collect::<String>();
-            eprintln!("self.code: {}", self.code);
+        self.code = insert_dead_branches(&self.code, 0);
+        for i in 1..attempts {
+            if self.is_syntax_ok()? {
+                self.reparse(ObfuscatorError::DeadCode)?;
+                return Ok(());
+            } else {
+                self.code = insert_dead_branches(&old, i);
+            }
         }
-        if self.reparse(ObfuscatorError::DeadCode).is_err() {
-            self.insert_dead_branches() // might hang a better solution must be found
-        } else {
+        if self.is_syntax_ok()? {
+            self.reparse(ObfuscatorError::DeadCode)?;
             Ok(())
+        } else {
+            Err(ObfuscatorError::DeadCode)
         }
     }
 }

src/obfuscator/error.rs

@@ -5,6 +5,7 @@ pub enum ObfuscatorError {
     Booleans,
     Strings,
     Numbers,
+    PythonSyntaxCheck(std::io::Error),
     Functions(String),
 }
 
@@ -21,6 +22,9 @@ impl std::fmt::Display for ObfuscatorError {
                 f,
                 "Obfuscator failed while obfuscation functions at function {s}"
             ),
+            Self::PythonSyntaxCheck(err) => {
+                write!(f, "The Python command for syntax failed : {err}")
+            }
             _ => write!(
                 f,
                 "Obfuscator encountered an error in {:?} obfuctation",
@@ -30,4 +34,11 @@ impl std::fmt::Display for ObfuscatorError {
     }
 }
 
-impl std::error::Error for ObfuscatorError {}
+impl std::error::Error for ObfuscatorError {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        match self {
+            Self::PythonSyntaxCheck(err) => Some(err),
+            _ => None
+        }
+    }
+}

src/obfuscator/mod.rs

@@ -1,12 +1,13 @@
 mod boolean;
 mod dead_code_entry_points;
-mod error;
+pub mod error;
 mod functions;
 mod init;
 mod intergers;
 mod print;
 mod random_identifiers;
 mod strings;
+mod syntax;
 pub use error::Result;
 
 use tree_sitter::{Parser, Tree};

src/obfuscator/syntax.rs

@@ -0,0 +1,28 @@
+use super::Obfuscator;
+use super::Result;
+use std::process::Command;
+
+impl Obfuscator {
+    pub fn is_syntax_ok(&self) -> Result<bool> {
+        let cmd_status = Command::new("python3")
+            .stdout(std::process::Stdio::null())
+            .stderr(std::process::Stdio::null())
+            .arg("-Bc")
+            .arg(format!(
+                "compile(r\'\'\'{}\'\'\',\"\", \"exec\")",
+                self.code
+            ))
+            .status();
+
+        let cmd_res = match cmd_status {
+            Ok(val) => val,
+            Err(err) => return Err(super::error::ObfuscatorError::PythonSyntaxCheck(err)),
+        };
+
+        if cmd_res.success() {
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
+}