nops_finder.py

import angr
import argparse
import numpy as np
from pwnlib.context import context
from pwnlib.elf import ELF

from extract import GadgetExtractor


class Encoder:

    _commandUninteresting = ['cmp', 'jne', 'ret', 'je', 'test', 'jmp', 'dword', 'call', 'qword', 'hlt', 'nop', 'jbe', 'jb']

    def _encodeOne(self, dstOP):
        if dstOP.startswith('ptr'):
            return 'mem[' + self._encodeOne(dstOP[5:-1]) + ']'
        args = dstOP.split(' ')
        ret = []
        for idx in range(len(args)):
            if idx % 2 == 0:
                if args[idx].isalpha():
                    ret.append('regs.' + args[idx])
                    continue
            ret.append(args[idx])
        return ' '.join(ret)


    def encode(self, gadget):
        instrs = [ instr.strip() for instr in gadget.split(';') ]
        constraints = []

        for instr in instrs:
            opcode, operands = instr.split(' ')[0], ' '.join(instr.split(' ')[1:])
            if operands.startswith('byte'):
                operands = operands[5:]
            if operands.startswith('dword') or operands.startswith('qword'):
                operands = operands[6:]

            dstOP = operands.split(',')[0]
            
            if opcode in self._commandUninteresting:
                continue

            constraints.append(self._encodeOne(dstOP))
        return constraints

    @staticmethod
    def decode(s, constraint, setVal = None):
        if constraint.endswith(']'):
            split_constraint = constraint.split('[')
            if setVal != None:
                getattr(s, split_constraint[0])[Encoder.decode(s, split_constraint[1][:-1])].uint64_t = setVal
                return None
            else:
                return getattr(s, split_constraint[0])[Encoder.decode(s, split_constraint[1][:-1])].uint64_t

        args = constraint.split(' ')
        operands = []
        operations = []

        if setVal != None:
            reg = args[0].split('.')[-1]
            setattr(getattr(s, 'regs'), reg, setVal)
            return None

        for arg in args:
            if arg.startswith('regs'):
                reg = arg.split('.')[-1].split('*')
                operands.append(getattr(getattr(s, 'regs'), reg[0]))
                if len(reg) == 2:
                    operands[-1] *= int(reg[1])
                continue

            if arg == '+' or arg == '-':
                operations.append(arg)
                continue

            operands.append(arg)

        result = operands[0]
        '''
        TODO1: Fix decoding memory addresses
        TODO2: Fix discrepancies between absolutes addresses in Ropgadget's output and the binary
               generated by angr
        '''
        #for op in range(len(operations)):
        #    if operations[op] == '+':
        #        result += operands[op + 1]
        #    else:
        #        result -= operands[op + 1]

        return result

            
class Predicate:
    def __init__(self, constraints):
        self.constraints = constraints

    def __call__(self):

        def is_true(s):
            for constraint in self.constraints:
                s.add_constraints(Encoder.decode(s, constraint[0]) == constraint[1])
            if s.satisfiable():
                return True
            for i in range(len(self.constraints)):
                try:
                    s.solver.constraints.pop()
                except:
                    pass
            s.solver.reload_solver()
            return False
        return is_true


class CFI:

    def __init__(self, srcs, dsts, Oracle):
        self.map = {}
        for s in srcs:
            allowed_targets = []
            for d in dsts:
                if Oracle(s, d):
                    allowed_targets.append(d)
            self.map[s] = allowed_targets

    def __call__(self):

        def CFI_mapper(state):
            dsts = self.map[state.solver.eval(state.regs.rip)]
            successors = []

            for d in dsts:
                s = state.copy()
                s.regs.rip = d
                s.history.jumpkind = 'Ijk_Boring'
                successors.append(s)

            return successors

        return CFI_mapper


class NOPsFinder:
    def __init__(self, path, extractor):
        self.extractor = extractor
        self._loadGadgets(path)
        self.proj = angr.project.load_shellcode(
            self.concatenated_gadgets, arch='amd64', support_selfmodifying_code=True)

        cfi = CFI(self.srcs, self.dsts, self.Oracle())()

        for s in self.srcs:
            self.proj.hook(s, cfi)

    def _loadGadgets(self, path):
        context.clear(arch='amd64')

        with open(path, 'r') as f:
            gadgets = [line.strip().split(':') for line in f.readlines()]
            self.gadgetOrigAdrrs = [ line[0][:-1] for line in gadgets ]
            gadgets = [ ':'.join(line[1:]).strip() for line in gadgets ] 

        self.concatenated_gadgets = ' ; '.join(gadgets)

        gadgets_len = []
        last_instr_len = []

        for g in gadgets:
            gadgets_len.append(len(ELF.from_assembly(
                g, arch='amd64').executable_segments[0].data()))
            try:
                last_instr_len.append(len(ELF.from_assembly(' ; '.join(g.split(';')[:-1]).strip(),
                                                            arch='amd64').executable_segments[0].data()))
            except:
                last_instr_len.append(0)

        self.dsts = self._reduce_len(gadgets_len)
        self.srcs = []

        self.srcsRaw = {}
        self.dstRaw = {}


        for i in range(len(self.dsts)):
            self.srcs.append(self.dsts[i] + last_instr_len[i])

        for sidx, s in enumerate(self.srcs):
            self.srcsRaw[s] = self.gadgetOrigAdrrs[sidx]
            self.dstRaw[self.dsts[sidx]] = self.gadgetOrigAdrrs[sidx]

        self.gadgetInfo = {}
        for gidx, addr in enumerate(self.gadgetOrigAdrrs):
            self.gadgetInfo[addr] = [self.dsts[gidx], self.srcs[gidx], gadgets[gidx]]

    def findNOPs(self, start_gadget, end_gadget):
        src = self.gadgetInfo[start_gadget][1]
        dst = self.gadgetInfo[end_gadget][0]
        self.state = self.proj.factory.blank_state(addr = src, add_options={angr.options.SYMBOL_FILL_UNCONSTRAINED_MEMORY,
                                                                angr.options.ZERO_FILL_UNCONSTRAINED_REGISTERS, 
                                                                angr.options.CONSTRAINT_TRACKING_IN_SOLVER})
        gadget = self.gadgetInfo[start_gadget][2]
        encoder = Encoder()
        constraints = encoder.encode(gadget)
        concrete_constraints = [ [constraint, np.random.randint(1, 100)] for constraint in constraints ]

        for cidx, constraint in enumerate(constraints):
            encoder.decode(self.state, constraint, concrete_constraints[cidx][1]) 

        concrete_constraints.append(['regs.rip', dst])

        self.sim = self.proj.factory.simgr(self.state)
        condition = Predicate(constraints = concrete_constraints)()
        return self.sim.explore(find=condition)

    def Oracle(self):
        def rawMapper(src, dst):
            rawSrc = self.srcsRaw[src]
            rawDst = self.dstRaw[dst]
            res = self.extractor(rawSrc)

            if '*' in res:
                return True

            res = [int('0x'+r, 16) for r in res]
            return int(rawDst, 16) in res 

        return rawMapper

    @staticmethod
    def _reduce_len(lst):
        out = []
        for i in range(len(lst)):
            if i == 0:
                out.append(0)
            else:
                out.append(out[i-1] + lst[i-1])
        return out

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--gadgets')
    parser.add_argument('--bname')

    args = parser.parse_args()

    extractor = GadgetExtractor(args.bname)
    NF = NOPsFinder(args.gadgets, extractor)
    NOP = NF.findNOPs('0x000000000040106e', '0x00000000004012cc')
    NOP = NOP.found[0]
    print(NOP.solver.eval(NOP.regs.rip))