gipu/assembler/assembler.py

import sys
import re
import struct
import IPython
import copy


class AssemblerException(Exception):
    pass


class InvalidRegister(AssemblerException):

    def __init__(self, register):
        super().__init__("Invalid register: {}".format(register))


class InvalidOperation(AssemblerException):

    def __init__(self, operation):
        super().__init__("Invalid operation: {}".format(operation))


class ExpectedImmediate(AssemblerException):

    def __init__(self, value):
        super().__init__("Expected immediate, got {}".format(value))


class ExpectedRegister(AssemblerException):

    def __init__(self, value):
        super().__init__("Expected register, got {}".format(value))


class IPOverwrite(AssemblerException):

    def __init__(self, instruction=None):
        if instruction:
            super().__init__("IP can't be overwritten. Instruction: {}".format(instruction))
        else:
            super().__init__("IP can't be overwritten.")


class InvalidValue(AssemblerException):

    def __init__(self, instruction):
        super().__init__("Invalid value while assembling: {}".format(instruction))


class VMAssembler:

    def __init__(self, key):
        self.assembled_code = bytearray()
        self.encrypt_ops(key)

    def parse(self, instruction):
        action = getattr(self, "{}".format(instruction.opcode.method))
        action(instruction)

    def process_code_line(self, line):
        components = [x for x in re.split('\W', line) if x]
        instruction = VMInstruction(components[0], components[1:])
        sys.stdout.write(str(instruction) + "\n")
        self.parse(instruction)

    def imm2reg(self, instruction):
        """
        Intel syntax -> REG, IMM
        """
        opcode = instruction.opcode
        reg = instruction.args[0]
        imm = instruction.args[1]
        if reg.name == "ip":
            raise IPOverwrite(instruction)
        if not imm.isimm():
            raise ExpectedImmediate(imm)
        if not reg.isreg():
            raise ExpectedRegister(reg)
        if not opcode.uint8() or not reg.uint8() or not imm.uint16():
            raise InvalidValue(instruction)
        self.assembled_code += opcode.uint8() + reg.uint8() + imm.uint16()
        return

    def reg2reg(self, instruction):
        """
        Intel syntax -> DST_REG, SRC_REG
        """
        opcode = instruction.opcode
        dst_reg = instruction.args[0]
        src_reg = instruction.args[1]
        if dst_reg.name == "ip" or src_reg.name == "ip":
            raise IPOverwrite(instruction)
        if not dst_reg.isreg():
            raise ExpectedRegister(dst_reg)
        if not src_reg.isreg():
            raise ExpectedRegister(src_reg)
        if not opcode.uint8() or not dst_reg.uint8() or not src_reg.uint8():
            raise InvalidValue(instruction)
        byte_with_nibbles = struct.pack("<B", dst_reg.uint8()[0] << 4 ^ (
            src_reg.uint8()[0] & 0b00001111))
        self.assembled_code += opcode.uint8() + byte_with_nibbles
        return

    def reg2imm(self, instruction):
        """
        Intel syntax -> IMM, REG
        """
        opcode = instruction.opcode
        imm = instruction.args[0]
        reg = instruction.args[1]
        if reg.name == "ip":
            raise IPOverwrite(instruction)
        if not imm.isimm():
            raise ExpectedImmediate(imm)
        if not reg.isreg():
            raise ExpectedRegister(reg)
        if not opcode.uint8() or not reg.uint8() or not imm.uint16():
            raise InvalidValue(instruction)
        self.assembled_code += opcode.uint8() + imm.uint16() + reg.uint8()
        return

    def byt2reg(self, instruction):
        """
        Intel syntax -> REG, [BYTE]IMM
        """
        opcode = instruction.opcode
        reg = instruction.args[0]
        imm = instruction.args[1]
        if reg.name == "ip":
            raise IPOverwrite(instruction)
        if not imm.isimm():
            raise ExpectedImmediate(imm)
        if not reg.isreg():
            raise ExpectedRegister(reg)
        if not opcode.uint8() or not reg.uint8() or not imm.uint8():
            raise InvalidValue(instruction)
        self.assembled_code += opcode.uint8() + reg.uint8() + imm.uint8()
        return

    def regonly(self, instruction):
        """
        Instruction with only an argument: a register
        """
        opcode = instruction.opcode
        reg = instruction.args[0]
        if reg.name == "ip":
            raise IPOverwrite(instruction)
        if not reg.isreg():
            raise ExpectedRegister(reg)
        if not opcode.uint8() or not reg.uint8():
            raise InvalidValue(instruction)
        self.assembled_code += opcode.uint8() + reg.uint8()
        return

    def immonly(self, instruction):
        """
        Instruction with only an argument: an immediate
        """
        opcode = instruction.opcode
        imm = instruction.args[0]
        if not imm.isimm():
            raise ExpectedImmediate(imm)
        if not opcode.uint8() or not imm.uint16():
            raise InvalidValue(instruction)
        self.assembled_code += opcode.uint8() + imm.uint16()
        return

    def jump(self, instruction):
        imm_op_re = re.compile(".*[iI]$")
        reg_op_re = re.compile(".*[rR]$")
        arg = instruction.args[0]
        section = next((x for x in functions if x.name == arg.name), None)
        # TODO this is due the VMComponent structure
        instruction.args[0].name = section.offset
        instruction.args[0].value = section.offset
        if imm_op_re.match(instruction.opcode.name):
            self.immonly(instruction)
        elif reg_op_re.match(instruction.opcode.name):
            self.regonly(instruction)
        else:
            raise AssemblerException()

    def single(self, instruction):
        """
        Instruction with no arguments
        """
        opcode = instruction.opcode
        self.assembled_code += opcode.uint8()
        return

    def encrypt_ops(self, key):
        key_ba = bytearray(key, 'utf-8')
        olds = copy.deepcopy(ops)

        # RC4 KSA! :-P
        arr = [i for i in range(256)]
        j = 0
        for i in range(len(arr)):
            j = (j + arr[i] + key_ba[i % len(key)]) % len(arr)
            arr[i], arr[j] = arr[j], arr[i]

        for i, o in enumerate(ops):
            o.set_value(arr[i])

        for o, n in zip(olds, ops):
            print("{} : {}->{}".format(o.name, hex(o.value), hex(n.value)))


