vm.py

import os
import struct
import sys

INTRANGE = 2 ** 15
REGCOUNT = 8

class Arch:
	def __init__(self):
		self.registers = [0] * REGCOUNT
		self.stack = []
		self.inbuffer = ""
		self.opcodes = [
			self.halt,
			lambda: self.writereg(self.read(), self.readreg()),
			lambda: self.stack.append(self.readreg()),
			lambda: self.writereg(self.read(), self.stack.pop()),
			lambda: self.writereg(self.read(), int(self.readreg() == self.readreg())), # eq
			lambda: self.writereg(self.read(), int(self.readreg() > self.readreg())), # gt
			lambda: self.setpc(self.readreg()), # jmp
			self.jt,
			self.jf,
			lambda: self.writereg(self.read(), self.readreg() + self.readreg()),
			lambda: self.writereg(self.read(), self.readreg() * self.readreg()),
			lambda: self.writereg(self.read(), self.readreg() % self.readreg()),
			lambda: self.writereg(self.read(), self.readreg() & self.readreg()),
			lambda: self.writereg(self.read(), self.readreg() | self.readreg()),
			lambda: self.writereg(self.read(), self.readreg() ^ 0x7FFF), # not
			lambda: self.writereg(self.read(), self.code[self.readreg()]), # rmem
			self.wmem,
			self.call,
			self.ret,
			self.out,
			self.input,
			lambda: False # noop
		]
		self.pc = 0
		pass
	def runfile(self, file):
		code = [0,] * INTRANGE
		i = 0
		with open(file, "rb") as f:
			short = f.read(2)
			while len(short) == 2:
				code[i] = struct.unpack("H", short)[0]
				i += 1
				short = f.read(2)
		self.run(code)
	def run(self, code):
		self.pc = 0
		self.code = code
		while True:
			opcode = code[self.pc]
#			print self.pc #, '>', opcode, code[self.pc:self.pc+4], self.registers, self.stack
			self.pc += 1
			if self.opcodes[opcode]:
				self.opcodes[opcode]()
			else:
				raise Exception("Unkown opcode: {}".format(opcode))
	def read(self):
		short = self.code[self.pc]
		self.pc += 1
		if short >= (INTRANGE + REGCOUNT):
			raise Exception("Invalid number: {}".format(short))
		return short
	def write(self, addr, val):
		assert type(val) is int
		if addr >= INTRANGE:
			raise Exception("Address out of range: {}".format(addr))
		self.code[addr] = val % INTRANGE
	def isreg(self, n):
		return (INTRANGE) <= n < (INTRANGE + REGCOUNT)
	def readreg(self):
		short = self.read()
		if self.isreg(short):
			return self.registers[short - INTRANGE]
		return short
	def writereg(self, addr, val):
		assert type(val) is int
		if self.isreg(addr):
			self.registers[addr - INTRANGE] = val % INTRANGE
		else:
			raise Exception("Not register")
	def setpc(self, addr):
		self.pc = addr
	def halt(self):
		raise Exception("End of file")
	def out(self):
		sys.stdout.write(chr(self.readreg()))
		sys.stdout.flush()
	def noop(self):
		pass
	def jmp(self):
		addr = self.readreg()
		self.pc = addr
	def jt(self):
		a = self.readreg()
		b = self.readreg()
		if a != 0:
			self.pc = b
	def jf(self):
		a = self.readreg()
		b = self.readreg()
		if a == 0:
			self.pc = b
	def call(self):
		jmpaddr = self.readreg()
		self.stack.append(self.pc)
		self.pc = jmpaddr
	def wmem(self):
		to_this_address = self.readreg()
		write_this_value = self.readreg()
		self.code[to_this_address] = write_this_value
	def ret(self):
		try:
			self.pc = self.stack.pop()
		except IndexError:
			raise Exception("I need custom exceptions")
	def input(self):
		if self.inbuffer == "":
			self.inbuffer = raw_input() + "\n"
		self.writereg(self.read(), ord(self.inbuffer[0]))
		self.inbuffer = self.inbuffer[1:]
arch = Arch()
arch.runfile("challenge.bin")