<commit_before>import angr
import heapq
import struct
import claripy

from . import rop_utils
from . import common

from .errors import RopException
from .rop_chain import RopChain
from .rop_gadget import RopGadget

import types
import logging
from collections import defaultdict

l = logging.getLogger("angrop.chain_builder")


class ChainBuilder(object):
    """
    This class provides functions to generate common ropchains based on existing gadgets.
    """

    def __init__(self, project, gadgets, duplicates, reg_list, base_pointer, badbytes, roparg_filler):
        """
        Initializes the chain builder.

        :param project: Angr project
        :param gadgets: a list of RopGadget gadgets
        :param duplicates:
        :param reg_list: A list of multipurpose registers
        :param base_pointer: The name ("offset") for base pointer register
        :param badbytes: A list with badbytes, which we should avaoid
        :param roparg_filler: An integer used when popping superfluous registers
        """
        self.project = project
        self._gadgets = gadgets
        # TODO get duplicates differently?
        self._duplicates = duplicates
        self._reg_list = reg_list
        self._base_pointer = base_pointer
        self.badbytes = badbytes
        self._roparg_filler = roparg_filler

        self._syscall_instruction = None
        if self.project.arch.linux_name == "x86_64":
            self._syscall_instructions = {b"\x0f\x05"}
        elif self.project.arch.linux_name == "i386":
            self._syscall_instructions = {b"\xcd\x80"}

        self._execve_syscall = None
        if "unix" in self.project.loader.main_object.os.lower():
            if self.project.arch.bits == 64:
                self._execve_syscall = 59
            elif self.project.arch.bits == 32:
                self._execve_syscall = 11
            else:
                raise RopException("unknown unix platform")

        # test state
        self._test_symbolic_state = rop_utils.make_symbolic_state(self.project, self._reg_list)

        # filtered gadget cache
        self._filtered_reg_gadgets = None

    def set_regs(self, modifiable_memory_range=None, use_partial_controllers=False, rebase_regs=None, **registers):
        """
        :param registers: dict of registers to values
        :return: a chain which will set the registers to the requested values

        example:
        chain = rop.set_regs(rax=0x1234, rcx=0x41414141)
        """

        if len(registers) == 0:
            return RopChain(self.project, self)

        if rebase_regs is None:
            rebase_regs = set()

        gadgets, best_stack_change, _ = self._find_reg_setting_gadgets(modifiable_memory_range,
                                                                       use_partial_controllers, **registers)
        if gadgets is None:
            raise RopException("Couldn't set registers :(")

        return self._build_reg_setting_chain(gadgets, modifiable_memory_range,
                                             registers, best_stack_change, rebase_regs)

    # TODO handle mess ups by _find_reg_setting_gadgets and see if we can set a register in a syscall preamble
    # or if a register value is explicitly set to just the right value
    def do_syscall(self, syscall_num, arguments, ignore_registers=None, modifiable_memory_range=None,
                   use_partial_controllers=False, rebase_regs=None, needs_return=True):
        """
        build a rop chain which performs the requested system call with the arguments set to 'registers' before
        the call is made
        :param syscall_num: the syscall number to execute
        :param arguments: the register values to have set at system call time
        :param ignore_registers: list of registers which shouldn't be set
        :return: a RopChain which makes the system with the requested register contents
        """

        registers = {}

        if ignore_registers is None:
            ignore_registers = []

        # set the system call number
        registers[self.project.arch.register_names[self.project.arch.syscall_num_offset]] = syscall_num

        cc = angr.SYSCALL_CC[self.project.arch.name]["default"](self.project.arch)

        # distinguish register arguments from stack arguments
        register_arguments = arguments
        stack_arguments = []
        if len(arguments) > len(cc.ARG_REGS):
            register_arguments = arguments[:len(cc.ARG_REGS)]
            stack_arguments = arguments[len(cc.ARG_REGS):]

        for arg, reg in zip(register_arguments, cc.ARG_REGS):
            registers[reg] = arg

        # remove any registers which have been asked to be ignored
        for reg in ignore_registers:
            if reg in registers:
                del registers[reg]

        # first find gadgets to the set the registers
        chain = self.set_regs(modifiable_memory_range, use_partial_controllers, rebase_regs, **registers)

        # find small stack change syscall gadget that also fits the stack arguments we want
        smallest = None
        for gadget in filter(lambda g: g.starts_with_syscall, self._gadgets):
            # adjust stack change for ret
            stack_change = gadget.stack_change - self.project.arch.bytes
            required_space = len(stack_arguments) * self.project.arch.bytes
            if stack_change >= required_space:
                if smallest is None or gadget.stack_change < smallest.stack_change:
                    smallest = gadget

        if smallest is None and not needs_return:
            syscall_locs = self._get_syscall_locations()
            if len(syscall_locs) > 0:
                smallest = RopGadget(syscall_locs[0])
                smallest.stack_change = self.project.arch.bits

        if smallest is None:
            raise RopException("No suitable syscall gadgets found")

        chosen_gadget = smallest
        # add the gadget to the chain
        chain.add_gadget(chosen_gadget)

        # now we'll just pad it out with zeroes, in the future we'll probably want a way to be smart about
        # the next gadget in the chain

        # add the syscall gadget's address
        chain.add_value(chosen_gadget.addr, needs_rebase=True)

        # remove one word to account for the ret
        padding_bytes = chosen_gadget.stack_change - self.project.arch.bytes<commit_msg>fix block_length on syscall gadgets<commit_after>import angr
import heapq
import struct
import claripy

from . import rop_utils
from . import common

from .errors import RopException
from .rop_chain import RopChain
from .rop_gadget import RopGadget

