from System.CodeDom import *
from System.CodeDom.Compiler import CodeDomProvider, CodeGeneratorOptions
from System.IO import StringWriter
import System.String, System.Int16

from Struct import Struct

CS = 0x2E
DS = 0x3E
ES = 0x26
FS = 0x64
GS = 0x65
SS = 0x36
OPSIZE = 0x66
ADSIZE = 0x67
LOCK = 0xF0
REPNE = 0xF2
REP = 0xF3

V = 1

def memAccess(expr):
	return CodeArrayIndexerExpression(
		CodeVariableReferenceExpression('Memory'),
		expr
	)

AX = CodeVariableReferenceExpression('ax')
CX = CodeVariableReferenceExpression('cx')
DX = CodeVariableReferenceExpression('dx')
BX = CodeVariableReferenceExpression('bx')
SP = CodeVariableReferenceExpression('sp')
BP = CodeVariableReferenceExpression('bp')
SI = CodeVariableReferenceExpression('si')
DI = CodeVariableReferenceExpression('di')
regs = AX, CX, DX, BX, SP, BP, SI, DI

Mod00 = (
	memAccess(CodeBinaryOperatorExpression(BX, CodeBinaryOperatorType.Add, SI)),
	memAccess(CodeBinaryOperatorExpression(BX, CodeBinaryOperatorType.Add, DI)),
	memAccess(CodeBinaryOperatorExpression(BP, CodeBinaryOperatorType.Add, SI)),
	memAccess(CodeBinaryOperatorExpression(BP, CodeBinaryOperatorType.Add, DI)),
	memAccess(SI),
	memAccess(DI),
	None,
	memAccess(BX)
)
Mod01_10 = (
	CodeBinaryOperatorExpression(BX, CodeBinaryOperatorType.Add, SI),
	CodeBinaryOperatorExpression(BX, CodeBinaryOperatorType.Add, DI),
	CodeBinaryOperatorExpression(BP, CodeBinaryOperatorType.Add, SI),
	CodeBinaryOperatorExpression(BP, CodeBinaryOperatorType.Add, DI),
	SI,
	DI,
	BP,
	BX
)

def formE(size, code, first, off, prefixes, (mod, reg, rm)):
	if mod == 0:
		ret = Mod00[rm]
		if ret == None:
			return off + 2, memAccess(CodePrimitiveExpression(Struct.uint16(code[off:off+2])))
		return off, ret
	elif mod == 1:
		ret = Mod01_10
		return off + 1, memAccess(
			CodeBinaryOperatorExpression(
				ret, 
				CodeBinaryOperatorType.Add, 
				CodePrimitiveExpression(ord(code[off]))
			)
		)
	elif mod == 2:
		ret = Mod01_10
		return off + 2, memAccess(
			CodeBinaryOperatorExpression(
				ret, 
				CodeBinaryOperatorType.Add, 
				CodePrimitiveExpression(Struct.uint16(code[off:off+2]))
			)
		)
	else:
		return off, regs[rm]

def formG(size, code, first, off, prefixes, (mod, reg, rm)):
	return off, regs[reg]

forms = {
	'E' : (True, formE),
	'G' : (True, formG),
}
sizes = {
	'v' : V,
}

ARITH_ADD = CodeBinaryOperatorType.Add
ARITH_SUB = CodeBinaryOperatorType.Subtract
ARITH_MUL = CodeBinaryOperatorType.Multiply
ARITH_DIV = CodeBinaryOperatorType.Divide
ARITH_XOR = '^'
ARITH_LSL = '<<'
ARITH_LSR = '>>'

class Decompiler(object):
	def __init__(self):
		self.provider = CodeDomProvider.CreateProvider('CSharp')
		self.options = CodeGeneratorOptions()
		self.options.BracingStyle = 'Block'
		
		add = lambda A, B: self.arith(ARITH_ADD, A, B)
		sub = lambda A, B: self.arith(ARITH_SUB, A, B)
		mul = lambda A, B: self.arith(ARITH_MUL, A, B)
		div = lambda A, B: self.arith(ARITH_DIV, A, B)
		xor = lambda A, B: self.arith(ARITH_XOR, A, B)
		lsl = lambda A, B: self.arith(ARITH_LSL, A, B)
		lsr = lambda A, B: self.arith(ARITH_LSR, A, B)
		move = CodeAssignStatement
		
		self.opcodes = {
			0x03 : (self.form('Gv,Ev'), lambda D, S: move(D, add(D, S)), False),
			0x33 : (self.form('Gv,Ev'), lambda D, S: move(D, xor(D, S)), False),
		}
	
	def compileExpr(self, expr):
		writer = StringWriter()
		self.provider.GenerateCodeFromExpression(expr, writer, self.options)
		ret = writer.ToString()
		writer.Close()
		return ret
	
	def arith(self, op, A, B):
		if isinstance(op, CodeBinaryOperatorType):
			return CodeBinaryOperatorExpression(
				A,
				op,
				B
			)
		else:
			return CodeSnippetExpression(
					'(%s %s %s)' % (
						self.compileExpr(A),
						op,
						self.compileExpr(B)
					)
				)

	def branch(L, C=None, cond=None, hintTaken=False):
		pass
	
	def comp(op, A, B):
		pass
	
	def form(self, formstr):
		outform = []
		for opc in formstr.split(','):
			outform.append((forms[opc[0]], sizes[opc[1]]))
		return tuple(outform)
	
	def decode(self, code, off):
		prefixes = []
		c = ord(code[off])
		off += 1
		while c in (CS, DS, ES, FS, GS, SS, OPSIZE, ADSIZE, LOCK, REPNE, REP):
			prefixes.append(c)
			c = ord(code[off])
			off += 1
		
		if c in self.opcodes:
			opcd = self.opcodes[c]
		else:
			c = (c << 8) | ord(code[off])
			off += 1
			
			if c in self.opcodes:
				opcd = self.opcodes[c]
			else:
				c = (c << 8) | ord(code[off])
				off += 1
				
				if c in self.opcodes:
					opcd = self.opcodes[c]
				else:
					raise DecoderError('Unknown opcode: %06x (may not be three bytes)' % c)
		
		first = off
		modrm = None
		
		ops = []
		forms, func, branch = opcd
		for form in forms:
			(hasModrm, form), size = form
			if hasModrm:
				modrm = ord(code[first])
				modrm = modrm >> 6, (modrm >> 3) & 7, modrm & 7
				off += 1
				break
		for form in forms:
			(hasModrm, form), size = form
			if hasModrm:
				off, stmt = form(size, code, first, off, prefixes, modrm)
			else:
				off, stmt = form(size, code, first, off, prefixes)
			ops.append(stmt)
		
		return func(*ops), off, branch
	
	def decompile(self, code, off):
		method = CodeMemberMethod()
		method.Name = 'Decompiled'
		method.ReturnType = CodeTypeReference(System.String)
		method.Attributes = MemberAttributes.Public
		
		branch = False
		while not branch and len(code) > off:
			stmt, off, branch = self.decode(code, off)
			method.Statements.Add(stmt)
		
		writer = StringWriter()
		self.provider.GenerateCodeFromMember(method, writer, self.options)
		print writer.ToString()
		writer.Close()

if __name__=='__main__':
	decomp = Decompiler()
	decomp.decompile('\x33\xC9\x33\xD2\x03\xD0', 0)