r2_backward_slicing.py

"""
Backward slice from a given sink.

Unfortunately, the BB definition used by Radare2 is too limiting:
    https://github.com/radareorg/radare2/issues/7170#issuecomment-631579110
"""
from functools import reduce

import r2pipe
import sys

#
# Utils
#

def _predecessors_from_codexrefs(codexrefs):
    def _block_top_address(address):
        result = r2.cmdj("abj %s" % address)
        assert len(result) == 1, "A `call` instruction should belong to a single block"
        return result[0]['addr']
    # TODO:
    # if len(codexrefs) != 0:
    #     address = codexrefs[0]['addr']
    #     if _block_top_address(address) == 0x401140:
    #         calls = list(filter(
    #             # Let's get the `call`s instructions in the node of the callsite, located *before* the current callsite.
    #             lambda x: x['type'] == 'CALL' and x['at'] < address,
    #             r2.cmdj("axffj @%s" % address)
    #         ))
    #         sorted_calls = sorted(calls, key=lambda e: e['at'], reverse=True)
    #         import ipdb; ipdb.set_trace()
    #         pass

    return list(map(
        lambda p: _block_top_address(p['addr']),
        codexrefs
    ))

def _codexrefs_from_function_address(address):
    result = r2.cmdj("afij %s" % address)
    assert len(result) == 1, "Function <%s> should have one set of info" % address
    return result[0]["codexrefs"]

PREDECESSORS_PER_ADDRESS = dict()
def _predecessors_from_address(address):
    if address in PREDECESSORS_PER_ADDRESS.keys():
        return PREDECESSORS_PER_ADDRESS[address]

    codexrefs = _codexrefs_from_function_address(address)
    return _predecessors_from_codexrefs(codexrefs)

def all_predecessors_from_address(address, callback):
    """
    :param int address:
    :param Function callback: A function to call for on each predecessor.
    """
    predecessors = _predecessors_from_address(address)

    if len(predecessors) == 0:
        return []

    for p in predecessors:
        _add_to_graph(address, p)

    return list(reduce(
        lambda acc, p: acc + all_predecessors_from_address(p, callback),
        predecessors,
        [address]
    ))

#
# Main
#

try:
    binary = sys.argv[1]
    sink_name = sys.argv[2]

    r2 = r2pipe.open(binary)
    r2.cmd('aaa')

    sink_address = int(r2.cmd("f~%s[0]" % sink_name)[:-1], 16)
except IndexError:
    exit("usage: %s <binary> <function_address>" % sys.argv[0])
except BrokenPipeError:
    exit("%s failed to open. Is this really an existing binary?" % binary)
except ValueError:
    exit("There is no function '%s' in the binary." % sys.argv[2])

def _add_to_graph(a, p):
    (a, p) = (hex(a), hex(p))
    r2.cmd("agn %s" % p)
    r2.cmd("age %s %s" % (p, a))

r2.cmd("agn %s" % hex(sink_address))
all_predecessors_from_address(sink_address, _add_to_graph)

result = r2.cmd('agg')
print(result)