Assemble statements into machine readable code.

    Args:
        code(str): The statements to assemble.
        addr(int): The memory address where the code will run.
        syntax(AsmSyntax): The input assembler syntax for x86. Defaults to
            nasm, ignored on other platforms.
        target(~pwnypack.target.Target): The target architecture. The
            global target is used if this argument is ``None``.
        gnu_binutils_prefix(str): When the syntax is AT&T, gnu binutils'
            as and ld will be used. By default, it selects
            ``arm-*-as/ld`` for 32bit ARM targets,
            ``aarch64-*-as/ld`` for 64 bit ARM targets,
            ``i386-*-as/ld`` for 32bit X86 targets and
            ``amd64-*-as/ld`` for 64bit X86 targets (all for various flavors
            of ``*``. This option allows you to pick a different toolchain.
            The prefix should always end with a '-' (or be empty).

    Returns:
        bytes: The assembled machine code.

    Raises:
        SyntaxError: If the assembler statements are invalid.
        NotImplementedError: In an unsupported target platform is specified.

    Example:
        >>> from pwny import *
        >>> asm('''
        ...     pop rdi
        ...     ret
        ... ''', target=Target(arch=Target.Arch.x86, bits=64))
        b'_\\xc3'

def asm(code, addr=0, syntax=None, target=None, gnu_binutils_prefix=None):
    """
    Assemble statements into machine readable code.

    Args:
        code(str): The statements to assemble.
        addr(int): The memory address where the code will run.
        syntax(AsmSyntax): The input assembler syntax for x86. Defaults to
            nasm, ignored on other platforms.
        target(~pwnypack.target.Target): The target architecture. The
            global target is used if this argument is ``None``.
        gnu_binutils_prefix(str): When the syntax is AT&T, gnu binutils'
            as and ld will be used. By default, it selects
            ``arm-*-as/ld`` for 32bit ARM targets,
            ``aarch64-*-as/ld`` for 64 bit ARM targets,
            ``i386-*-as/ld`` for 32bit X86 targets and
            ``amd64-*-as/ld`` for 64bit X86 targets (all for various flavors
            of ``*``. This option allows you to pick a different toolchain.
            The prefix should always end with a '-' (or be empty).

    Returns:
        bytes: The assembled machine code.

    Raises:
        SyntaxError: If the assembler statements are invalid.
        NotImplementedError: In an unsupported target platform is specified.

    Example:
        >>> from pwny import *
        >>> asm('''
        ...     pop rdi
        ...     ret
        ... ''', target=Target(arch=Target.Arch.x86, bits=64))
        b'_\\xc3'
    """

    if target is None:
        target = pwnypack.target.target

    if syntax is None and target.arch is pwnypack.target.Target.Arch.x86:
        syntax = AsmSyntax.nasm

    if HAVE_KEYSTONE and WANT_KEYSTONE:
        ks_mode = 0
        ks_syntax = None

        if target.arch is pwnypack.target.Target.Arch.x86:
            ks_arch = keystone.KS_ARCH_X86
            if target.bits is pwnypack.target.Target.Bits.bits_32:
                ks_mode |= keystone.KS_MODE_32
            else:
                ks_mode |= keystone.KS_MODE_64
            if syntax is AsmSyntax.nasm:
                ks_syntax = keystone.KS_OPT_SYNTAX_NASM
            elif syntax is AsmSyntax.intel:
                ks_syntax = keystone.KS_OPT_SYNTAX_INTEL
            else:
                ks_syntax = keystone.KS_OPT_SYNTAX_ATT

        elif target.arch is pwnypack.target.Target.Arch.arm:
            if target.bits is pwnypack.target.Target.Bits.bits_32:
                ks_arch = keystone.KS_ARCH_ARM

                if target.mode & pwnypack.target.Target.Mode.arm_thumb:
                    ks_mode |= keystone.KS_MODE_THUMB
                else:
                    ks_mode |= keystone.KS_MODE_ARM

                if target.mode & pwnypack.target.Target.Mode.arm_v8:
                    ks_mode |= keystone.KS_MODE_V8

                if target.mode & pwnypack.target.Target.Mode.arm_m_class:
                    ks_mode |= keystone.KS_MODE_MICRO

                if target.endian is pwnypack.target.Target.Endian.little:
                    ks_mode |= keystone.KS_MODE_LITTLE_ENDIAN
                else:
                    ks_mode |= keystone.KS_MODE_BIG_ENDIAN
            else:
                ks_arch = keystone.KS_ARCH_ARM64
                ks_mode |= keystone.KS_MODE_LITTLE_ENDIAN
        else:
            raise NotImplementedError('Unsupported syntax or target platform.')

        ks = keystone.Ks(ks_arch, ks_mode)
        if ks_syntax is not None:
            ks.syntax = ks_syntax
        try:
            data, insn_count = ks.asm(code, addr)
        except keystone.KsError as e:
            import traceback
            traceback.print_exc()
            raise SyntaxError(e.message)
        return b''.join(six.int2byte(b) for b in data)

    if target.arch is pwnypack.target.Target.Arch.x86 and syntax is AsmSyntax.nasm:
        with tempfile.NamedTemporaryFile() as tmp_asm:
            tmp_asm.write(('bits %d\norg %d\n%s' % (target.bits.value, addr, code)).encode('utf-8'))
            tmp_asm.flush()

            tmp_bin_fd, tmp_bin_name = tempfile.mkstemp()
            os.close(tmp_bin_fd)

            try:
                p = subprocess.Popen(
                    [
                        'nasm',
                        '-o', tmp_bin_name,
                        '-f', 'bin',
                        tmp_asm.name,
                    ],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                )
                stdout, stderr = p.communicate()

                if p.returncode:
                    raise SyntaxError(stderr.decode('utf-8'))

                tmp_bin = open(tmp_bin_name, 'rb')
                result = tmp_bin.read()
                tmp_bin.close()
                return result
            finally:
                try:
                    os.unlink(tmp_bin_name)
                except OSError:
                    pass
    elif target.arch in (pwnypack.target.Target.Arch.x86, pwnypack.target.Target.Arch.arm):
        preamble = ''
        as_flags = []
        ld_flags = []

        if target.arch is pwnypack.target.Target.Arch.x86:
            if target.bits == 32:
                binutils_arch = 'i386'
            else:
                binutils_arch = 'amd64'

            if syntax is AsmSyntax.intel:
                preamble = '.intel_syntax noprefix\n'

            ld_flags.extend(['--oformat', 'binary'])
        else:
            if target.bits == 32:
                binutils_arch = 'arm'
                if target.mode & pwnypack.target.Target.Mode.arm_v8:
                    as_flags.append('-march=armv8-a')
                elif target.mode & pwnypack.target.Target.Mode.arm_m_class:
                    as_flags.append('-march=armv7m')
            else:
                binutils_arch = 'aarch64'

            if target.endian is pwnypack.target.Target.Endian.little:
                as_flags.append('-mlittle-endian')
                ld_flags.append('-EL')
            else:
                as_flags.append('-mbig-endian')
                ld_flags.append('-EB')

            if target.mode & pwnypack.target.Target.Mode.arm_thumb:
                as_flags.append('-mthumb')

        if gnu_binutils_prefix is None:
            gnu_binutils_prefix = find_binutils_prefix(binutils_arch)

        tmp_out_fd, tmp_out_name = tempfile.mkstemp()
        try:
            os.close(tmp_out_fd)

            p = subprocess.Popen(
                [
                    '%sas' % gnu_binutils_prefix,
                    '-o', tmp_out_name
                ] + as_flags,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )
            stdout, stderr = p.communicate((preamble + code).encode('utf-8'))

            if p.returncode:
                raise SyntaxError(stderr.decode('utf-8'))

            tmp_bin_fd, tmp_bin_name = tempfile.mkstemp()
            try:
                os.close(tmp_bin_fd)

                p = subprocess.Popen(
                    [
                        '%sld' % gnu_binutils_prefix,
                        '-Ttext', str(addr),
                    ] + ld_flags + [
                        '-o', tmp_bin_name,
                        tmp_out_name,
                    ],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                )
                stdout, stderr = p.communicate()

                if p.returncode:
                    raise SyntaxError(stderr.decode('utf-8'))

                if 'binary' in ld_flags:
                    tmp_bin = open(tmp_bin_name, 'rb')
                    result = tmp_bin.read()
                    tmp_bin.close()
                    return result
                else:
                    tmp_bin = ELF(tmp_bin_name)
                    return tmp_bin.get_section_header('.text').content
            finally:
                try:
                    os.unlink(tmp_bin_name)
                except OSError:
                    pass
        finally:
            try:
                os.unlink(tmp_out_name)
            except OSError:
                pass  # pragma: no cover

    else:
        raise NotImplementedError('Unsupported syntax or target platform.')