Simple rules for claripy AST simplification for human readability (total ordering, Reverse() simplification, ...)

simplify.py

import itertools
import claripy
import operator

top = 'T' # lambda a: True
ident = '=' # lambda a: a
inv = '!' # lambda a: not a
bot = '_' # lambda a: False

cmp_info = {
    # (A == B), (A ULT B), (A SLT B)
    '__eq__': (ident, bot, bot),
    '__ne__': (inv, bot, bot),
    #'__ge__': (ident, inv, bot),
    #'__lt__': (inv, ident, bot),
    #'__le__': (ident, ident, bot),
    #'__gt__': (inv, inv, bot),
    'UGE': (ident, inv, bot),
    'ULT': (inv, ident, bot),
    'ULE': (ident, ident, bot),
    'UGT': (inv, inv, bot),
    'SGE': (ident, bot, inv),
    'SLT': (inv, bot, ident),
    'SLE': (ident, bot, ident),
    'SGT': (inv, bot, inv),
}

ALL_BOT = (bot, bot, bot)

rev_cmp_info = dict((v, k) for k, v in cmp_info.items())

# assume both operator works on same value
or_map = {
    (top, top): top, (top, ident): top, (top, inv): top, (top, bot): top,
    (ident, top): top, (ident, ident): ident, (ident, inv): top, (ident, bot): ident,
    (inv, top): top, (inv, ident): top, (inv, inv): inv, (inv, bot): inv,
    (bot, top): top, (bot, ident): ident, (bot, inv): inv, (bot, bot): bot,
}

inv_map = {top: bot, ident: inv, inv: ident, bot: top}

def get_op_info(expr, inv=False):
    while expr.op == 'Not':
        expr, = expr.args
        inv = not inv
    k = cmp_info.get(expr.op, None)
    if k is not None:
        return expr, (inv_map[k[0]], inv_map[k[1]], inv_map[k[2]]) if inv else k
    else:
        return expr, None

def ast_equivalent(a, b):
    return claripy.is_true(a == b)

def simplify_or_cmppair(left, right, left_inv, right_inv, inv):
    left, left_op = get_op_info(left, left_inv)
    if left_op is None:
        #print("simplify_or_cmppair(%r, %r, %r, %r, %r): left simplification failed" % (left, right, left_inv, right_inv, inv))
        return None

    right, right_op = get_op_info(right, right_inv)
    if right_op is None:
        #print("simplify_or_cmppair(%r, %r, %r, %r, %r): right simplification failed" % (left, right, left_inv, right_inv, inv))
        return None

    (la, lb), (ra, rb) = left.args, right.args

    # normalize (A left_op B) or (B right_op A)
    if ast_equivalent(la, rb) and ast_equivalent(lb, ra):
        ra, rb = rb, ra
        right_op = (right_op[0], inv_map[right_op[1]], inv_map[right_op[2]])
    elif not (ast_equivalent(la, ra) and ast_equivalent(lb, rb)):
        #print("simplify_or_cmppair(%r, %r, %r, %r, %r): two args non-equiv" % (left, right, left_inv, right_inv, inv))
        return None

    # now we need to simplify the form of [inv xor {(A left_op B) or (A right_op B)}]
    #print("%s[%r <%r> %r] OR [%r <%r> %r]" % ("inv " if inv else "", la, left_op, lb, ra, right_op, rb))
    (leq, lul, lsl), (req, rul, rsl) = left_op, right_op
    new_op = (or_map[leq, req], or_map[lul, rul], or_map[lsl, rsl])
    if inv:
        new_op = (inv_map[new_op[0]], inv_map[new_op[1]], inv_map[new_op[2]])
    if top in new_op:
        return claripy.false if inv else claripy.true
    elif top == ALL_BOT:
        return claripy.true if inv else claripy.false
    else:
        final_opnam = rev_cmp_info.get(new_op, None)
        if final_opnam is not None:
            return getattr(claripy, final_opnam)(la, lb)

    return None

