Merge pull request #42 from bliutech/add-var-counter

bliutech · Jul 20, 2024 · 57669e5 · 57669e5
2 parents 5e17b4c + 094e5c5
commit 57669e5
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 11 deletions.
diff --git a/experiments/experiment.py b/experiments/experiment.py
@@ -9,7 +9,7 @@
 
 from solver import Solver
 
-from utils.metrics import OpCounter
+from utils.metrics import OpCounter, VarCounter
 
 
 def count_ops(ast: Expr) -> int:
@@ -18,6 +18,12 @@ def count_ops(ast: Expr) -> int:
     return oc.getCount()
 
 
+def count_vars(ast: Expr) -> int:
+    vc: VarCounter = VarCounter()
+    ast.accept(vc)
+    return vc.getCount()
+
+
 def run_experiment(
     passes: list[str], n: int = 100
 ) -> list[tuple[int, str, int, str, int, str, int]]:
@@ -54,28 +60,37 @@ def run_experiment(
         orig_ast: Expr = p.parse(l.tokens)
         print(f"Original expression: {orig_ast}", file=sys.stderr)
 
-        orig_count = count_ops(orig_ast)
-        print(f"Original count: {orig_count}", file=sys.stderr)
+        orig_op_count = count_ops(orig_ast)
+        print(f"Original Operation count: {orig_op_count}", file=sys.stderr)
+        orig_var_count = count_vars(orig_ast)
+        print(f"Original Variable count: {orig_var_count}", file=sys.stderr)
 
         obf_ast = MBAObfuscator().obfuscate(orig_ast)
         print(f"Obfuscated expression: {obf_ast}", file=sys.stderr)
-        obf_count = count_ops(obf_ast)
-        print(f"Obfuscated count: {obf_count}", file=sys.stderr)
+        obf_op_count = count_ops(obf_ast)
+        print(f"Obfuscated Operation count: {obf_op_count}", file=sys.stderr)
+        obf_var_count = count_vars(obf_ast)
+        print(f"Obfuscated Variable count: {obf_var_count}", file=sys.stderr)
 
         simplified_ast = s.run(obf_ast)
         print(f"Simplified expression: {simplified_ast}", file=sys.stderr)
-        simplified_count = count_ops(simplified_ast)
-        print(f"Simplified count: {simplified_count}", file=sys.stderr)
+        simplified_op_count = count_ops(simplified_ast)
+        print(f"Simplified Operation count: {simplified_op_count}", file=sys.stderr)
+        simplified_var_count = count_vars(simplified_ast)
+        print(f"Simplified Variable count: {simplified_var_count}", file=sys.stderr)
 
         res.append(
             (
                 i,
                 orig_ast,
-                orig_count,
+                orig_op_count,
+                orig_var_count,
                 obf_ast,
-                obf_count,
+                obf_op_count,
+                obf_var_count,
                 simplified_ast,
-                simplified_count,
+                simplified_op_count,
+                simplified_var_count,
             )
         )
 

diff --git a/utils/metrics.py b/utils/metrics.py
@@ -1,6 +1,6 @@
 from typing import override
 
-from parser.ast import OrExpr, AndExpr, XorExpr, NotTerm
+from parser.ast import OrExpr, AndExpr, XorExpr, NotTerm, VarVar
 from parser.visitor import Visitor
 
 
@@ -38,3 +38,17 @@ def visitXorExpr(self, node: XorExpr) -> None:
     def visitNotTerm(self, node: NotTerm) -> None:
         self._count += 1
         node.first.accept(self)
+
+
+class VarCounter(Visitor):
+    """Counts the number of boolean operators visited"""
+
+    def __init__(self) -> None:
+        self._count: int = 0
+
+    def getCount(self) -> int:
+        return self._count
+
+    @override
+    def visitVarVar(self, node: "VarVar") -> None:
+        self._count += 1