class VMComponent:
    """
    Represents a register, operation or an immediate the VM recognizes
    """

    def __init__(self, name, value, method=None):
        self.name = name.casefold()
        self.value = value
        self.method = method

    def __repr__(self):
        return "{}".format(self.name)

    def set_name(self, name):
        self.name = name

    def set_value(self, value):
        self.value = value

    def uint8(self):
        numre = re.compile("^[0-9]+$")
        if isinstance(self.value, int):
            return struct.pack("<B", self.value)
        elif self.value.startswith("0x"):
            return struct.pack("<B", int(self.value, 16))
        elif numre.match(self.value):  # only numbers
            return struct.pack("<B", int(self.value))
        return None

    def uint16(self):
        numre = re.compile("^[0-9]+$")
        if isinstance(self.value, int):
            return struct.pack("<H", self.value)
        elif self.value.startswith("0x"):
            return struct.pack("<H", int(self.value, 16))
        elif numre.match(self.value):  # only numbers
            return struct.pack("<H", int(self.value))
        return None

    def isreg(self):
        if self.name not in [x.casefold() for x in reg_names]:
            return False
        return True

    def isop(self):
        if self.name not in [x[0].casefold() for x in op_names]:
            return False
        return True

    def isimm(self):
        if not immediate_re.match(str(self.name)):
            return False
        return True


class VMInstruction:
    """
    Represents an instruction the VM recognizes.
    e.g: MOVI [R0, 2]
          ^       ^
        opcode  args
    """

    def __init__(self, opcode, instr_list):
        self.opcode = None
        self.args = None
        self.size = 1

        self.opcode = next((x for x in ops if x.name == opcode), None)
        if self.opcode == None:
            raise InvalidOperation(opcode)
        self.args = []
        for el in instr_list:
            if immediate_re.match(el):
                # directly append the immediate
                self.args.append(VMComponent(el, el))
                self.size += 2
                continue
            elif register_re.match(el):
                # create a VM component for a register
                reg_comp = next((x for x in regs if x.name == el), None)
                self.args.append(reg_comp)
                self.size += 1
                continue
            else:
                # section
                print(el)
                sec_comp = next((x for x in functions if x.name == el), None)
                if sec_comp:
                    self.args.append(VMComponent(
                        sec_comp.name, sec_comp.offset))
                    self.size += 2
                    continue
                raise AssemblerException()

    def __repr__(self):
        return "{} {}".format(self.opcode.name, ", ".join([x.name for x in self.args]))


class VMFunction:
    """
    Represents a code section or "label" such as "main:"
    """

    def __init__(self, name, line_start):
        self.name = name
        self.size = 0
        self.offset = 0
        self.line_start = line_start
        self.line_end = 0
        self.labels = []

    def set_size(self, size):
        self.size = size

    def set_offset(self, offset):
        self.offset = offset
    
    def set_line_start(self, start):
        self.line_start = start
    
    def set_line_end(self, end):
        self.line_end = end
    

    def __repr__(self):
        return "{} | ls: {}, le: {}, s: {}, o: {}".format(self.name, hex(self.line_start), hex(self.line_end),  hex(self.size), hex(self.offset))

