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

class DecoderError(Exception):
	pass

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

B = 1
V = 2
Z = 3

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

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

class Decompiler(object):
	def __init__(self):
		self.provider = CodeDomProvider.CreateProvider('CSharp')
		self.options = CodeGeneratorOptions()
		self.options.BracingStyle = 'Block'
		
		self.forms = {
			'E' : (True, self.formE),
			'G' : (True, self.formG),
			'I' : (False, self.formI),
		}
		self.sizes = {
			'b' : B,
			'v' : V,
			'z' : Z,
		}
		
		self.Mod00 = (
			self.memAccess(CodeBinaryOperatorExpression(BX, CodeBinaryOperatorType.Add, SI)),
			self.memAccess(CodeBinaryOperatorExpression(BX, CodeBinaryOperatorType.Add, DI)),
			self.memAccess(CodeBinaryOperatorExpression(BP, CodeBinaryOperatorType.Add, SI)),
			self.memAccess(CodeBinaryOperatorExpression(BP, CodeBinaryOperatorType.Add, DI)),
			self.memAccess(SI),
			self.memAccess(DI),
			None,
			self.memAccess(BX)
		)
		self.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
		)
		
		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)
		and_= lambda A, B: self.arith(ARITH_AND, A, B)
		or_ = lambda A, B: self.arith(ARITH_OR , 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
		
		group1 = (
			lambda D, S: move(D, add(D, S)),
			lambda D, S: move(D, or_(D, S)),
			lambda D, S: move(D, adc(D, S)),
			lambda D, S: move(D, sbb(D, S)),
			lambda D, S: move(D, and_(D, S)),
			lambda D, S: move(D, sub(D, S)),
			lambda D, S: move(D, xor(D, S)),
			lambda D, S: move(D, cmp_(D, S)),
		)
		
		self.opcodes = {
			0x03 : (self.form('Gv,Ev'), lambda D, S: move(D, add(D, S)), False, False),
			0x33 : (self.form('Gv,Ev'), lambda D, S: move(D, xor(D, S)), False, False),
			0x80 : (self.form('Eb,Ib'), lambda (mod, reg, rm), D, S: group1[reg](D, S), False, True),
			0x81 : (self.form('Eb,Ib'), lambda (mod, reg, rm), D, S: group1[reg](D, S), False, True),
			0x82 : (self.form('Ev,Iz'), lambda (mod, reg, rm), D, S: group1[reg](D, S), False, True),
			0x83 : (self.form('Ev,Ib'), lambda (mod, reg, rm), D, S: group1[reg](D, S), False, True),
			0x89 : (self.form('Ev,Gv'), lambda D, S: move(D, S), False, False),
		}
	
	def memAccess(self, expr):
		return CodeArrayIndexerExpression(
			CodeVariableReferenceExpression('Memory'),
			expr
		)
	
	def formE(self, size, code, first, off, prefixes, (mod, reg, rm)):
		if mod == 0:
			ret = self.Mod00[rm]
			if ret == None:
				return off + 2, self.memAccess(CodePrimitiveExpression(Struct.uint16(code[off:off+2])))
			return off, ret
		elif mod == 1:
			ret = self.Mod01_10
			return off + 1, self.memAccess(
				CodeBinaryOperatorExpression(
					ret, 
					CodeBinaryOperatorType.Add, 
					CodePrimitiveExpression(ord(code[off]))
				)
			)
		elif mod == 2:
			ret = self.Mod01_10
			return off + 2, self.memAccess(
				CodeBinaryOperatorExpression(
					ret, 
					CodeBinaryOperatorType.Add, 
					CodePrimitiveExpression(Struct.uint16(code[off:off+2]))
				)
			)
		else:
			return off, regs[rm]
	
	def formG(self, size, code, first, off, prefixes, (mod, reg, rm)):
		return off, regs[reg]
	
	def formI(self, size, code, first, off, prefixes):
		if size == B:
			return off+1, CodePrimitiveExpression(ord(code[off]))
		elif size == Z:
			return off+2, CodePrimitiveExpression(Struct.uint16(code[off:off+2]))

	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((self.forms[opc[0]], self.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, instModrm = 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)
		
		if instModrm:
			return func(modrm, *ops), off, branch
		else:
			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\x89\x1E\xBE\x1E\x83\xC3\x0F', 0)