from struct import *
import binascii

VM_REGISTER_SET = ('R0', 'R1', 'R2', 'R3', 'R4', 'R5', 'R6', 'R7', 'R8', 'R9', 'R10', 'R11', 'R12', 'R13', 'R14', 'R15')
VM_SIZE = ("BYTE", "WORD", "DWORD")
VM_MATH = ("XOR", "ADD", "SUB", "SHL", "SHR", "ROL", "ROR", "DIV")

def parse_src(CODE, DATA, p0, p1, p2):
	size = 0
	mnemonic = ""
	
	if p0 == 0:
		mnemonic = "%s %s"%(VM_SIZE[p2], VM_REGISTER_SET[p1])
	elif p0 == 1:
		mnemonic = "%s &%s"%(VM_SIZE[p2], VM_REGISTER_SET[p1])
	elif p0 == 2:
		if p2 == 0:
			mnemonic = "0x%02X"%CODE[0]
			size += 1
		elif p2 == 1:
			mnemonic = "0x%04X"%unpack('<H', CODE[0:2])
			size += 2
		elif p2 == 2:
			mnemonic = "0x%08X"%unpack('<L', CODE[0:4])
			size += 4
	elif p0 == 3:
		(offset, ) = unpack('<L', CODE[0:4])
		mnemonic = "&data_%X"%offset
		size += 4

	return (size, mnemonic)

def parse_dst(CODE, p0, p1, p2):
	size = 0
	mnemonic = ""
	if p0 == 0:
		mnemonic = "DWORD %s"%VM_REGISTER_SET[p2]
	elif p0 == 1:
		mnemonic = "%s &%s"%(VM_SIZE[p1], VM_REGISTER_SET[p2])

	return (size, mnemonic)

ROM_DATA = open(input("Filepath to the ROM image: "), "rb").read()

# standard ROM
if unpack('<L', ROM_DATA[14:18]) == 0:
	(Signature, AddressOfEntryPoint, BaseOfData, ImageSize, Reserved) = unpack('<HLLLL', ROM_DATA[0:18])
# driver ROM
else:
	(Signature, ImageSize, AddressOfEntryPoint, BaseOfData, Reserved) = unpack('<HLLLL', ROM_DATA[0:18])
	Signature = (((Signature&0xFF)<<8)|((Signature&0xFF00)>>8)) # BSWAP, because of the wrong endian

# assign the CODE section
CODE = b""
if Signature == 0x1337:
	print("Code is not encrypted - Disassembling the entry point")
	CODE = ROM_DATA[AddressOfEntryPoint:BaseOfData]
else:
	print("Code is encrypted - Disassembling the loader stub")
	CODE = ROM_DATA[calcsize('<HLLLL'):AddressOfEntryPoint]

# assign the DATA section
DATA = ROM_DATA[BaseOfData:]