op_names = [["MOVI", "imm2reg"],
            ["MOVR", "reg2reg"],
            ["LOAD", "imm2reg"],
            ["STOR", "reg2imm"],
            ["ADDI", "imm2reg"],
            ["ADDR", "reg2reg"],
            ["SUBI", "imm2reg"],
            ["SUBR", "reg2reg"],
            ["ANDB", "byt2reg"],
            ["ANDW", "imm2reg"],
            ["ANDR", "reg2reg"],
            ["YORB", "byt2reg"],
            ["YORW", "imm2reg"],
            ["YORR", "reg2reg"],
            ["XORB", "byt2reg"],
            ["XORW", "imm2reg"],
            ["XORR", "reg2reg"],
            ["NOTR", "regonly"],
            ["MULI", "imm2reg"],
            ["MULR", "reg2reg"],
            ["DIVI", "imm2reg"],
            ["DIVR", "reg2reg"],
            ["PUSH", "regonly"],
            ["POOP", "regonly"],
            ["CMPI", "imm2reg"],
            ["CMPR", "reg2reg"],
            ["JMPI", "jump"],
            ["JMPR", "jump"],
            ["JPAI", "jump"],
            ["JPAR", "jump"],
            ["JPBI", "jump"],
            ["JPBR", "jump"],
            ["JPEI", "jump"],
            ["JPER", "jump"],
            ["JPNI", "jump"],
            ["JPNR", "jump"],
            ["RETN", "single"],
            ["SHIT", "single"],
            ["NOPE", "single"],
            ["GRMN", "single"]]

reg_names = ["R0", "R1", "R2", "R3", "S0", "S1", "S2", "S3", "IP", "BP", "SP"]
ops = [VMComponent(le[0], i, le[1]) for i, le in enumerate(op_names)]
regs = [VMComponent(s.casefold(), i) for i, s in enumerate(reg_names)]
functions = []
instruction_re = re.compile("([\w]{4})(?:\ +(?:([\w]+)\ *(?:,[\ ]*([\w]+))*))?") # 1: opcode 2+: args
function_re = re.compile("(?:def\ )([a-zA-Z]*)\:")
immediate_re = re.compile("(?:0x)?[0-9]*[0-9]$")
register_re = re.compile("(^[rRsS]{1}[0-4]{1}$)|([iIrRsS]{1}[pP]{1}$)")
labeldef_re = re.compile("([a-zA-Z]*)\:")
labelcall_re = re.compile("(?:[jJ]{1}[pPmM]{1}[pPaAbBeEnN]{1}[iIrR]{1}\ *)([\w]*)")

def parse_functions(lines):
    current_size = 0
    cur_func = None

    # first parsing to get functions' names
    for i, line in enumerate(lines):
        match = function_re.match(line)
        if match:
            if cur_func:
                tmp = next(x for x in functions if x.name == cur_func)
                tmp.set_line_end(i-1)
            cur_func = match.group(2)
            functions.append(VMFunction(cur_func, i + 1))
            continue
    tmp = next(x for x in functions if x.name == cur_func)
    tmp.set_line_end(i)

    # calculating sizes and offsets
    for line in lines:
        match = function_re.match(line)
        if match:
            if cur_func:
                tmp = next(x for x in functions if x.name == cur_func)
                tmp.set_size(current_size)
            cur_func = match.group(2)
            current_size = 0
            continue
        components = [x for x in instruction_re.match(line).groups() if x is not None]
        current_size += VMInstruction(components[0], components[1:]).size
    tmp = next(x for x in functions if x.name == cur_func)
    tmp.set_size(current_size)

    # if not, main as to be the first entry
    for i in range(len(functions)):
        if functions[i].name == "main" and i is not 0:
            functions[0], functions[i] = functions[i], functions[0]
            break

    calc_fun_offsets()

def calc_fun_offsets():
    current_offset = 0
    for i in range(1, len(functions)):
        prev_size = functions[i - 1].size
        current_offset += prev_size
        functions[i].set_offset(current_offset)


def main():
    if len(sys.argv) < 4:
        print("Usage: {} opcodes_key file_to_assemble output".format(
            sys.argv[0]))
        return

    vma = VMAssembler(sys.argv[1])
    with open(sys.argv[2], 'r') as f:
        filedata = f.readlines()
    filedata = [x.strip() for x in filedata if x.strip()]

    # let's parse the whole file for labels
    parse_functions(filedata)

    if "main" not in [x.name for x in functions]:
        sys.stderr.write("No main specified!")
        return
    
    for s in functions:
        section_code = filedata[s.line_start:s.line_end+1]
        for line in section_code:
            vma.process_code_line(line)

    with open(sys.argv[3], 'wb') as f:
        f.write(vma.assembled_code)

if __name__ == '__main__':
    main()