Synacor challenge - Implemented a VM to run the binary

virtual_machine.py

from enum import IntEnum
REGISTER_COUNT = 8

class Ops(IntEnum):
    HALT = 0
    SET  = 1
    PUSH = 2
    POP  = 3
    EQ   = 4
    GT   = 5
    JMP  = 6
    JT   = 7
    JF   = 8
    ADD  = 9
    MULT = 10
    MOD  = 11
    AND  = 12
    OR   = 13
    NOT  = 14
    RMEM = 15
    WMEM = 16
    CALL = 17
    RET  = 18
    OUT  = 19
    IN   = 20
    NOOP = 21


class VirtualMachine:
    def __init__(self, memory, registers=None, stack=None):
        if registers == None:
            registers = [0 for _ in range(REGISTER_COUNT)]
        if stack == None:
            stack = []

        self.memory = memory
        self.registers = registers
        self.stack = stack
        self.p = 0

    def get_op(self, p):
        return self.memory[p]

    def get_ascii_val_from_mem(self, p):
        return chr(self.get_val_from_mem(p))

    def get_val_from_mem(self, p):
        val = self.memory[p]
        if val >= 32768:
            val = self.registers[val%32768]
        return val

    def set_register(self, i, value):
        self.registers[i] = value

    def get_register(self, p):
        register =  self.memory[p]%32768
        return register

    def stack_push(self, val):
        self.stack.append(val)

    def stack_pop(self):
        if not len(self.stack):
            raise Exception("Stack is empty")
        return self.stack.pop()

    def get_register_value(self, i):
        return self.registers[i]

    def execute_program(self, program):
        inp = ''
        while(True):
            op = self.get_op(self.p)
            if op == Ops.HALT:
                print("Halting the machine...")
                break
            elif op == Ops.SET:
                a = self.get_register(self.p+1)
                b = self.get_val_from_mem(self.p+2)
                self.set_register(a, b)
                self.p += 3
            elif op == Ops.PUSH:
                a = self.get_val_from_mem(self.p+1)
                self.stack_push(a)
                self.p += 2
            elif op == Ops.POP:
                a = self.get_register(self.p+1)
                val = self.stack_pop()
                self.set_register(a, val)
                self.p += 2
            elif op == Ops.EQ:
                a = self.get_register(self.p+1)
                b = self.get_val_from_mem(self.p+2)
                c = self.get_val_from_mem(self.p+3)
                val = 0
                if b == c:
                    val = 1
                self.set_register(a, val)
                self.p += 4   
            elif op == Ops.GT:
                a = self.get_register(self.p+1)
                b = self.get_val_from_mem(self.p+2)
                c = self.get_val_from_mem(self.p+3)
                val = 0
                if b > c:
                    val = 1
                self.set_register(a, val)
                self.p += 4
            elif op == Ops.JMP:
                self.p = self.get_val_from_mem(self.p+1)
            elif op == Ops.JT:
                a = self.get_val_from_mem(self.p+1)
                b = self.get_val_from_mem(self.p+2)
                if a:
                    self.p = b
                else:
                    self.p += 3
            elif op == Ops.JF:
                a = self.get_val_from_mem(self.p+1)
                b = self.get_val_from_mem(self.p+2)
                if not a:
                    self.p = b
                else:
                    self.p += 3
            elif op == Ops.ADD:
                a = self.get_register(self.p+1)
                b = self.get_val_from_mem(self.p+2)
                c = self.get_val_from_mem(self.p+3)
                sum = (b + c)%32768
                self.set_register(a, sum)
                self.p += 4
            elif op == Ops.MULT:
                a = self.get_register(self.p+1)
                b = self.get_val_from_mem(self.p+2)
                c = self.get_val_from_mem(self.p+3)
                prod = (b * c)%32768
                self.set_register(a, prod)
                self.p += 4
            elif op == Ops.MOD:
                a = self.get_register(self.p+1)
                b = self.get_val_from_mem(self.p+2)
                c = self.get_val_from_mem(self.p+3)
                mod = b % c
                self.set_register(a, mod)
                self.p += 4
            elif op == Ops.AND:
                a = self.get_register(self.p+1)
                b = self.get_val_from_mem(self.p+2)
                c = self.get_val_from_mem(self.p+3)
                val = b & c
                self.set_register(a, val)
                self.p += 4
            elif op == Ops.OR:
                a = self.get_register(self.p+1)
                b = self.get_val_from_mem(self.p+2)
                c = self.get_val_from_mem(self.p+3)
                val = b | c
                self.set_register(a, val)
                self.p += 4
            elif op == Ops.NOT:
                a = self.get_register(self.p+1)
                b = self.get_val_from_mem(self.p+2)
                val = 2**15 - b - 1
                self.set_register(a, val)
                self.p += 3
            elif op == Ops.RMEM:
                a = self.get_register(self.p+1)
                b = self.get_val_from_mem(self.p+2)
                c = self.get_val_from_mem(b)
                self.set_register(a, c)
                self.p += 3
            elif op == Ops.WMEM:
                a = self.get_val_from_mem(self.p+1)
                b = self.get_val_from_mem(self.p+2)
                self.memory[a] = b
                self.p += 3
            elif op == Ops.CALL:
                self.stack_push(self.p+2)
                self.p = self.get_val_from_mem(self.p+1)
            elif op == Ops.RET:
                val = self.stack_pop()
                self.p = val
            elif op == Ops.OUT:
                print(f"{self.get_ascii_val_from_mem(self.p+1)}", end="")
                self.p += 2
            elif op == Ops.IN:
                if not len(inp):
                    inp = input() + "\n"
                ch = ord(inp[0])
                inp = inp[1:]
                a = self.get_register(self.p+1)
                self.set_register(a, ch)
                self.p += 2
            elif op == Ops.NOOP:
                self.p += 1
            else:
                print(f"Op not implemented {op}")
                break


def little_endian(a, b):
    le_bits = bin(b)[2:].zfill(8) + bin(a)[2:].zfill(8)
    return int(le_bits, 2)


def read_program(file_name):
    f = open(file_name, "rb")
    l = list(bytearray(f.read()))
    for i, ele in enumerate(l):
        if not i%2:
            l[i] = little_endian(ele, l[i+1])
    del l[1::2]
    return l


def get_val(memory, pointer):
    val = memory[pointer]


if __name__ == "__main__":
    file_name = "challenge.bin"
    program = read_program(file_name)
    vm = VirtualMachine(program)
    vm.execute_program(program)