def simplify_logic(expr):
    if not isinstance(expr, claripy.ast.Bool):
        return expr
    inv = False
    while expr.op == 'Not':
        expr, = expr.args
        inv = not inv
    if expr.op == 'And':
        alist = list(expr.args)
        if len(alist) == 2:
            a, b = alist
            k = simplify_or_cmppair(a, b, True, True, not inv)
            if k is not None:
                return k
        while True:
            for (i, a), (j, b) in itertools.combinations(enumerate(alist), r=2):
                k = simplify_or_cmppair(a, b, True, True, True)
                if k is not None:
                    del alist[i], alist[j]
                    alist.append(k)
                    break
            else:
                break
        expr = claripy.And(*alist)
    elif expr.op == 'Or':
        alist = list(expr.args)
        if len(alist) == 2:
            a, b = alist
            k = simplify_or_cmppair(a, b, False, False, inv)
            if k is not None:
                return k
        while True:
            for (i, a), (j, b) in itertools.combinations(enumerate(alist), r=2):
                k = simplify_or_cmppair(a, b, False, False, False)
                if k is not None:
                    del alist[i], alist[j]
                    alist.append(k)
                    break
            else:
                break
        expr = claripy.Or(*alist)
    return claripy.Not(expr) if inv else expr

def simplify_if(expr):
    if expr.op != 'If':
        return expr

    cond, iftrue, iffalse = expr.args
    mod = True

    while mod:
        mod = False
        if iftrue.op == 'If':
            iftrue_cond, iftrue_iftrue, iftrue_iffalse = iftrue.args
            if ast_equivalent(iffalse, iftrue_iftrue): # If(x, If(y, A, B), A) => If(x & ~y, B, A)
                cond = claripy.And(cond, claripy.Not(iftrue_cond))
                iftrue = iftrue_iffalse
                mod = True
            elif ast_equivalent(iffalse, iftrue_iffalse): # If(x, If(y, A, B), B) => If(x & y, A, B)
                cond = claripy.And(cond, iftrue_cond)
                iftrue = iftrue_iftrue
                mod = True
        if iffalse.op == 'If':
            iffalse_cond, iffalse_iftrue, iffalse_iffalse = iffalse.args
            if ast_equivalent(iftrue, iffalse_iftrue): # If(x, A, If(y, A, B)) => If(x | y, A, B)
                cond = claripy.Or(cond, iffalse_cond)
                iffalse = iffalse_iffalse
                mod = True
            elif ast_equivalent(iftrue, iffalse_iftrue): # If(x, A, If(y, B, A)) => If(x | ~y, A, B)
                cond = claripy.And(cond, claripy.Not(iffalse_cond))
                iffalse = iffalse_iftrue
                mod = True

    return claripy.If(cond, iftrue, iffalse)

def _simplify(expr):
    expr = simplify_logic(expr)
    expr = simplify_if(expr)
    if expr.op in claripy.operations.leaf_operations:
        return expr
    else:
        a = _simplify(expr.args[0])
        bargs = map(_simplify, expr.args[1:])
        try:
            method = getattr(a, expr.op)
        except AttributeError:
            return getattr(claripy.ast.all_operations, expr.op)(a, *bargs)
        else:
            return method(*bargs)

def reverse_deep(expr):
    if not isinstance(expr, claripy.ast.Base):
        raise TypeError("expr must be an instance of claripy.ast.Base")
    if expr.op == 'If':
        cond, iftrue, iffalse = expr.args
        return claripy.If(cond, reverse_deep(iftrue), reverse_deep(iffalse))
    else:
        return claripy.Reverse(expr)

def simplify(expr):
    if not isinstance(expr, claripy.ast.Base):
        return expr
    if expr.op == 'Reverse':
        assert len(expr.args) == 1
        expr = reverse_deep(expr.args[0])
    return _simplify(claripy.simplify(expr))