IP = 0
ZFLAG = ""
while IP < len(CODE):

	# save the instruction's OPCODE 
	OPCODE = CODE[IP]

	# save the current instruction pointer address
	ADDRESS = IP

	# increment the IP to skip the instruction OPCODE
	IP += 1

	# read the additional bytes for each instruction
	if OPCODE in (0x00, ):
		pass
	elif OPCODE in (0x02, 0x04, 0x07, 0x0B, 0x0C):
		a0 = CODE[IP]
		IP += 1		
	elif OPCODE in (0x01, 0x03, 0x06, 0x0A, 0x0D):
		a0, a1 = unpack("<BB", CODE[IP:IP+2])
		IP += 2		
	else:
		print("Invalid instruction: %X"%OPCODE)
		break

	# start reading the instructions
	INSTRUCTION = ""
		
	# RET
	if OPCODE == 0x00:
		INSTRUCTION = "RET"
	# MOV
	elif OPCODE == 0x01:
		OPERAND_SRC = parse_src(CODE[IP:], DATA, (a1&3), (a0>>4&0x0F), (a1>>4&3))
		IP += OPERAND_SRC[0]

		OPERAND_DST = parse_dst(CODE[IP:], (a1>>2&3), (a1>>4&3), (a0&0x0F))
		IP += OPERAND_DST[0]

		INSTRUCTION = "MOV %s, %s"%(OPERAND_DST[1], OPERAND_SRC[1])
	# ??
	elif OPCODE == 0x02:
		INSTRUCTION = "$0x02"
	# CALL
	elif OPCODE == 0x03:
		OPERAND_SRC = parse_src(CODE[IP:], DATA, (a1&3), (a0&0x0F), 2)
		IP += OPERAND_SRC[0]

		OPERAND_DST = parse_src(CODE[IP:], DATA, (a1>>2&3), (a0>>4&0x0F), 2)
		IP += OPERAND_DST[0]

		INSTRUCTION = "CALL %s->%s"%(OPERAND_DST[1], OPERAND_SRC[1])
	# PUSH
	elif OPCODE == 0x04:
		OPERAND_SRC = parse_src(CODE[IP:], DATA, (a0>>4&3), (a0&0x0F), (a0>>6&3))
		IP += OPERAND_SRC[0]
		
		INSTRUCTION = "PUSH %s"%OPERAND_SRC[1]
	# CMP
	elif OPCODE == 0x06:
		ZFLAG = ("JZ", "JNZ", "JNB")[a1>>4&7]
	
		OPERAND_SRC = parse_src(CODE[IP:], DATA, (a1&3), (a0>>4&0x0F), (a1>>2&3))
		IP += OPERAND_SRC[0]

		INSTRUCTION = "CMP DWORD %s, %s"%(VM_REGISTER_SET[a0&0x0F], OPERAND_SRC[1])
	# JMP
	elif OPCODE == 0x07:
		OPERAND_SRC = parse_src(CODE[IP:], DATA, 2, 0, 2)
		IP += OPERAND_SRC[0]

		if a0 == 1:
			INSTRUCTION = "%s %08X"%(ZFLAG, int(OPERAND_SRC[1], 16))
		else:
			INSTRUCTION = "JMP %08X"%int(OPERAND_SRC[1], 16)
			
		ZFLAG = ""
	# ARITHMETIC
	elif OPCODE == 0x0A:
		OPERAND_SRC = parse_src(CODE[IP:], DATA, (a1&3), (a0>>4&0x0F), (a1>>2&3))
		IP += OPERAND_SRC[0]

		OPERAND_DST = "%s %s"%(VM_SIZE[a1>>2&3], VM_REGISTER_SET[a0&0x0F])

		INSTRUCTION = "%s %s, %s"%(VM_MATH[a1>>4&7], OPERAND_DST, OPERAND_SRC[1])
	# MALLOC
	elif OPCODE == 0x0B:
		OPERAND_SRC = parse_src(CODE[IP:], DATA, (a0>>4&3), (a0&0x0F), (a0>>6&3))
		IP += OPERAND_SRC[0]

		INSTRUCTION = "MOV DWORD R0, MALLOC(%s)"%OPERAND_SRC[1]
	# FREE
	elif OPCODE == 0x0C:
		OPERAND_SRC = parse_src(CODE[IP:], DATA, (a0>>4&3), (a0&0x0F), (a0>>6&3))
		IP += OPERAND_SRC[0]

		INSTRUCTION = "FREE(%s)"%OPERAND_SRC[1]
	# LOCAL CALL
	elif OPCODE == 0x0D:
		OPERAND_SRC = parse_src(CODE[IP:], DATA, (a1&3), (a0&0x0F), 2)
		IP += OPERAND_SRC[0]

		OPERAND_DST = parse_src(CODE[IP:], DATA, (a1>>2&3), (a0>>4&0x0F), (a1>>4&3))
		IP += OPERAND_DST[0]

		INSTRUCTION = "CALL data_%s"%OPERAND_DST[1]
	else:
		print("MORE RE NEEDED!! 0x%02X"%OPCODE)
		break

	print("%08X | "%ADDRESS, end="")
	print((" ".join("{:02X}".format(x) for x in CODE[ADDRESS:IP])).ljust(34), end="")
	print(INSTRUCTION)