class Patcher:
    """Patch processor, applies patches to the target binary

    Args:
        binary: Either a path to the target binary or the raw bytes
            of the binary
        cle_opts: An optional replacement set of options to pass to
            cle.Loader
    """

    def __init__(
        self,
        binary: Union[Path, str, bytes],
        use_angr: bool = True,
        cle_opts: Optional[Dict[str, bool]] = None,
        cfg_opts: Optional[Dict[str, bool]] = None,
    ) -> None:
        """Set up patcher with the target binary"""
        self.binary = BinaryManager(binary, use_angr, cle_opts, cfg_opts)

    def apply(self, patch: PatchType) -> None:
        """Apply a patch to the target binary

        Args:
            patch: The patch to apply, can be any of the patch types
        """

        def camel_to_snake(name: str) -> str:
            """
            Convert a camel case string to snake case
            """
            name = sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
            return sub("([a-z0-9])([A-Z])", r"\1_\2", name).lower()

        patch_dispatch_func_name = f"apply_{camel_to_snake(patch.__class__.__name__)}"
        if not hasattr(self, patch_dispatch_func_name):
            raise NotImplementedError(
                f"Patch type {patch.__class__.__name__} is not supported."
            )

        patch_dispatch_func = getattr(self, patch_dispatch_func_name)
        patch_dispatch_func(patch)

    def save(self, path: Union[str, Path]) -> None:
        """
        Save the patched binary to the given path
        """
        if isinstance(path, str):
            path = Path(path)

        self.binary.save(path)

        path.chmod(0o755)

    def apply_nop_patch(self, patch: NopPatch) -> None:
        """
        Apply a nop patch to the target binary
        """
        if self.binary.use_angr:
            nop_len = len(cast(Arch, self.binary.cle_binary.arch).nop_instruction)

            for address_range in patch.address_ranges:
                for address in range(
                    address_range.start,
                    max(address_range.end, address_range.start + nop_len),
                    nop_len,
                ):
                    self.binary.write(
                        address, cast(Arch, self.binary.cle_binary.arch).nop_instruction
                    )
        else:
            for address_range in patch.address_ranges:
                self.binary.write(address_range.start, b"\x90" * address_range.length)

    def apply_invert_branch_patch(self, patch: InvertBranchPatch) -> None:
        """
        Apply a branch patch to the target binary
        """
        raise NotImplementedError("Branch patches are not supported.")

    def apply_always_branch_patch(self, patch: AlwaysBranchPatch) -> None:
        """
        Apply a branch patch to the target binary
        """

    def apply_never_branch_patch(self, patch: NeverBranchPatch) -> None:
        """
        Apply a branch patch to the target binary
        """
        raise NotImplementedError("Branch patches are not supported.")

    def apply_skip_and_return_patch(self, patch: SkipAndReturnPatch) -> None:
        """
        Apply a branch patch to the target binary
        """
        raise NotImplementedError("Branch patches are not supported.")

    def apply_function_replace_patch(self, patch: FunctionReplacePatch) -> None:
        """
        Apply a function replace patch to the target binary
        """
        raise NotImplementedError("Function replace patches are not supported.")

    def apply_caller_replace_patch(self, patch: CallerReplacePatch) -> None:
        """
        Apply a caller replace patch to the target binary
        """
        raise NotImplementedError("Caller replace patches are not supported.")

    def apply_init_patch(self, patch: InitPatch) -> None:
        """
        Apply an init patch to the target binary
        """
        raise NotImplementedError("Init patches are not supported.")

    def apply_fini_patch(self, patch: FiniPatch) -> None:
        """
        Apply a fini patch to the target binary
        """
        raise NotImplementedError("Fini patches are not supported.")

    def apply_data_patch(self, patch: DataPatch) -> None:
        """
        Apply a data patch to the target binary
        """
        self.binary.add_data(patch.data, patch.label)

    def apply_add_code_patch(self, patch: AddCodePatch) -> None:
        """
        Apply an add code patch to the target binary
        """
        self.binary.add_code(patch.code, patch.label)

    def apply_replace_code_patch(self, patch: ReplaceCodePatch) -> None:
        """
        Apply a replace code patch to the target binary
        """
        self.binary.write(patch.address, patch.code)

@dataclass
class AddCodePatch(CodePatch):
    """Patch that adds some code to the binary at some labeled location

    Attributes:
        label: The label to add the code at
    """

    label: Optional[str] = None

