"""
This module contains functions to assemble and disassemble code for a given
target platform. By default the keystone engine assembler will be used if it
is available. If it's not available (or if the ``WANT_KEYSTONE`` environment
variable is set and it's not ``1``, ``YES`` or ``TRUE`` (case insensitive)),
pwnypack falls back to using the *nasm* assembler for nasm syntax on X86 or
*GNU as* for any other supported syntax / architecture. Disassembly is
performed by *ndisasm* on x86 for nasm syntax. *capstone* is used for any
other supported syntax / architecture.
Currently, the only supported architectures are
:attr:`~pwnypack.target.Target.Arch.x86` (both 32 and 64 bits variants) and
:attr:`~pwnypack.target.Target.Arch.arm` (both 32 and 64 bits variants).
"""
from __future__ import print_function
try:
import shutilwhich
except ImportError:
pass
import argparse
import os
import subprocess
import sys
from enum import IntEnum
import shutil
from pwnypack.elf import ELF
import pwnypack.target
import pwnypack.main
import pwnypack.codec
import tempfile
import six
try:
import capstone
HAVE_CAPSTONE = True
except ImportError:
HAVE_CAPSTONE = False
try:
import keystone
HAVE_KEYSTONE = True
except ImportError:
HAVE_KEYSTONE = False
WANT_KEYSTONE = os.environ.get('WANT_KEYSTONE', '1').upper() in ('1', 'YES', 'TRUE')
__all__ = [
'AsmSyntax',
'asm',
'disasm',
]
BINUTILS_SUFFIXES = [
'none-eabi-',
'unknown-linux-gnu-',
'linux-gnu-',
'linux-gnueabi-',
]
BINUTILS_PREFIXES = {}
def find_binutils_prefix(arch):
global BINUTILS_PREFIXES
prefix = BINUTILS_PREFIXES.get(arch)
if prefix is not None:
return prefix
for suffix in BINUTILS_SUFFIXES:
prefix = '%s-%s' % (arch, suffix)
if shutil.which('%sas' % prefix) and \
shutil.which('%sld' % prefix):
BINUTILS_PREFIXES[arch] = prefix
return prefix
else:
raise RuntimeError('Could not locate a suitable binutils for %s.' % arch)
class AsmSyntax(IntEnum):
"""
This enumeration is used to specify the assembler syntax.
"""
nasm = 0 #: Netwide assembler syntax
intel = 1 #: Intel assembler syntax
att = 2 #: AT&T assembler syntax
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.')
def prepare_capstone(syntax=AsmSyntax.att, target=None):
"""
Prepare a capstone disassembler instance for a given target and syntax.
Args:
syntax(AsmSyntax): The assembler syntax (Intel or AT&T).
target(~pwnypack.target.Target): The target to create a disassembler
instance for. The global target is used if this argument is
``None``.
Returns:
An instance of the capstone disassembler.
Raises:
NotImplementedError: If the specified target isn't supported.
"""
if not HAVE_CAPSTONE:
raise NotImplementedError('pwnypack requires capstone to disassemble to AT&T and Intel syntax')
if target is None:
target = pwnypack.target.target
if target.arch == pwnypack.target.Target.Arch.x86:
if target.bits is pwnypack.target.Target.Bits.bits_32:
md = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_32)
else:
md = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)
elif target.arch == pwnypack.target.Target.Arch.arm:
mode = 0
if target.bits is pwnypack.target.Target.Bits.bits_32:
arch = capstone.CS_ARCH_ARM
if target.mode and pwnypack.target.Target.Mode.arm_thumb:
mode = capstone.CS_MODE_THUMB
else:
mode = capstone.CS_MODE_ARM
if target.mode and pwnypack.target.Target.Mode.arm_m_class:
mode |= capstone.CS_MODE_MCLASS
if target.mode and pwnypack.target.Target.Mode.arm_v8:
mode |= capstone.CS_MODE_V8
else:
arch = capstone.CS_ARCH_ARM64
if target.endian is pwnypack.target.Target.Endian.little:
mode |= capstone.CS_MODE_LITTLE_ENDIAN
else:
mode |= capstone.CS_MODE_BIG_ENDIAN
md = capstone.Cs(arch, mode)
else:
raise NotImplementedError('Only x86 is currently supported.')
md.skipdata = True
if syntax is AsmSyntax.att:
md.syntax = capstone.CS_OPT_SYNTAX_ATT
elif syntax is AsmSyntax.intel:
md.skipdata_setup(('db', None, None))
else:
raise NotImplementedError('capstone engine only implements AT&T and Intel syntax.')
return md
def disasm(code, addr=0, syntax=None, target=None):
"""
Disassemble machine readable code into human readable statements.
Args:
code(bytes): The machine code that is to be disassembled.
addr(int): The memory address of the code (used for relative
references).
syntax(AsmSyntax): The output assembler syntax. This defaults to
nasm on x86 architectures, AT&T on all other architectures.
target(~pwnypack.target.Target): The architecture for which the code
was written. The global target is used if this argument is
``None``.
Returns:
list of str: The disassembled machine code.
Raises:
NotImplementedError: In an unsupported target platform is specified.
RuntimeError: If ndisasm encounters an error.
Example:
>>> from pwny import *
>>> disasm(b'_\\xc3', target=Target(arch=Target.Arch.x86, bits=64))
['pop rdi', 'ret']
"""
if target is None:
target = pwnypack.target.target
if syntax is None:
if target.arch is pwnypack.target.Target.Arch.x86:
syntax = AsmSyntax.nasm
else:
syntax = AsmSyntax.att
if syntax is AsmSyntax.nasm:
if target.arch is not pwnypack.target.Target.Arch.x86:
raise NotImplementedError('nasm only supports x86.')
p = subprocess.Popen(
[
'ndisasm',
'-b',
str(target.bits.value),
'-o',
str(addr),
'-',
],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
stdout, stderr = p.communicate(code)
if p.returncode:
raise RuntimeError(stderr.decode('utf-8'))
return [
line.split(None, 2)[2]
for line in stdout.decode('utf-8').split('\n')
if line and not line.startswith(' ')
]
elif syntax in (AsmSyntax.intel, AsmSyntax.att):
md = prepare_capstone(syntax, target)
statements = []
total_size = 0
for (_, size, mnemonic, op_str) in md.disasm_lite(code, addr):
statements.append((mnemonic + ' ' + op_str).strip())
total_size += size
return statements
else:
raise NotImplementedError('Unsupported syntax for host platform.')
@pwnypack.main.register('asm')
def asm_app(parser, cmd, args): # pragma: no cover
"""
Assemble code from commandline or stdin.
Please not that all semi-colons are replaced with carriage returns
unless source is read from stdin.
"""
parser.add_argument('source', help='the code to assemble, read from stdin if omitted', nargs='?')
pwnypack.main.add_target_arguments(parser)
parser.add_argument(
'--syntax', '-s',
choices=AsmSyntax.__members__.keys(),
default=None,
)
parser.add_argument(
'--address', '-o',
type=lambda v: int(v, 0),
default=0,
help='the address where the code is expected to run',
)
args = parser.parse_args(args)
target = pwnypack.main.target_from_arguments(args)
if args.syntax is not None:
syntax = AsmSyntax.__members__[args.syntax]
else:
syntax = None
if args.source is None:
args.source = sys.stdin.read()
else:
args.source = args.source.replace(';', '\n')
return asm(
args.source,
syntax=syntax,
target=target,
addr=args.address,
)
@pwnypack.main.register('disasm')
def disasm_app(_parser, cmd, args): # pragma: no cover
"""
Disassemble code from commandline or stdin.
"""
parser = argparse.ArgumentParser(
prog=_parser.prog,
description=_parser.description,
)
parser.add_argument('code', help='the code to disassemble, read from stdin if omitted', nargs='?')
pwnypack.main.add_target_arguments(parser)
parser.add_argument(
'--syntax', '-s',
choices=AsmSyntax.__members__.keys(),
default=None,
)
parser.add_argument(
'--address', '-o',
type=lambda v: int(v, 0),
default=0,
help='the address of the disassembled code',
)
parser.add_argument(
'--format', '-f',
choices=['hex', 'bin'],
help='the input format (defaults to hex for commandline, bin for stdin)',
)
args = parser.parse_args(args)
target = pwnypack.main.target_from_arguments(args)
if args.syntax is not None:
syntax = AsmSyntax.__members__[args.syntax]
else:
syntax = None
if args.format is None:
if args.code is None:
args.format = 'bin'
else:
args.format = 'hex'
if args.format == 'hex':
code = pwnypack.codec.dehex(pwnypack.main.string_value_or_stdin(args.code))
else:
code = pwnypack.main.binary_value_or_stdin(args.code)
print('\n'.join(disasm(code, args.address, syntax=syntax, target=target)))
@pwnypack.main.register(name='symbol-disasm')
def disasm_symbol_app(_parser, _, args): # pragma: no cover
"""
Disassemble a symbol from an ELF file.
"""
parser = argparse.ArgumentParser(
prog=_parser.prog,
description=_parser.description,
)
parser.add_argument(
'--syntax', '-s',
choices=AsmSyntax.__members__.keys(),
default=None,
)
parser.add_argument('file', help='ELF file to extract a symbol from')
parser.add_argument('symbol', help='the symbol to disassemble')
args = parser.parse_args(args)
if args.syntax is not None:
syntax = AsmSyntax.__members__[args.syntax]
else:
syntax = None
elf = ELF(args.file)
symbol = elf.get_symbol(args.symbol)
print('\n'.join(disasm(symbol.content, symbol.value, syntax=syntax, target=elf)))
