from __future__ import print_function
import argparse
import re
import capstone
import six
import sys
import pwnypack.codec
import pwnypack.elf
import pwnypack.main
import pwnypack.asm
import pwnypack.target
import pwnypack.util


__all__ = [
    'find_gadget',
]


INVALID_GROUPS = set((capstone.CS_GRP_CALL, capstone.CS_GRP_JUMP))


def find_gadget(elf, gadget, align=1, unique=True):
    if not isinstance(elf, pwnypack.elf.ELF):
        elf = pwnypack.elf.ELF(elf)

    matches = []
    gadgets = []

    if isinstance(gadget, six.binary_type):
        gadget = re.compile(re.escape(gadget))

    for section in elf.section_headers:
        if section.type != section.Type.progbits:
            continue

        for match in gadget.finditer(section.content):
            match_index = match.start()
            if match_index % align != 0:
                continue

            match_gadget = match.group()

            if match_gadget in gadgets:
                continue

            match_addr = section.addr + match_index

            md = pwnypack.asm.prepare_capstone(syntax=pwnypack.asm.AsmSyntax.intel, target=elf)
            md.detail = True
            match_asm = []

            for insn in md.disasm(match_gadget, match_addr):
                if insn.id == capstone.CS_OP_INVALID or set(insn.groups) & INVALID_GROUPS:
                    # Don't try to disassemble this particular gadget again.
                    gadgets.append(match_gadget)
                    break
                match_asm.append((insn.mnemonic + ' ' + insn.op_str).strip())
            else:
                matches.append({
                    'section': section,
                    'offset': match_index,
                    'addr': match_addr,
                    'gadget': match_gadget,
                    'asm': match_asm,
                })
                if unique:
                    gadgets.append(match_gadget)

    return matches


@pwnypack.main.register('gadget')
def gadget_app(_parser, cmd, args):  # pragma: no cover
    """
    Find ROP gadgets in an ELF binary.
    """

    parser = argparse.ArgumentParser(
        prog=_parser.prog,
        description=_parser.description,
    )
    parser.add_argument('file', help='ELF file to find gadgets in')
    parser.add_argument('gadget', help='the assembler source or reghex expression')
    parser.add_argument(
        '--reghex', '-r',
        dest='mode',
        action='store_const',
        const='reghex',
        help='use reghex expression (hex bytes interspaced with ? for wildcard)',
    )
    parser.add_argument(
        '--asm', '-a',
        dest='mode',
        action='store_const',
        const='asm',
        help='use assembler expression (separate lines with semi-colon)',
    )
    parser.add_argument(
        '--all', '-l',
        dest='unique',
        action='store_const',
        const=False,
        default=True,
        help='also show non-unique gadgets',
    )
    args = parser.parse_args(args)

    if args.mode is None:
        try:
            pwnypack.util.reghex(args.gadget)
            args.mode = 'reghex'
        except SyntaxError:
            args.mode = 'asm'

    elf = pwnypack.elf.ELF(args.file)

    if args.mode == 'reghex':
        try:
            gadget = pwnypack.util.reghex(args.gadget)
        except SyntaxError:
            print('Invalid reghex pattern.')
            sys.exit(1)
    else:
        try:
            gadget = pwnypack.asm.asm(args.gadget.replace(';', '\n'), target=elf)
        except SyntaxError as e:
            print('Could not assemble:', e.msg)
            sys.exit(1)

    matches = find_gadget(
        elf,
        gadget,
        unique=args.unique
    )

    if not matches:
        print('No gadgets found.', file=sys.stdout)
        return

    longest_gadget = max(len(m['gadget']) for m in matches)
    fmt = '  0x%%0%dx: [ %%-%ds ] %%s' % (elf.bits / 4, longest_gadget * 3 - 1)

    current_section = None

    for match in matches:
        if match['section'].name != current_section:
            if current_section is not None:
                print()
            print('Section: %s' % match['section'].name)
            current_section = match['section'].name

        hex_gadget = pwnypack.codec.enhex(match['gadget'])
        print(fmt % (
            match['addr'],
            ' '.join(
                hex_gadget[i:i+2]
                for i in range(0, len(hex_gadget), 2)
            ),
            ' ; '.join(match['asm'])
        ))

    print()

@@ -1,8 +1,123 @@
+%22%22%22%0AThe ROP module contains a function to find gadgets in ELF binaries that can%0Abe used to create ROP chains.%0A%22%22%22%0A%0A
 from __f
@@ -491,24 +491,1142 @@
 ique=True):%0A
+    %22%22%22%0A    Find a ROP gadget in a the executable sections of an ELF executable or%0A    library. The ROP gadget can be either a set of bytes for an exact match%0A    or a (bytes) regular expression. Once it finds gadgets, it uses the%0A    capstone engine to verify if the gadget consists of valid instructions%0A    and doesn't contain any call or jump instructions.%0A%0A    Args:%0A        elf(:class:%60~pwnypack.elf.ELF%60): The ELF instance to find a gadget in.%0A        gadget(bytes or regexp): The gadget to find.%0A        align(int): Make sure the gadget starts at a multiple of this number%0A        unique(bool): If true, only unique gadgets are returned.%0A%0A    Returns:%0A        dict: A dictionary containing a description of the found%0A            gadget. Contains the following fields:%0A%0A            - section: The section the gadget was found in.%0A            - offset: The offset inside the segment the gadget was found at.%0A            - addr: The virtual memory address the gadget will be located at.%0A            - gadget: The machine code of the found gadget.%0A            - asm: A list of disassembled instructions.%0A%0A    %22%22%22%0A%0A
     if not i

Add docstrings for pwnypack.rop.