@dataclass
class AlwaysBranchPatch(BranchPatch):
    """Patch that converts a conditional branch into an unconditional branch
    that always takes the "true" branch
    """

@dataclass
class BranchPatch:
    """Base patch that modifies a conditional branch at a particular address

    Attributes:
        address: The address of the branch to modify
    """

    address: int

@dataclass
class CallerReplacePatch:
    """Patch that will redirect all or some callers of a function elsewhere

    Attributes:
        new_code: The new code to replace the function with
        function_address: The name of the function to replace
        callers: The set of callers to redirect, optional

    """

    new_code: Code
    function_address: Optional[int] = None
    callers: Set[int] = field(default_factory=set)

    def __post_init__(self) -> None:
        """
        Check that we either got a function address or callers, or both
        """
        if self.function_address is None and not self.callers:
            raise ValueError(
                "A function address and/or set of callers must be provided."
            )

@dataclass
class CodePatch:
    """Base patch that holds some code

    Attributes:
        code: The code to add

    """

    code: Code

@dataclass
class DataPatch:
    """Patch that adds some data with some protections

    Attributes:
        data: The data to add
        label: The label to give the data, optional
        read: Whether the data should be readable, defaults to True
        write: Whether the data should be writable, defaults to False
        execute: Whether the data should be executable, defaults to False
    """

    data: bytes
    read: bool = True
    write: bool = False
    exec: bool = False
    label: Optional[str] = None

@dataclass
class FiniPatch:
    """Patch that will run some code upon exit from the program

    Attributes:
        code: The code to run
        priority: The priority of the code to run, optional
    """

    code: Code
    priority: int = 0

@dataclass
class FunctionReplacePatch:
    """Patch that will replace a function's contents with new code

    Attributes:
        function_address: The name of the function to replace
        new_code: The new code to replace the function with
    """

    function_address: int
    new_code: Code

@dataclass
class InitPatch:
    """Patch that will run some code immediately on entry to the program

    Attributes:
        code: The code to run
        priority: The priority of the code to run, optional
    """

    code: Code
    priority: int = 0

@dataclass
class InvertBranchPatch(BranchPatch):
    """Patch that inverts the true/false branch targets of a conditional branch"""

@dataclass
class NeverBranchPatch(BranchPatch):
    """Patch that converts a conditional branch into an unconditional branch
    that always takes the "false" branch
    """

@dataclass
class NopPatch:
    """Patch that converts specific addresses or ranges of addresses to no-operations

    Attributes:
        address_ranges: The address ranges to convert to no-operations

    Args:
        addresses: Individual addresses to convert to no-operations, optional.
        address_ranges: The address ranges to convert to no-operations, optional.

    """

    addresses: InitVar[Optional[List[int]]] = None
    address_ranges: Set[AddressRange] = field(default_factory=set)

    def __post_init__(self, addresses: Optional[List[int]]) -> None:
        """Convert optionally provided addresses to address ranges

        Args:
            addresses: Addresses that can be provided in lieu of address ranges
        """
        if addresses is None:
            return

        for address in addresses:
            self.address_ranges.add(AddressRange(start=address, end=address))

@dataclass
class ReplaceCodePatch(CodePatch):
    """Patch that replaces some code with some other code

    Attributes:
        address: The address to replace the code at, optional. This "address" can either
            be a label, an actual address, a list of addresses, a function that takes a
            [TransformInfo][TransformInfo] and returns an address, or a function that
            takes a [TransformInfo][TransformInfo] and returns a list of addresses.

    """

    address: Union[
        str,
        Callable[[TransformInfo], int],
        int,
        List[int],
        Callable[[TransformInfo], List[int]],
    ]

@dataclass
class RuntimeResolverPatch(CodePatch):
    """Patch that adds a runtime resolver to the binary"""

@dataclass
class SkipAndReturnPatch:
    """Patch that skips a call to a subroutine at a particular address and
    fakes a return value as if the called function had returned it using
    the default calling convention

    Attributes:
        address: The address of the call to skip
        return_value: The value to return from the call
    """

    caller_address: int
    return_value: int

