path: root/sploit/payload/rop.py

from graphlib import TopologicalSorter

from sploit.arch import arch, btoi, itob
from sploit.payload.gadhint import GadHint
from sploit.payload.payload import Payload
from sploit.payload.payload_entry import padalign, padlen

_POP_MAGIC = 0xdead
_SPM_MAGIC = b"\x69"
_ERROR_MAGIC = 0xbaadc0de

class ROP:
    """
    ROP chain generation tool

    This class contains methods for automating basic return-oriented programming
    workloads, such as loading register values and calling into arbitrary
    functions or syscalls.  The tools are currently designed to work on x86
    (32 or 64 bit) and ARM (32 bit only).

    The main appeal of the ROP class is the ability to abstract away the manual
    construction of ROP chain data, and instead make declarative statements
    like "call this function with these arguments".  The ROP class will also
    utilize its supplied binary objects to automatically find and use trivial
    gadgets.

    The user is able to provide annotations for more complicated gadgets, which
    help instruct the class how to incorporate them into a ROP chain.  This is
    done with the GadHint dataclass.  GadHint objects are provided to a ROP
    instance by including them in the Symtbl of one of the binary objects it is
    constructed with.  If applicable, a user-supplied gadget will take
    precedence over automatic gadget searching.  See the GadHint module to learn
    more about the descriptive attributes that are supported.

    objects (list[ELF]): The binary objects this ROP instance will consider for
    gadget searching.  If one of these is the target executable binary, it
    should appear first in the list.

    safe_syscalls (bool): If True, require that automatically found syscall
    instructions are immediately followed by a return instruction.

    align_calls (bool): If True, ensure that the stack return address into
    function calls is aligned according to the architecture alignment property.

    clean_stack (bool): If True, attempt to locate a cleaning gadget to "pop"
    stack data that is leftover from a function call.  Required if attempting
    to make multiple calls that involve stack-based arguments.
    """

    def __init__(self, *objects, safe_syscalls=True, align_calls=True,
                 clean_stack=True):
        """Construct new ROP builder."""
        self.objects = objects
        self.safe_syscalls = safe_syscalls
        self.align_calls = align_calls
        self.clean_stack = clean_stack

    def search_gadgets(self, *regexes, cont=False):
        """Return a list of matching gadgets, considering all objects."""
        results = []
        for obj in self.objects:
            results += obj.gadgets(*regexes, cont=cont)
        return results

    def search_gadget(self, *regexes):
        """Return the first matching gadget, considering all objects."""
        for obj in self.objects:
            try:
                return obj.gadget(*regexes)
            except:
                pass
        raise LookupError(
            f"ROP: Need to define gadget symbol for {'; '.join(regexes)}")

    def gadget(self, gadget):
        """
        Return a generic ROP payload.

        gadget (GadHint): Annotated gadget to prepare a chain from.
        """
        return self.__build_chain(gadget, {})

    def assign(self, **sets):
        """
        Return a ROP payload to control given registers.

        **sets (str:int): Keyword arguments specify register assignments to
        perform with this ROP chain.  Argument names correspond to register
        names.
        """
        return self.gadget(GadHint(requirements=sets))

    def call(self, func, *params):
        """
        Return a ROP payload to call function.

        func (int): Entry address of function to call.
        *params (int): Remaining positional args are passed to func.
        """
        register_params = dict(zip(arch.funcargs, params))
        stack_params = params[len(register_params):]
        gadget = GadHint(func, requirements=register_params, stack=stack_params,
                         align=self.align_calls)
        return self.gadget(gadget)

    def syscall(self, *params):
        """
        Return a ROP payload to call kernel.

        *params (int): The first argument is the syscall number.  Remaining
        positional arguments are passed to the syscall.
        """
        if len(params) > len(arch.kernargs):
            raise TypeError("ROP: Too many arguments passed to syscall. "
                f"Target architecture supports up to {len(arch.kernargs)-1}.")

        register_params = dict(zip(arch.kernargs, params))
        sc = self.__get_gadget("syscall", {})
        return self.gadget(sc.with_requirements(register_params))

    def memcpy(self, dst, src):
        """
        Return a ROP payload to write data into memory.

        dst (int): The destination memory address.
        src (bytes): The content to write.
        """
        data = Payload()
        for idx in range(0, len(src), arch.wordsize):
            word = btoi(src[idx:idx+arch.wordsize])
            data(self.gadget(self.__get_write(dst+idx, word)))
        return data

    def __get_hints(self):
        """Return all user-supplied gadget hints."""
        return [h for obj in self.objects for _,h in obj.sym if type(h) is GadHint]

    def __discover_requirements(self, seen, graph, current):
        """
        Populate gadget dependency graph.

        This function recursively looks up gadgets to ensure all necessary
        required gadgets can be found, and stores this information into the
        given graph object.  Established dependencies encode the order that the
        chain builder should attempt to satisfy register requirements.
        Dependency loops are detected by the TopologicalSorter.

        seen (set): Set of (register,value) tuples we have already discovered.
        graph (TopologicalSorter): Dependency graph model object.
        current (GadHint): Current gadget we are processing.
        """
        for r, v in current.requirements.items():
            # We key on register name _and_ value because some gadgets may
            # only be capable of storing specific values in a target register.
            # Requiring a register to store different values may require the
            # use of multiple gadgets.
            if (r, v) not in seen:
                gadget = self.__get_gadget(r, current.requirements)

                # Add gadget's requirements to the dependency graph.
                # We say that each requirement is a 'successor' to this
                # current gadget 'r', so that the chain builder will satisfy
                # 'r' first.  This prevents the fulfillment of 'r' from
                # clobbering targets it requires, as the builder will satisfy
                # them afterward.
                for x in gadget.requirements:
                    graph.add(x, r)

                # Treat gadget's load immediates as pseudo-requirements for
                # the sake of target ordering, following the same logic
                # as above.
                for x in gadget.imms:
                    graph.add(x, r)

                # Mark node as visited
                seen.add((r, v))
                self.__discover_requirements(seen, graph, gadget)

    def __get_gadget(self, target, sets):
        """
        Get context-specific gadget.

        target (str): Either "ret", "syscall", or the name of a register we
        would like to modify.

        sets (dict{str:int}): The set of other register requirements we are
        trying to fulfill in parallel.  Values may affect the gadget we decide
        to use.
        """
        # First, consider user-provided hints before automatically locating
        # gadgets.
        for hint in self.__get_hints():
            # Setup additional requirements based on hint's register moves.
            # If a mov target is in sets, require to set the src to the 'sets'
            # value.
            addl_reqs = { src:sets[dst] for dst, src in hint.movs.items() if dst in sets }
            hint = hint.with_requirements(addl_reqs)

            # Pops will be accounted for by the chain builder.
            # Immediates will be handled by gadget ordering in chain builder.
            # Writes are a non-issue here.

            if hint.syscall:
                # Only consider syscalls if the target is syscall.
                if target == "syscall":
                    return hint
            elif target in hint.imms:
                if hint.imms[target] == sets[target]:
                    return hint
            elif target in hint.pops:
                return hint
            elif target in hint.movs:
                return hint

        # Automatically locate simple gadgets
        if target == "ret":
            return GadHint(self.search_gadget(arch.ret))

        if target == "syscall":
            insns = [arch.syscall, arch.ret] if self.safe_syscalls else [arch.syscall]
            return GadHint(self.search_gadget(*insns), syscall=True)

        # target == register
        insns = [ i.format(target) for i in arch.popgad ]
        return GadHint(self.search_gadget(*insns), pops=[target])

    def __get_clean(self, size):
        """
        Get a stack cleaning gadget that moves sp by _at least_ size.

        size (int): Minimum stack pointer move.
        """
        # spm values of zero (the default) can't be trusted, as in this case
        # the user likely hasn't annotated the GadHint properly.  Returning a
        # larger move than requested is fine, since the chain builder can insert
        # junk to be popped.
        for hint in self.__get_hints():
            if hint.spm >= size and hint.spm > 0:
                return hint

        results = self.search_gadgets(*arch.cleangad)
        table = { int(g.asm[0].group(1), 0): g for g in results }
        sizes = sorted([ x for x in table.keys() if x >= size ])

        if len(sizes) > 0:
            return GadHint(table[sizes[0]], spm=sizes[0])

        raise LookupError(
            f"ROP: Need to define a stack move gadget of at least {size}")

    def __get_write(self, dst, src):
        """
        Get a memory write gadget, injected with requirements for user data.

        dst (int): The intended memory write location.
        src (int): The intended value to write.
        """
        # If any exist, take the first write provided by user hints, assuming
        # the user's intent to specifically use _this_ write.  Follow-on gadgets
        # to setup the dst and src registers must be findable.
        for hint in self.__get_hints():
            if hint.writes:
                d, s = list(hint.writes.items())[0]
                return hint.with_requirements({d:dst, s:src})

        # Only take an automatic write gadget if we can prove up front that its
        # requirements can be met, otherwise move on.  A later search result may
        # pass the test.
        results = self.search_gadgets(*arch.writegad)

        for gad in results:
            d = gad.asm[0].group("dst")
            s = gad.asm[0].group("src")

            try:
                # Assert requirements are met.
                gadget = GadHint(gad, writes={d: s}, requirements={d:dst, s:src})
                self.__discover_requirements(set(), TopologicalSorter(), gadget)
                return gadget
            except:
                pass

        raise LookupError("ROP: Need to define gadgets for memory write / deps")

    def __build_chain(self, gadget, sets):
        """
        Generate ROP chain data for a given gadget.

        This function recursively builds a ROP chain for the given gadget and
        its requirements, returning the result as a Payload.

        gadget (GadHint): Current gadget to process.

        sets (dict{str:int}): The set of other register requirements we are
        trying to fulfill in parallel.
        """
        # Form a to-do-list of registers from our immediate requirements,
        # attempting to order them such that we avoid overwriting/conflicting
        # values.  This may not be possible, in which case graph.static_order()
        # will raise an exception.
        reqs = gadget.requirements
        graph = TopologicalSorter({ r:set() for r in reqs })
        self.__discover_requirements(set(), graph, gadget)
        to_do_list = [ x for x in graph.static_order() if x in reqs ]

        chain = Payload()

        # Start chain by satisfying to-do-list requirements.
        if len(to_do_list) > 0:
            chain.requirements = Payload()

        while len(to_do_list) > 0:
            r = to_do_list[0]
            g = self.__get_gadget(r, reqs)
            c = self.__build_chain(g, reqs)
            chain.requirements[f"{r}_{reqs[r]}"] = c

            # This gadget may satisfy multiple items in the to-do-list.
            # Specifically, all of its pop and mov targets, and any load
            # immediates that match our requirements.  Non-matching
            # immediates will be handled by a later gadget.
            imms = g.imms.keys() & reqs.keys()
            imms = [ x for x in imms if g.imms[x] == reqs[x] ]
            done = g.pops + list(g.movs) + imms
            to_do_list = [ x for x in to_do_list if x not in done ]

        # Append chain data to execute this gadget, but respect offset == 0
        # as a way to disable this gadget (perform a NULL gadget).
        if gadget.offset != 0:
            # Stack alignment if required.
            if gadget.align:
                ret = self.__get_gadget("ret", {})
                chain.alignment = padalign(0, itob(ret))

            # "Return address" entry into this gadget.
            chain.gadget = gadget.offset

            # The gadget's "inner stack data" will be values to be popped and
            # additional junk data to be deallocated by the gadget itself.
            if gadget.pops or gadget.spm > 0:
                chain.inner = Payload()
                chain.inner(*[ sets.get(p, _POP_MAGIC) for p in gadget.pops ])
                if gadget.spm > 0:
                    chain.inner.pad = padlen(gadget.spm, _SPM_MAGIC)

            # The gadget's "outer stack data" will be the additional values
            # explicitly specified by the gadget.  Append a separate gadget
            # to clean up these values.
            if gadget.stack:
                size = len(gadget.stack) * arch.wordsize

                if self.clean_stack:
                    clean = self.__get_clean(size)
                    chain.cleanup = clean.offset
                    pad = padlen(clean.spm, _SPM_MAGIC)
                else:
                    chain.cleanup = _ERROR_MAGIC
                    pad = None

                chain.outer = Payload()
                chain.outer(*gadget.stack)
                if pad: chain.outer.pad = pad

        return chain