import types
import logging
from collections import defaultdict

l = logging.getLogger("angrop.chain_builder")


class ChainBuilder(object):
    """
    This class provides functions to generate common ropchains based on existing gadgets.
    """

    def __init__(self, project, gadgets, duplicates, reg_list, base_pointer, badbytes, roparg_filler):
        """
        Initializes the chain builder.

        :param project: Angr project
        :param gadgets: a list of RopGadget gadgets
        :param duplicates:
        :param reg_list: A list of multipurpose registers
        :param base_pointer: The name ("offset") for base pointer register
        :param badbytes: A list with badbytes, which we should avaoid
        :param roparg_filler: An integer used when popping superfluous registers
        """
        self.project = project
        self._gadgets = gadgets
        # TODO get duplicates differently?
        self._duplicates = duplicates
        self._reg_list = reg_list
        self._base_pointer = base_pointer
        self.badbytes = badbytes
        self._roparg_filler = roparg_filler

        self._syscall_instruction = None
        if self.project.arch.linux_name == "x86_64":
            self._syscall_instructions = {b"\x0f\x05"}
        elif self.project.arch.linux_name == "i386":
            self._syscall_instructions = {b"\xcd\x80"}

        self._execve_syscall = None
        if "unix" in self.project.loader.main_object.os.lower():
            if self.project.arch.bits == 64:
                self._execve_syscall = 59
            elif self.project.arch.bits == 32:
                self._execve_syscall = 11
            else:
                raise RopException("unknown unix platform")

        # test state
        self._test_symbolic_state = rop_utils.make_symbolic_state(self.project, self._reg_list)

        # filtered gadget cache
        self._filtered_reg_gadgets = None

    def set_regs(self, modifiable_memory_range=None, use_partial_controllers=False, rebase_regs=None, **registers):
        """
        :param registers: dict of registers to values
        :return: a chain which will set the registers to the requested values

        example:
        chain = rop.set_regs(rax=0x1234, rcx=0x41414141)
        """

        if len(registers) == 0:
            return RopChain(self.project, self)

        if rebase_regs is None:
            rebase_regs = set()

        gadgets, best_stack_change, _ = self._find_reg_setting_gadgets(modifiable_memory_range,
                                                                       use_partial_controllers, **registers)
        if gadgets is None:
            raise RopException("Couldn't set registers :(")

        return self._build_reg_setting_chain(gadgets, modifiable_memory_range,
                                             registers, best_stack_change, rebase_regs)

    # TODO handle mess ups by _find_reg_setting_gadgets and see if we can set a register in a syscall preamble
    # or if a register value is explicitly set to just the right value
    def do_syscall(self, syscall_num, arguments, ignore_registers=None, modifiable_memory_range=None,
                   use_partial_controllers=False, rebase_regs=None, needs_return=True):
        """
        build a rop chain which performs the requested system call with the arguments set to 'registers' before
        the call is made
        :param syscall_num: the syscall number to execute
        :param arguments: the register values to have set at system call time
        :param ignore_registers: list of registers which shouldn't be set
        :return: a RopChain which makes the system with the requested register contents
        """

        registers = {}

        if ignore_registers is None:
            ignore_registers = []

        # set the system call number
        registers[self.project.arch.register_names[self.project.arch.syscall_num_offset]] = syscall_num

        cc = angr.SYSCALL_CC[self.project.arch.name]["default"](self.project.arch)

        # distinguish register arguments from stack arguments
        register_arguments = arguments
        stack_arguments = []
        if len(arguments) > len(cc.ARG_REGS):
            register_arguments = arguments[:len(cc.ARG_REGS)]
            stack_arguments = arguments[len(cc.ARG_REGS):]

        for arg, reg in zip(register_arguments, cc.ARG_REGS):
            registers[reg] = arg

        # remove any registers which have been asked to be ignored
        for reg in ignore_registers:
            if reg in registers:
                del registers[reg]

        # first find gadgets to the set the registers
        chain = self.set_regs(modifiable_memory_range, use_partial_controllers, rebase_regs, **registers)

        # find small stack change syscall gadget that also fits the stack arguments we want
        smallest = None
        for gadget in filter(lambda g: g.starts_with_syscall, self._gadgets):
            # adjust stack change for ret
            stack_change = gadget.stack_change - self.project.arch.bytes
            required_space = len(stack_arguments) * self.project.arch.bytes
            if stack_change >= required_space:
                if smallest is None or gadget.stack_change < smallest.stack_change:
                    smallest = gadget

        if smallest is None and not needs_return:
            syscall_locs = self._get_syscall_locations()
            if len(syscall_locs) > 0:
                smallest = RopGadget(syscall_locs[0])
                smallest.block_length = self.project.factory.block(syscall_locs[0]).size
                smallest.stack_change = self.project.arch.bits

        if smallest is None:
            raise RopException("No suitable syscall gadgets found")

        chosen_gadget = smallest
        # add the gadget to the chain
        chain.add_gadget(chosen_gadget)

        # now we'll just pad it out with zeroes, in the future we'll probably want a way to be smart about
        # the next gadget in the chain

        # add the syscall gadget's address
        chain.add_value(chosen_gadget.addr, needs_rebase=True)

        # remove one word to account for the ret
        padding_bytes = chosen_gadget.stack_change - self.project.arch.bytes
{'fraction_non_alphanumeric': 0.06621694967821974, 'fraction_numerical': 0.005815305885089556, 'mean_word_length': 2.423944783647465, 'pattern_counts': {'":': 4, '<': 5, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1966669', 'n_tokens_mistral': 3670, 'n_tokens_neox': 3438, 'n_words': 1065}