class BinaryManager:
    """Wrapper for a binary program to provide information for patching

    Attributes:
        path: The path to the binary, if it exists on disk
        blob: The binary blob of the binary
        lief_binary: The LIEF binary object, which provides the majority of the program
            information for modification
        angr_project: The angr project, which provides the CFG and other information
            for analysis and deep inspection
        cle_binary: The CLE binary object, which provides angr's loader information
        cle_opts: The kwargs passed to `cle.Loader`, which can be used to override
            the default loader options
        cfg_opts: The kwargs passed to `angr.Project.analyses.CFGFast`, which can be
            used to override the default CFG options
        writes: A list of `WriteOperation`s that are queued for application to the
            binary
        code_to_add: A dictionary of `Code` objects to add to the binary, keyed by
            the label they are associated with
        alignment: The alignment to use when aligning addresses
        data_to_add: A dictionary of data to add to the binary, keyed by the label
            the data is associated with

    Args:
        binary: If `binary` is a `Path` or `str` we will attempt to load it from that
            path on disk. If it is `bytes`, we will load the bytes as a binary blob.
        cle_opts: If provided, will override the kwargs passed to `cle.Loader`, defaults
            to sane options.
        cfg_opts: If provided, will override the kwargs passed to
            `angr.Project.analyses.CFGFast`, defaults to sane options.
        silence_angr_logs: If `True`, will silence angr's logging output, defaults to
            `False`
        alignment: The alignment to use when aligning addresses, defaults to 8, but will
            be overridden by lief's alignment if it discovers a different alignment when
            loading the binary

    """

    path: Optional[Path] = None
    blob: BytesIO
    lief_binary: LIEFBinary
    angr_project: Project
    cle_binary: Backend
    cle_opts: Dict[str, bool] = {
        "auto_load_libs": False,
        "use_system_libs": False,
    }
    cfg_opts: Dict[str, bool] = {
        "normalize": True,
        "data_references": False,
        "cross_references": False,
        "skip_unmapped_addrs": True,
        "force_complete_scan": False,
    }
    writes: List[WriteOperation] = []
    code_to_add: Dict[str, Code] = {}
    alignment: int = 0x1000
    data_to_add: Dict[str, bytes] = {}

    def __init__(
        self,
        binary: Union[Path, str, bytes],
        use_angr: bool = True,
        cle_opts: Optional[Dict[str, bool]] = None,
        cfg_opts: Optional[Dict[str, bool]] = None,
        silence_angr_logs: bool = False,
        alignment: int = 8,
    ) -> None:
        """Initialize the binary wrapper via one of several methods."""

        self.use_angr = use_angr

        if isinstance(binary, Path):
            self.path = binary
        elif isinstance(binary, str):
            self.path = Path(binary)
        elif isinstance(binary, bytes):
            self.path = None
        else:
            raise TypeError(
                f"Requested binary is of type {type(binary)}"
                ", expected Path, str, or bytes."
            )

        if self.path is not None and not self.path.is_file():
            raise FileNotFoundError(
                f"Requested binary {self.path} was not found or could not be opened."
            )

        if cle_opts is not None:
            self.cle_opts = cle_opts

        if cfg_opts is not None:
            self.cfg_opts = cfg_opts

        if silence_angr_logs:
            for logger_name in ("angr", "pyvex", "cle", "archinfo", "claripy"):
                getLogger(logger_name).setLevel("ERROR")

        self.alignment = alignment

        if self.path is None:
            self.blob = BytesIO(cast(bytes, binary))
        else:
            self.blob = BytesIO(self.path.read_bytes())

        self.load_lief_binary()

        self.load_angr_project()

    def reload_blob_from_lief(self) -> None:
        """Reload the blob from the current LIEF binary"""
        tempfile = NamedTemporaryFile(delete=False)
        temppath = Path(tempfile.name)

        try:
            tempfile.close()
            self.lief_binary.write(str(temppath.resolve()))
            self.blob = BytesIO(temppath.read_bytes())
        except Exception as e:
            raise e
        finally:
            temppath.unlink(missing_ok=True)

        self.load_lief_binary()

        self.load_angr_project()

    def load_lief_binary(self) -> None:
        """Load the binary into LIEF"""
        self.lief_binary = parse(self.blob.getbuffer())

    def load_angr_project(self) -> None:
        """Load the angr project from the binary blob"""

        if not self.use_angr:
            self.angr_project = None
            self.cle_binary = None
            return

        if self.proj_opts is None:
            self.proj_opts = {
                "main_opts": {"base_addr": self.lief_binary.imagebase},
                "load_options": self.cle_opts,
            }
        self.angr_project = Project(self.blob, **self.proj_opts)
        self.angr_project.analyses.CFGFast(  # type: ignore
            **self.cfg_opts,
        )

        if self.angr_project.loader.main_object is None:
            raise FileNotFoundError(
                f"Requested binary {self.path} " "was not found or could not be opened."
            )

        self.cle_binary = self.angr_project.loader.main_object

    def align(
        self,
        address: int,
        alignment: Optional[int] = None,
    ) -> int:
        """Align an address to the specified alignment. If no alignment is
        provided, the default alignment will be used.

        Args:
            address: The address to align
            alignment: The alignment to use, defaults to the default alignment of the
                binary
        """
        if alignment is None:
            alignment = self.alignment

        address = (address + (alignment - 1)) & ~(alignment - 1)

        return address

    def write(
        self,
        where: Union[
            str,
            Callable[[TransformInfo], int],
            int,
            List[int],
            Callable[[TransformInfo], List[int]],
        ],
        data: Union[bytes, Callable[[Dict[str, int]], bytes], Code],
    ) -> None:
        """Write data to the binary at the given virtual address

        Args:

            where: Either an address or a label that will resolve to an address, or
                a function that takes the transform info after segment modification and
                returns an address
            data: The bytes or a function that takes a dictionary of labels
                and addresses and returns bytes to write to the patch
        """

        operation = WriteOperation(data, where)
        logger.info(f"Queuing write operation: {operation}")
        self.writes.append(operation)

    def read(self, vaddr: int, size: int) -> bytes:
        """Read data from the binary at the given virtual address

        Args:
            vaddr: The virtual address to read from
            size: The number of bytes to read
        """
        try:
            if self.use_angr:
                offset = next(
                    filter(lambda s: s.contains_addr(vaddr), self.cle_binary.sections)
                ).addr_to_offset(vaddr)
            else:
                offset = self.lief_binary.virtual_address_to_offset(vaddr)

        except StopIteration as e:
            logger.error(f"Could not find section for address {vaddr}")
            raise NoSectionError(f"Could not find section for address {vaddr}") from e

        self.blob.seek(offset)
        return self.blob.read(size)

    def add_code(self, code: Code, label: Optional[str] = None) -> None:
        """Mark some code as being added on next save

        Args:
            code: The binary code being added
            label: The label to associate with the code for resolution of
                other patches referencing it
        """
        if label is None:
            label = f"code_{len(self.code_to_add)}"

        logger.info(f"Queueing code addition at label {label}: {code}")

        code.label = label

        self.code_to_add[label] = code

    def add_data(self, data: bytes, label: Optional[str] = None) -> None:
        """Add some data to the binary

        Args:
            data: The data to add
            label: The label to associate with the data location
        """
        if label is None:
            label = f"data_{len(self.data_to_add)}"

        logger.info(f"Queueing data addition at label {label}: ({len(data)} bytes)")

        self.data_to_add[label] = data

    def apply(self) -> None:
        """Apply patches to the binary"""
        # Offsets is filled in twice:
        # - first time fills in offset in the new section
        # - second time fills in offset from base address
        transform_info = TransformInfo(self.lief_binary, self.angr_project)

        transform_info.code_size = 0
        transform_info.data_size = 0
        transform_info.all_data = b""
        transform_info.data_offsets = {}
        transform_info.code_offsets = {}

        # Figure out how much space we need for code
        # Each code patch is dummy compiled and placed in a subsection that is aligned
        # to the alignment of the binary.
        for label, code in self.code_to_add.items():
            code.reset()
            code.dummy()
            compiled = code.compile(
                "dummy",
            )

            # TODO: There might be a smarter way but we don't really prioritize size
            aligned_size = self.align(len(compiled), self.alignment) * 2

            transform_info.code_offsets[label] = transform_info.code_size
            transform_info.code_size += aligned_size

        for label, data in self.data_to_add.items():
            transform_info.data_offsets[label] = len(data)
            transform_info.data_size += len(data)
            transform_info.all_data += data

        # Round up code and data sizes to the next multiple of self.alignment
        transform_info.code_size = (
            transform_info.code_size + (self.alignment - 1)
        ) & ~(self.alignment - 1)

        transform_info.data_size = (
            transform_info.data_size + (self.alignment - 1)
        ) & ~(self.alignment - 1)

        # Create new sections

        if transform_info.data_size > 0:
            new_data_segment = Segment()
            new_data_segment.content = list(transform_info.all_data)
            new_data_segment.type = SEGMENT_TYPES.LOAD
            new_data_segment.alignment = self.alignment
            new_data_segment.flags = SEGMENT_FLAGS(SEGMENT_FLAGS.R | SEGMENT_FLAGS.W)
            new_data_segment = self.lief_binary.add(new_data_segment)

            transform_info.data_base = new_data_segment.virtual_address

            # Fix up offsets to they point to the actual address
            for label, addr in transform_info.data_offsets.items():
                transform_info.data_offsets[label] = (
                    cast(int, transform_info.data_base) + addr
                )

        if transform_info.code_size > 0:
            new_code_segment = Segment()
            new_code_segment.content = list(b"\x00" * transform_info.code_size)
            new_code_segment.type = SEGMENT_TYPES.LOAD
            new_code_segment.alignment = self.alignment
            new_code_segment.flags = SEGMENT_FLAGS(SEGMENT_FLAGS.X | SEGMENT_FLAGS.R)
            new_code_segment = self.lief_binary.add(new_code_segment)

            transform_info.code_base = new_code_segment.virtual_address

            # Fix up offsets to they point to the actual address
            for label, addr in transform_info.code_offsets.items():
                transform_info.code_offsets[label] = (
                    cast(int, transform_info.code_base) + addr
                )

            # Queue writes for code
            for label, code in self.code_to_add.items():
                self.writes.append(WriteOperation(code, label))

        # Reload the angr project to get the new sections
        self.reload_blob_from_lief()
        transform_info.lief_binary = self.lief_binary
        transform_info.angr_project = self.angr_project

        for write in self.writes:
            offsets = []

            if isinstance(write.where, str):
                offsets.append(
                    transform_info.data_offsets.get(
                        write.where, transform_info.code_offsets.get(write.where, None)
                    )
                )

            elif isinstance(write.where, int):
                offsets.append(write.where)

            elif isinstance(write.where, list):
                offsets.extend(write.where)
            else:
                new_offsets = write.where(transform_info)

                if not isinstance(offsets, list):
                    offsets.append(new_offsets)
                else:
                    offsets.extend(new_offsets)

            for offset in offsets:
                logger.debug(f"Writing {write.data} to {offset:#0x}")
                transform_info.current_offset = offset

                if isinstance(write.data, bytes):
                    data = write.data

                elif isinstance(write.data, Code):

                    write.data.reset()
                    write.data.build(transform_info)
                    data = write.data.compile(
                        cast(str, write.data.label), transform_info
                    )

                    if self.use_angr:
                        disassembly = self.angr_project.arch.disasm(data, offset)
                    else:
                        logger.debug(f"Disassembly for {data} @ {offset}")
                        disassembler = Cs(CS_ARCH_X86, CS_MODE_64)
                        disassembler.detail = True
                        disassembly = "\n".join(
                            map(
                                lambda i: f"{i.address:#x}: {i.mnemonic} {i.op_str}",
                                disassembler.disasm(data, offset),
                            )
                        )

                    logger.debug(f"Disassembly of data for label {write.data.label}:")

                    for disas_line in disassembly.splitlines():
                        logger.debug(f"  {disas_line}")
                else:
                    data = write.data(transform_info.code_offsets)

                logger.info(f"Writing {len(data)} bytes to {offset:#0x}")
                self.lief_binary.patch_address(offset, list(data))

    def save(self, where: Path) -> None:
        """Apply any pending operations and save the binary to a file

        Args:
            where: The path to the destination to save the binary
        """
        logger.info("Applying pending operations")
        self.apply()

        logger.info(f"Saving binary to {where}")
        self.lief_binary.write(str(where.resolve()))

    def asm(self, asm: str, vaddr: int) -> bytes:
        """Assemble the given assembly code at the given virtual address

        Args:
            asm: The assembly code to assemble
            vaddr: The virtual address to assemble at
        """
        if self.use_angr:
            return self.cle_binary.arch.asm(asm, vaddr, as_bytes=True)  # type: ignore
        else:
            assembler = Ks(KS_ARCH_X86, KS_MODE_64)
            return assembler.asm(asm, vaddr, as_bytes=True)  # type: ignore

    def disasm(self, vaddr: int) -> CsInsn:
        """Disassemble one instruction from the binary at an address

        Args:
            vaddr: The virtual address to disassemble at
        """

        if self.use_angr:
            block = (
                self.binary.angr_project.kb.cfgs["CFGFast"]
                .get_any_node(vaddr, anyaddr=True)
                .block
            )

            for instr in block.capstone.insns:
                if instr.address == vaddr:
                    return instr

            raise ValueError(f"Could not find instruction at {vaddr:#0x}")
        else:
            logger.warning("Disassembly is not supported without angr")