#! /usr/bin/env python3

from pwn import *

ADD = 0
ADDI = 1
SUB = 2
COPY = 3
LOADI = 4

A = 0
B = 1
C = 2
D = 3
E = 4
F = 5
G = 6
H = 7
I = 8
J = 9
K = 10
L = 11
M = 12
N = 13

INSTR_LEN = 8


def instr_i(opcode, reg1, imm: int):
    assert (opcode == ADDI or opcode == LOADI)
    return bytes([opcode, reg1, 0, 0]) + imm.to_bytes(4, byteorder='little')


def instr_r(opcode, reg1, reg2):
    assert (opcode == ADD or opcode == SUB or opcode == COPY)
    return bytes([opcode, reg1, 0, 0, reg2, 0, 0, 0])

context.log_level = 'debug'
with remote("localhost", 1337, fam="ipv4") as p:
    p.recvuntil(b"Password: ")
    p.sendline(b"1234")

    program = instr_i(LOADI, A, 0x50)

    print(p.recvuntil(b"COPaaS - Compiler-oriented programming as a service\n").decode())
    print(p.recvuntil(b"? (y/N):").decode())
    p.sendline(b"N")
    print(p.recvuntil(b"Using the demo version!").decode())
    print(p.recvuntil(b"should it bee?").decode())
    len_msg = str(len(program) // INSTR_LEN).encode()
    log.info(f"Sending: {len_msg}")
    p.sendline(len_msg)
    print(p.recvuntil(b"Now your program:").decode())
    log.info(f"Sending program: {list(program)}")
    pause()
    p.send(program)
    p.interactive()
    print(p.recvuntil(b"Your program exited with "))
    exit_code = int(p.recvuntil(b"!", drop=True))
    log.info(f"Retrieved exit code {exit_code}!")
    p.interactive()