pwndbg/pwndbg/arguments.py

"""
Allows describing functions, specifically enumerating arguments which
may be passed in a combination of registers and stack values.
"""
import gdb
from capstone import CS_GRP_CALL
from capstone import CS_GRP_INT

import pwndbg.abi
import pwndbg.arch
import pwndbg.chain
import pwndbg.color.nearpc as N
import pwndbg.constants
import pwndbg.disasm
import pwndbg.funcparser
import pwndbg.functions
import pwndbg.ida
import pwndbg.memory
import pwndbg.regs
import pwndbg.symbol
import pwndbg.typeinfo

ida_replacements = {
    "__int64": "signed long long int",
    "__int32": "signed int",
    "__int16": "signed short",
    "__int8": "signed char",
    "__uint64": "unsigned long long int",
    "__uint32": "unsigned int",
    "__uint16": "unsigned short",
    "__uint8": "unsigned char",
    "_BOOL_1": "unsigned char",
    "_BOOL_2": "unsigned short",
    "_BOOL_4": "unsigned int",
    "_BYTE": "unsigned char",
    "_WORD": "unsigned short",
    "_DWORD": "unsigned int",
    "_QWORD": "unsigned long long",
    "__pure": "",
    "__hidden": "",
    "__return_ptr": "",
    "__struct_ptr": "",
    "__array_ptr": "",
    "__fastcall": "",
    "__cdecl": "",
    "__thiscall": "",
    "__userpurge": "",
}


def get_syscall_name(instruction):
    if CS_GRP_INT not in instruction.groups:
        return None

    syscall_register = pwndbg.abi.ABI.syscall().syscall_register

    # If we are on x86/x64, return no syscall name for other instructions than syscall and int 0x80
    if syscall_register in ("eax", "rax"):
        mnemonic = instruction.mnemonic
        if not (mnemonic == "syscall" or (mnemonic == "int" and instruction.op_str == "0x80")):
            return None

    syscall_number = getattr(pwndbg.regs, syscall_register)
    return pwndbg.constants.syscall(syscall_number) or "<unk_%d>" % syscall_number


def get(instruction):
    """
    Returns an array containing the arguments to the current function,
    if $pc is a 'call' or 'bl' type instruction.

    Otherwise, returns None.
    """
    n_args_default = 4

    if instruction is None:
        return []

    if instruction.address != pwndbg.regs.pc:
        return []

    if CS_GRP_CALL in instruction.groups:
        try:
            abi = pwndbg.abi.ABI.default()
        except KeyError:
            return []

        # Not sure of any OS which allows multiple operands on
        # a call instruction.
        assert len(instruction.operands) == 1

        target = instruction.operands[0].int

        if not target:
            return []

        name = pwndbg.symbol.get(target)
        if not name:
            return []
    elif CS_GRP_INT in instruction.groups:
        # Get the syscall number and name
        name = get_syscall_name(instruction)
        abi = pwndbg.abi.ABI.syscall()
        target = None

        if name is None:
            return []
    else:
        return []

    result = []
    name = name or ""

    sym = gdb.lookup_symbol(name)
    name = name.replace("isoc99_", "")  # __isoc99_sscanf
    name = name.replace("@plt", "")  # getpwiod@plt

    # If we have particular `XXX_chk` function in our database, we use it.
    # Otherwise, we show args for its unchecked version.
    # We also lstrip `_` in here, as e.g. `__printf_chk` needs the underscores.
    if name not in pwndbg.functions.functions:
        name = name.replace("_chk", "")
        name = name.strip().lstrip("_")  # _malloc

    func = pwndbg.functions.functions.get(name, None)

    # Try to extract the data from GDB.
    # Note that this is currently broken, pending acceptance of
    # my patch: https://sourceware.org/ml/gdb-patches/2015-06/msg00268.html
    if sym and sym[0]:
        try:
            n_args_default = len(sym[0].type.fields())
        except TypeError:
            pass

    # Try to grab the data out of IDA
    if not func and target:
        typename = pwndbg.ida.GetType(target)

        if typename:
            typename += ";"

            # GetType() does not include the name.
            typename = typename.replace("(", " function_name(", 1)

            for k, v in ida_replacements.items():
                typename = typename.replace(k, v)

            func = pwndbg.funcparser.ExtractFuncDeclFromSource(typename + ";")

    if func:
        args = func.args
    else:
        args = (pwndbg.functions.Argument("int", 0, argname(i, abi)) for i in range(n_args_default))

    for i, arg in enumerate(args):
        result.append((arg, argument(i, abi)))

    return result


def argname(n, abi=None):
    abi = abi or pwndbg.abi.ABI.default()
    regs = abi.register_arguments

    if n < len(regs):
        return regs[n]

    return "arg[%i]" % n


def argument(n, abi=None):
    """
    Returns the nth argument, as if $pc were a 'call' or 'bl' type
    instruction.
    Works only for ABIs that use registers for arguments.
    """
    abi = abi or pwndbg.abi.ABI.default()
    regs = abi.register_arguments

    if n < len(regs):
        return getattr(pwndbg.regs, regs[n])

    n -= len(regs)

    sp = pwndbg.regs.sp + (n * pwndbg.arch.ptrsize)

    return int(pwndbg.memory.poi(pwndbg.typeinfo.ppvoid, sp))


def arguments(abi=None):
    """
    Yields (arg_name, arg_value) tuples for arguments from a given ABI.
    Works only for ABIs that use registers for arguments.
    """
    abi = abi or pwndbg.abi.ABI.default()
    regs = abi.register_arguments

    for i in range(len(regs)):
        yield argname(i, abi), argument(i, abi)


def format_args(instruction):
    result = []
    for arg, value in get(instruction):
        code = arg.type != "char"
        pretty = pwndbg.chain.format(value, code=code)

        # Enhance args display
        if arg.name == "fd" and isinstance(value, int):
            path = pwndbg.file.readlink("/proc/%d/fd/%d" % (pwndbg.proc.pid, value))
            if path:
                pretty += " (%s)" % path

        result.append("%-10s %s" % (N.argument(arg.name) + ":", pretty))
    return result