obfuscate.py

"""
Lambda Calculus Obfuscator

CPS Protocol:
  Statement = λ(s, d) -> s(success_branch, fail_branch)
  Branch = λn -> n(should_continue, value)
  Chaining: stmt1(TRUE, None)(stmt2)(stmt3)...(final_continuation)
"""
import ast
from typing import List, Optional

import astor

# ============== AST Helpers ==============


def lam(params: List[str], body: ast.expr) -> ast.Lambda:
    """Create λparams. body"""
    return ast.Lambda(
        args=ast.arguments(args=[ast.arg(arg=p) for p in params]),
        body=body
    )


def call(func: ast.expr, *args: ast.expr) -> ast.Call:
    """Create func(args...)"""
    return ast.Call(func=func, args=list(args))


def name(id: str) -> ast.Name:
    """Create variable reference"""
    return ast.Name(id=id, ctx=ast.Load())


def const(value) -> ast.Constant:
    """Create constant"""
    return ast.Constant(value=value)


def assign_expr(var: str, val: ast.expr) -> ast.NamedExpr:
    """Create (var := val)"""
    return ast.NamedExpr(target=ast.Name(id=var, ctx=ast.Store()), value=val)


def ternary(test: ast.expr, if_true: ast.expr, if_false: ast.expr) -> ast.IfExp:
    """Create test ? if_true : if_false"""
    return ast.IfExp(test=test, body=if_true, orelse=if_false)

# ============== Church Encodings ==============


def TRUE():
    return lam(['x', 'y'], name('x'))    # λx,y. x


def FALSE():
    return lam(['x', 'y'], name('y'))   # λx,y. y


def FIX():
    return ast.parse(
        "lambda f: (lambda x: f(lambda *y: x(x)(*y))) (lambda x: f(lambda *y: x(x)(*y)))",
        mode='eval'
    ).body


def RET_NONE():
    # λs, d. cond(None, k)
    # instead of λs, d . λn. n(..)
    # returns λs, d. VALUE
    return ast.parse("lambda cond, k: cond(None, k)", mode='eval').body

# ============== CPS Builders ==============


def cps_branch(should_continue: ast.expr, value: ast.expr) -> ast.Lambda:
    """λn. n(should_continue, value)"""
    return lam(['n'], call(name('n'), should_continue, value))


def cps_stmt(value: ast.expr, continues: bool = True) -> ast.Lambda:
    """
    λ(s, d). s(λn. n(TRUE/FALSE, value), λn. n(FALSE, d))
    """
    cont = TRUE() if continues else FALSE()
    return lam(['s', 'd'], call(
        name('s'),
        cps_branch(cont, value),
        cps_branch(FALSE(), name('d'))
    ))

# ============== Statement Transformers ==============


class StmtToLambda(ast.NodeTransformer):
    """Transforms individual statements to λ(s, d) form."""

    def visit_Assign(self, node: ast.Assign) -> ast.Lambda:
        if len(node.targets) != 1 or not isinstance(node.targets[0], ast.Name):
            raise NotImplementedError
        return cps_stmt(assign_expr(node.targets[0].id, node.value), continues=True)

    def visit_Expr(self, node: ast.Expr) -> ast.Lambda:
        return cps_stmt(node.value, continues=True)

    def visit_Return(self, node: ast.Return) -> ast.Lambda:
        val = node.value if node.value is not None else const(None)
        return cps_stmt(val, continues=False)

# ============== Main Transformer ==============


class Obfuscator(ast.NodeTransformer):
    def __init__(self):
        super().__init__()
        self._lambdas = {}

    def stmts_to_lambda(self, stmts: List[ast.stmt], initial_d: Optional[ast.expr] = None) -> ast.expr:
        """
        Transform statements to chained expression.

        Empty: returns λn. n(TRUE, None) (a branch)
        Non-empty: returns stmt1(TRUE, initial_d)(stmt2)... (a call chain)
        """
        if initial_d is None:
            initial_d = const(None)

        if not stmts:
            return cps_branch(TRUE(), initial_d)

        exprs = [StmtToLambda().visit(stmt) for stmt in stmts]
        res = call(exprs[0], TRUE(), initial_d)
        for expr in exprs[1:]:
            res = call(res, expr)
        return res

    def visit_FunctionDef(self, node: ast.FunctionDef):
        self.generic_visit(node)

        func = lam(
            [arg.arg for arg in node.args.args],
            call(self.stmts_to_lambda(node.body), RET_NONE())
        )

        wrapped = lam([node.name], func)
        self._lambdas[node.name] = call(FIX(), wrapped)
        return None

    def visit_Assign(self, node: ast.Assign):
        self.generic_visit(node)

        if isinstance(node.value, ast.Lambda) and isinstance(node.targets[0], ast.Name):
            func_name = node.targets[0].id
            wrapped = lam([func_name], node.value)
            self._lambdas[func_name] = call(FIX(), wrapped)
            return None
        return node

    def visit_Call(self, node: ast.Call):
        self.generic_visit(node)

        if isinstance(node.func, ast.Name) and node.func.id in self._lambdas:
            return call(self._lambdas[node.func.id], *node.args)
        return node

    def visit_If(self, node: ast.If):
        self.generic_visit(node)

        condition = ternary(node.test, TRUE(), FALSE())
        body = self.stmts_to_lambda(node.body)
        orelse = self.stmts_to_lambda(node.orelse)

        # λ(s, d). condition(body, orelse)
        return lam(['s', 'd'], call(
            name('s'),
            call(condition, body, orelse),
            cps_branch(FALSE(), name('d'))
        ))
    
    def visit_For(self, node: ast.For):
        raise NotImplementedError("For loops are not yet supported.")

    def visit_While(self, node: ast.While):
        self.generic_visit(node)

        condition = ternary(node.test, TRUE(), FALSE())
        body = self.stmts_to_lambda(node.body)

        # TODO: Support while-else semanatics

        # loop_impl builds a recursive statement:
        # loop(s, d) = s( condition( body(n_loop), continue_branch ),
        #                fail_branch )
        # where n_loop = λ(s, d). loop(s, d)
        loop_impl = lam(['loop'], lam(['s', 'd'], call(
            name('s'),
            call(call(
                condition,
                # if condition is true: run body, then recurse via n_loop
                lam([], call(body, lam(['s', 'd'], call(name('loop'), name('s'), name('d'))))),
                # if condition is false: exit loop and continue with current d
                lam([], cps_branch(TRUE(), name('d')))
            )),
            # fail branch: propagate early return
            cps_branch(FALSE(), name('d'))
        )))

        # Return the fixed-point loop as the while statement
        return call(FIX(), loop_impl)