[flareon2018 level12] tiny disassembler for RSSB VM ( edit: I thought it was a subleq2 VM, but its RSSB instead )

disasm2.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
from pwn import *
import numpy
from collections import OrderedDict
import json
import inst2

_120d = 0
with open("./dump1", "rb") as f:
    data = f.read()

assert len(data) == 0x8400

code = dict()

def Word(addr):
    global data
    return u16(data[addr:addr+2:])

def patch_word(addr, d):
    global data
    tmp = list(data)
    d1 = d & 0xff
    d2 = (d>>8)&0xff;
    tmp[addr] = chr(d1)
    tmp[addr+1] = chr(d2)
    data = ''.join(tmp)

def subleq2():
    j = dict()
    _120f = Word(0x120f)
    cur_addr = (((0x7f6+ _120f )<<1)+0x223)&0xffff
    _1201 = Word(cur_addr)
    if _1201 == 0xffff:
        _120f += 1
        patch_word(0x120f, _120f)
        return

    cur = (((0x7f6+ _1201 )<<1)+0x223)&0xffff
    _120b = Word(cur)
    j['op1'] = cur
    _11ff = numpy.int16((_120b - Word(0x1211)))
    patch_word(0x1211, _11ff) # 0x1211 = acc
    
    s = "{:#x}: ".format(cur_addr)
    s += "subleq2 [{:#x}]".format(cur)

    if _11ff < 0:
        _120f += 2
        s += " [{:#x}]".format(cur+2)
        j['op2'] = cur+2
    else:
        _120f += 1
    
    j['asm'] = s
    code[cur_addr] = json.dumps(j)
    patch_word(0x120f, _120f)

    return s

def get_inst():
    while True:
        _120f = Word(0x120f)
        if _120f >= 0x2b57:
            return 

        tmp = Word((((0x7f6+ _120f )<<1)+0x223)&0xffff)

        if tmp == -2:
            return

        d = subleq2()
        if d:
            pass
            #print hex(Word(0x120f))
            #print d

# starting point, might need to change in order to disassemble the whole thing
patch_word(0x120f, 0x15b)
patch_word(0x120d, 0x15b)

try:
    get_inst()
except:
    pass
#exit(0)
od = OrderedDict(sorted(code.items()))
dis = []
black = []
for pos, code in od.iteritems():
    if pos not in black:
        dis.append((pos, code))
        
token, rev_token = dict(), dict()
cur_token = None
Z = 0x1213
acc = 0x1211
pc = 0x120f

def trans(code):
    def get_token(op):
        global token, cur_token, rev_token
        ret = None
        if op == Z:
            return hex(Z)
        if op == acc:
            return hex(acc)
        if op == pc:
            return hex(pc)

        if op not in token:
            ret = chr(cur_token)
            token[op] = ret
            rev_token[ret] = op
            cur_token += 1
        else:
            ret = token[op]
        return ret

    tk2 = None
    tk1 = get_token(code["op1"])
    
    return tk1

now = 0
while now < len(dis):
    matched = False
    for ins in inst2.instruction: # traverse all instructions
        sz = ins.sz
        token.clear()
        rev_token.clear()
        cur_token = ord('a')
        patt = []
        
        if now+sz-1 >= len(dis): continue

        for idx in xrange(sz):
            code = json.loads(dis[now+idx][1])
            patt.append(trans(code))
        
        if (ins.check(patt)):
            print("{:#x}: {}".format(dis[now][0], ins.get_inst(rev_token, data)))
            matched = True
            now += sz
            break

    if not matched:
        print(json.loads(dis[now][1])["asm"])
        now += 1
    

inst2.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from pwn import *

a = "a"
Z = "0x1213"
acc = "0x1211"
pc = "0x120f"
b = "b"

def Word(addr, data):
    return u16(data[addr:addr+2:])

class Inst(object):
    def __init__(self, name, patt, sz):
        """
        name: instruction name
        patt: subleq pattern ( an array )
        sz: subleq instruction number
        """
        self.name = name
        self.patt = patt
        self.sz = sz
    def check(self, input_patt):
        """
        check if a input pattern ( an array ) equals this instruction
        """
        res = input_patt == self.patt
        #print input_patt
        #print self.patt
        return res

class Jmp(Inst):
    def __init__(self, patt):
        super(Jmp, self).__init__("JMP", patt, 5)
    def get_inst(self, *args):
        rev_token = args[0]
        data = args[1]
        jmp_addr = rev_token["a"]
        return "JMP $+{}".format(Word(jmp_addr, data))

class Move(Inst):
    def __init__(self, patt):
        super(Move, self).__init__("MOV", patt, 10)
    def get_inst(self, *args):
        rev_token = args[0]
        data = args[1]
        dst = rev_token["a"]
        src = rev_token["b"]
        return "MOV [{:#x}], [{:#x}] ; {:#x}".format(dst, src, Word(src, data))

class Add(Inst):
    def __init__(self, patt):
        super(Add, self).__init__("ADD", patt, 4)
    def get_inst(self, *args):
        rev_token = args[0]
        data = args[1]
        dst = rev_token["b"]
        src = rev_token["a"]
        return "ADD [{:#x}], [{:#x}] ; {:#x}".format(dst, src, Word(src, data))

class Sub(Inst):
    def __init__(self, patt):
        super(Sub, self).__init__("SUB", patt, 5)
    def get_inst(self, *args):
        rev_token = args[0]
        data = args[1]
        dst = rev_token["b"]
        src = rev_token["a"]
        return "SUB [{:#x}], [{:#x}] ; {:#x}".format(dst, src, Word(src, data))

class Clr(Inst):
    def __init__(self, patt):
        super(Clr, self).__init__("CLR", patt, 5)
    def get_inst(self, *args):
        rev_token = args[0]
        dst = rev_token["a"]
        return "CLR [{:#x}], [{:#x}]".format(dst, dst)

instruction = []
instruction.append(Move([a, Z, Z, a, acc, b, Z, a, acc, acc]))
instruction.append(Add([a, Z, b, acc]))
instruction.append(Sub([a, Z, Z, b, acc]))
instruction.append(Clr([a, Z, Z, a, acc]))
instruction.append(Jmp([a, Z, pc, Z, acc]))