feat(experiments): implemented experiment framework

Translated LLVM Obfuscator's instruction substitution to an obfuscation pass and updated solver framework to allow easy running of experiments on solver passes.

experiments/experiment.py

@@ -12,6 +12,12 @@
 from utils.metrics import OpCounter
 
 
+def count_ops(ast: Expr) -> int:
+    oc: OpCounter = OpCounter()
+    ast.accept(oc)
+    return oc.getCount()
+
+
 def run_experiment(
     passes: list[str], n: int = 100
 ) -> list[tuple[int, str, int, str, int, str, int]]:
@@ -34,31 +40,31 @@ def run_experiment(
 
     bg: BooleanGenerator = BooleanGenerator()
     s: Solver = Solver(passes)
-    oc: OpCounter = OpCounter()
 
     for i in range(2, n + 1):
+        print("========================================", file=sys.stderr)
         print(f"Running experiment with n={i}", file=sys.stderr)
 
         expr: str = bg.generate_expr(i)
-        print(f"Original expression: {expr}", file=sys.stderr)
 
-        l: Lexer = Lexer(expr)
-        l.lex()
+        l: Lexer = Lexer()
+        l.lex(expr)
 
-        p: Parser = Parser(l.tokens)
-        orig_ast: Expr = p.parse()
+        p: Parser = Parser()
+        orig_ast: Expr = p.parse(l.tokens)
+        print(f"Original expression: {orig_ast}", file=sys.stderr)
 
-        orig_count = orig_ast.acceptRet(oc)
+        orig_count = count_ops(orig_ast)
         print(f"Original count: {orig_count}", file=sys.stderr)
 
         obf_ast = MBAObfuscator().obfuscate(orig_ast)
         print(f"Obfuscated expression: {obf_ast}", file=sys.stderr)
-        obf_count = obf_ast.acceptRet(oc)
+        obf_count = count_ops(obf_ast)
         print(f"Obfuscated count: {obf_count}", file=sys.stderr)
 
         simplified_ast = s.run(obf_ast)
         print(f"Simplified expression: {simplified_ast}", file=sys.stderr)
-        simplified_count = simplified_ast.acceptRet(oc)
+        simplified_count = count_ops(simplified_ast)
         print(f"Simplified count: {simplified_count}", file=sys.stderr)
 
         res.append(

experiments/obfuscator.py

@@ -7,11 +7,27 @@
 at https://github.com/obfuscator-llvm/obfuscator/wiki/Instructions-Substitution
 """
 
+from typing import override
+
 from parser.visitor import RetVisitor
 
+from parser.ast import (
+    Node,
+    Expr,
+    TermExpr,
+    OrExpr,
+    AndExpr,
+    XorExpr,
+    ParenTerm,
+    NotTerm,
+    VarVar,
+    TrueConst,
+    FalseConst,
+)
+
 
-class MBAObfuscator(RetVisitor):
-    def obfuscate(self, ast, n: int = 1):
+class MBAObfuscator(RetVisitor[Node]):
+    def obfuscate(self, ast: Expr, n: int = 1):
         """
         Obfuscates the given AST.
 
@@ -23,4 +39,64 @@ def obfuscate(self, ast, n: int = 1):
 
         return ast
 
-    # TODO: Override the visit methods to apply the obfuscation rules
+    @override
+    def visitTermExpr(self, node: TermExpr) -> Node:
+        first: Node = node.first.acceptRet(self)
+        if node.second:
+            second: Node = node.second.first.acceptRet(self)
+            if isinstance(node.second, OrExpr):
+                # a = b | c => a = (b & c) | (b ^ c)
+                return TermExpr(
+                    ParenTerm(TermExpr(first, AndExpr(second))),
+                    OrExpr(ParenTerm(TermExpr(first, XorExpr(second)))),
+                )
+            elif isinstance(node.second, AndExpr):
+                # a = b & c => a = (b ^ ~c) & b
+                return TermExpr(
+                    ParenTerm(TermExpr(first, XorExpr(NotTerm(second)))), AndExpr(first)
+                )
+            elif isinstance(node.second, XorExpr):
+                # a = a ^ b => a = (~a & b) | (a & ~b)
+                return TermExpr(
+                    ParenTerm(TermExpr(NotTerm(first), AndExpr(second))),
+                    OrExpr(ParenTerm(TermExpr(first, AndExpr(NotTerm(second))))),
+                )
+        return TermExpr(first)
+
+    @override
+    def visitOrExpr(self, node: OrExpr) -> Node:
+        first: Node = node.first.acceptRet(self)
+        second: Node = node.second.acceptRet(self)
+        return OrExpr(first, second)
+
+    @override
+    def visitAndExpr(self, node: AndExpr) -> Node:
+        first: Node = node.first.acceptRet(self)
+        second: Node = node.second.acceptRet(self)
+        return AndExpr(first, second)
+
+    @override
+    def visitXorExpr(self, node: XorExpr) -> Node:
+        first: Node = node.first.acceptRet(self)
+        second: Node = node.second.acceptRet(self)
+        return XorExpr(first, second)
+
+    @override
+    def visitParenTerm(self, node: ParenTerm) -> Node:
+        return ParenTerm(node.first.acceptRet(self))
+
+    @override
+    def visitNotTerm(self, node: NotTerm) -> Node:
+        return NotTerm(node.first.acceptRet(self))
+
+    @override
+    def visitVarVar(self, node: VarVar) -> Node:
+        return VarVar(node.name)
+
+    @override
+    def visitTrueConst(self, node: TrueConst) -> Node:
+        return TrueConst()
+
+    @override
+    def visitFalseConst(self, node: FalseConst) -> Node:
+        return FalseConst()

run_experiment.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python
+
+from experiments.experiment import run_experiment
+
+if __name__ == "__main__":
+    passes: list[str] = ["sympy_pass"]
+
+    res = run_experiment(passes, 10)
+
+    print("========================================")
+    print(res)

solver/__init__.py

@@ -9,14 +9,20 @@
 import sys
 from importlib import import_module
 
+from parser.ast import Expr
+
 
 class Solver:
     def __init__(self, passes: list[str]):
         self.passes: list[str] = passes
 
+    def get_module(self, name: str):
+        print(f"Importing {name}")
+        return import_module(f"solver.passes.{name}")
+
     # TODO: Add type annotation for AST once classes are finished.
-    def run(self, ast):
-        for m in map(import_module, f"passes.{self.passes}"):
+    def run(self, ast: Expr) -> Expr:
+        for m in map(self.get_module, self.passes):
             try:
                 ast = m.run_pass(ast)
             except Exception as e:

solver/passes/sympy_pass.py

@@ -30,7 +30,7 @@ def run_pass(ast: Expr) -> Expr:
     p: sympy.Basic = ast.acceptRet(tv)
 
     simplifiedExpr: sympy.Basic = sympy.simplify_logic(p)
-    simplifiedStr: str = str(simplifiedExpr)
+    simplifiedStr: str = str(simplifiedExpr).replace("~", "!")
 
     l: Lexer = Lexer()
     l.lex(simplifiedStr)