exploit_template.py

#!/usr/bin/env python2

from pwn import *

context.update(arch="amd64", # arch="i386", arch="mips", arch="arm",
               endian="little", os="linux", log_level="debug",
               terminal=["tmux", "split-window", "-v", "-p 85"],)

LOCAL, REMOTE, SSH = False, False, False
TARGET=os.path.realpath("./bof")
elf = ELF(TARGET)

def attach(r):
    if LOCAL:
        bkps = [elf.symbols["main"], ]
        cmds = []
        gdb.attach(r, '\n'.join(["break *{:#x}".format(x) for x in bkps] + cmds))
    return

def exploit(r):
    attach(r)
    r.interactive()
    return

if __name__ == "__main__":
    if len(sys.argv)==6 and sys.argv[1]=="ssh":
        SSH = True
        sh = ssh(host=sys.argv[2], port=sys.argv[3], user=sys.argv[4], password=sys.argv[5])
        sh.set_working_directory(os.path.dirname(TARGET))
        r = sh.process([TARGET, ])
    elif len(sys.argv)==3:
        REMOTE = True
        r = remote(sys.argv[1], int(sys.argv[2]))
    else:
        LOCAL = True
        r = process([TARGET, ])
    exploit(r)
    sys.exit(0)