import os
import sys
import pickle
import zlib
import hashlib
import copy
from termcolor import colored
from base64 import *
from colorama import init
init()

OS = "Linux"
#OS = "Windows"

doc = """
0000 NOP
0001 LDC par > A (max 20bit)
0010 LDV [par] > A
0011 STV A > [par]
0100 ADD A + [par] > A
0101 AND A AND [par] > A
0110 LOR A OR [par] > A
0111 XOR A XOR [par] > A
1000 EQL if A==[par]: -1 else 0 > A
1001 JMP Jump to [par]
1010 JMN if A<0: Jump to [par]
1011 NOT NOT A > A
1100 RAR A >> 1 > A
1101 LDI [[par]] > A
1110 STI A > [[par]]
1111 HLT Halt
"""

asmMap = ["NOP","LDC","LDV","STV","ADD","AND","LOR","XOR","EQL","JMP","JMN","NOT","RAR","LDI","STI","HLT"]

bits = 24
opBits = 4
parBits = 20

pointer = 0

registerA = 0
out = 0

color = [["green","on_grey"],["blue","on_grey"],["red","on_grey"],["magenta","on_grey"],["cyan","on_grey"],
		 ["grey","on_white"],["green","on_white"],["blue","on_white"],["red","on_white"],["magenta","on_white"],["cyan","on_white"],
		 ["grey","on_green"],["white","on_green"],["blue","on_green"],["red","on_green"],["magenta","on_green"],["cyan","on_green"],
		 ["grey","on_blue"],["white","on_blue"],["green","on_blue"],["red","on_blue"],["magenta","on_blue"],["cyan","on_blue"],
		 ["grey","on_red"],["white","on_red"],["green","on_red"],["blue","on_red"],["magenta","on_red"],["cyan","on_red"],
		 ["grey","on_magenta"],["white","on_magenta"],["green","on_magenta"],["blue","on_magenta"],["red","on_magenta"],["cyan","on_magenta"],
		 ["grey","on_cyan"],["white","on_cyan"],["green","on_cyan"],["blue","on_cyan"],["red","on_cyan"],["magenta","on_cyan"]]

memory = []

for i in range(2**parBits):
	memory.append(0)

def clear():
	if OS=="Windows":
		os.system("cls")
	elif OS=="Linux":
		os.system("clear")
	else:
		print("Error: Unknown OS")
		exit()

def toDec(binary):
	try:
		return int(binary, 2)
	except:
		return -1

def toBin(decimal):
	try:
		return int(bin(decimal)[2:])
	except:
		print("Error: The dec-number "+decimal+" could not be converted to binary")
		return decimal

def padBin(binary,bit=bits):
	return str(binary).zfill(bit)[-bit:]

def bin_input(display,bit=bits,allowAsm=False):
	raw = input(display)
	if raw=="x":
		return -1
	if allowAsm:
		for i in range(len(asmMap)):
			if raw.find(asmMap[i])!=-1:
				raw = raw.replace(asmMap[i],"")
				raw = raw.replace(" ","")
				if raw.find("_")!=-1:
					raw = raw.replace("_","0"*(bit-(len(raw)-1)-4))
				raw = str(toBin(i)).zfill(opBits)+raw.zfill(bits-opBits)
				break
		else:
			raw = raw.replace(" ","")
			if raw.find("_")!=-1:
				raw = raw.replace("_","0"*(bit-(len(raw)-1)))
	else:
		raw = raw.replace(" ","")
		if raw.find("_")!=-1:
			raw = raw.replace("_","0"*(bit-(len(raw)-1)))
	if raw.find("_")!=-1:
		print("[!] input contains multiple spread-zero characters ( _ )")
		return bin_input(display)
	ret = toDec(raw)
	if ret>2**bits:
		print("[!] input exceted maximal size for "+str(bit)+"-bit integer!")
		return bin_input(display,allowAsm=allowAsm)
	if ret==-1:
		print("[!] "+raw+" is not a valid binary number!")
		return bin_input(display,allowAsm=allowAsm)
	return ret

def parse(cmd,par,verbose=False,adress=0):
	global registerA
	global memory
	global out
	global pointer
	v=verbose
	cmd = cmd[-1*opBits:]
	if cmd=="0000": #nop
		if v: print(padBin(toBin(adress),bit=parBits)+": NOP")
	elif cmd=="0001": #ldc
		registerA = par
		if v: print(padBin(toBin(adress),bit=parBits)+": LDC "+padBin(toBin(par),bit=parBits)+" > Akku")
	elif cmd=="0010": #ldv
		registerA = memory[par]
		if v: print(padBin(toBin(adress),bit=parBits)+": LDV ["+padBin(toBin(par),bit=parBits)+"] > Akku")
	elif cmd=="0011": #stv
		memory[par] = registerA
		if v: print(padBin(toBin(adress),bit=parBits)+": STA Akku > ["+padBin(toBin(par),bit=parBits)+"]")
	elif cmd=="0100": #add
		registerA = (registerA + memory[par]) % 2**parBits
		if v: print(padBin(toBin(adress),bit=parBits)+": ADD Akku + ["+padBin(toBin(par),bit=parBits)+"]")
	elif cmd=="0101": #and
		newRegA = ""
		for b in range(len(padBin(toBin(registerA)))):
			if padBin(toBin(int(registerA)))[b]=="1" and padBin(toBin(int(memory[par])))[b]=="1":
				newRegA+="1"
			else:
				newRegA+="0"
		registerA = toDec(newRegA)
		if v: print(padBin(toBin(adress),bit=parBits)+": AND Akku, ["+padBin(toBin(par),bit=parBits)+"]")
	elif cmd=="0110": #lor
		newRegA = ""
		for b in range(len(str(registerA))):
			newRegA+=str(int( int(padBin(toBin(int(registerA)))[b]) or int(padBin(toBin(int(memory[par])))[b]) ))
		registerA = toDec(newRegA)
		if v: print(padBin(toBin(adress),bit=parBits)+": OR Akku, ["+padBin(toBin(par),bit=parBits)+"]")
	elif cmd=="0111": #XOR
		newRegA = ""
		for b in range(len(padBin(toBin(registerA)))):
			if padBin(toBin(registerA))[b] == padBin(toBin(memory[par]))[b]:
				newRegA+="0"
			else:
				newRegA+="1"
		registerA = toDec(newRegA)
		if v: print(padBin(toBin(adress),bit=parBits)+": XOR Akku, ["+padBin(toBin(par),bit=parBits)+"]")
	elif cmd=="1000": #EQL
		if registerA==memory[par]:
			registerA = toDec("1"*bits)
			if v: print(padBin(toBin(adress),bit=8)+": EQL {TRUE} Akku ?= "+padBin(toBin(par),bit=parBits))
		else:
			registerA = 0
			if v: print(padBin(toBin(adress),bit=8)+": EQL {FALSE} Akku ?= "+padBin(toBin(par),bit=parBits))
	elif cmd=="1001": #jmp
		pointer = par - 1
		if v: print(padBin(toBin(adress),bit=8)+": JMP "+padBin(toBin(par),bit=parBits))
	elif cmd=="1010": #jmn
		if registerA>=int("1"+"0"*(bits-1)):
			pointer = par - 1
			if v: print(padBin(toBin(adress),bit=8)+": JMN {TRUE} > "+padBin(toBin(par),bit=8))
		else:
			if v: print(padBin(toBin(adress),bit=8)+": JMN {FALSE} > "+padBin(toBin(par),bit=8))
	elif cmd=="1011": #not
		newRegA = ""
		for b in range(len(padBin(toBin(registerA)))):
			if padBin(toBin(registerA))[b]=="1":
				newRegA+="0"
			elif padBin(toBin(registerA))[b]=="0":
				newRegA+="1"
			else:
				print("FUUUUUUUUUU")
				exit()
		registerA = toDec(newRegA)
		if v: print(padBin(toBin(adress),bit=parBits)+": NOT Akku > Akku")
	elif cmd=="1100": #rar
		newRegA = padBin(toBin(registerA))
		newRegA = newRegA[-1]+newRegA[:-1]
		registerA = toDec(newRegA)
		if v: print(padBin(toBin(adress),bit=parBits)+": RAR Akku > Akku")
	elif cmd=="1101": #LDI
		mid = memory[par]
		registerA = memory[mid % parBits]
		if v: print(padBin(toBin(adress),bit=parBits)+": LDI [["+padBin(toBin(par),bit=parBits)+"]] > Akku")
	elif cmd=="1110": #STI
		mid = memory[par]
		memory[mid] = registerA
		if v: print(padBin(toBin(adress),bit=parBits)+": STI Akku > [["+padBin(toBin(par),bit=parBits)+"]]")
	elif cmd=="1111": #hlt
		return -1
	return True

if __name__=="__main__":
	print("[asmTron by CyberCodeDragon]")
	while True:
		x = ""
		x=input("#> ")
		x = x.split(" ")
		for i in range(2):
			x.append("")
		y = x[1]
		z = x[2]
		x = x[0]
		if x=="help" or x=="h":
			print("Avaible Commands:")
			print("docs              - display avaible operants")
			print("write [w]         - write value to adress")
			print("read [r]          - read value from adress")
			print("dump [d]          - dump from adress with length")
			print("run [u]           - execute programm")
			print("edit [e]          - edit from adress til >x<")
			print("cls               - clear screen")
			print("clear / reset     - clear memory,register,pointer")
			print("step [s]          - run one command from programm")
			print("debug             - execute programm with verbose output")
			print("pointer           - display pointer")
			print("setPointer        - sets pointer")
			print("resetPointer [rp] - sets pointer to zero")
			print("register [reg]    - displays A register")
			print("setRegister       - sets A register")
			print("save              - saves memory to local file")
			print("load              - loads memory from local file")
			print("code / exec       - execute one line of shellcode")
			print("guiDebug [gd]     - graphical debugger <not yet implemented for mima-arch!>")
			print("asm               - shellcode visualizer (assembly / jump-arrows) <not yet implemented for mima-arch!>")
			print("halting           - tries to determin, if algorythm will diverge")
		elif x=="docs":
			print(doc)
		elif x=="write" or x=="w":
			if y=="" or z=="":
				y=bin_input("adress> ")
				z=bin_input("value> ")
			try:
				memory[y] = z
			except:
				print("[!] Unknown error while writing")
		elif x=="read" or x=="r":
			if y=="":
				y=bin_input("adress> ")
			print(padBin(toBin(memory[y])))
		elif x=="dump" or x=="d":
			if y=="" and z=="":
				y=bin_input("adress> ")
				z=bin_input("length> ")
			elif z=="":
				z=toDec(y)
				y=0
			for m in range(z):
				print(padBin(toBin(y+m),bit=parBits)+": "+padBin(toBin(memory[y+m])))
		elif x=="run" or x=="u":
			while True:
				cmd = padBin(toBin(memory[pointer]))[:opBits]
				par = int(toDec(padBin(toBin(memory[pointer]))[-parBits:]))
				if parse(cmd,par)==-1:
					pointer = 0
					break
				pointer = (pointer + 1) % 2**parBits
		elif x=="verboseRun" or x=="vu":
			while True:
				cmd = padBin(toBin(memory[pointer]))[:opBits]
				par = int(toDec(padBin(toBin(memory[pointer]))[-parBits:]))
				if parse(cmd,par,verbose=True,adress=pointer)==-1:
					pointer = 0
					break
				pointer = (pointer + 1) % 2**parBits
		elif x=="edit" or x=="e":
			adress=bin_input("offset> ",bit=parBits)
			while True:
				value=bin_input(padBin(toBin(adress),bit=parBits)+": "+padBin(toBin(memory[adress]))+" > ",allowAsm=True)
				if value==-1:
					break
				try:
					memory[adress] = value
				except:
					print("[!] Unknown error while writing")
				if adress==len(memory)-1:
					adress = -1
					print("[*] End of mem")
				adress = adress + 1
		elif x=="cls":
			clear()
		elif x=="clear" or x=="reset":
			pointer = 0
			registerA = 0
			out = 0
			memory = []
			for i in range(2**parBits):
				memory.append(0)
		elif x=="varClear" or x=="vc":
			pointer = 0
			registerA = 0
			out = 0
			for i in range((2**parBits)/2):
				memory[(2**parBits)/2+i] = 0
		elif x=="step" or x=="s":
			cmd = padBin(toBin(memory[pointer]))[:opBits]
			par = int(toDec(padBin(toBin(memory[pointer]))[-parBits:]))
			if parse(cmd,par,verbose=True,adress=pointer)==-1:
				pointer = 0
				print("[*] Programm halted")
			pointer = (pointer + 1) % (2**parBits)
		elif x=="debug":
			while True:
				cmd = padBin(toBin(memory[pointer]))[:opBits]
				par = int(toDec(padBin(toBin(memory[pointer]))[-parBits:]))
				if parse(cmd,par,verbose=True,adress=pointer)==-1:
					pointer = -1
					print("[*] Programm halted")
					break
				pointer = (pointer + 1) % (2**parBits)
		elif x=="pointer":
			print("POINTER: "+padBin(toBin(pointer),bit=parBits))
		elif x=="setPointer":
			if y=="":
				pointer = bin_input("pointer> ",bit=parBits)
			else:
				pointer = y
		elif x=="resetPointer" or x=="rp":
			pointer = 0
		elif x=="register" or x=="reg":
			print("REGISTER: "+padBin(toBin(registerA)))
		elif x=="setRegister":
			if y=="":
				registerA = bin_input("register> ")
			else:
				registerA = y
		elif x=="save":
			if y=="":
				print("[*] Current saves:")
				for f in os.listdir("saves"):
					if f[-8:]==".asmTron":
						print("> "+f[:-8])
				name = input("name> ")
			else:
				name = y
			with open("saves/"+name+".asmTron", 'wb') as f:
				t=pickle.dumps(["mima24_1",memory])
				t=zlib.compress(t)
				f.write(t)
		elif x=="load" or x=="l":
			if y=="":
				print("[*] Avaible saves:")
				for f in os.listdir("saves"):
					if f[-8:]==".asmTron":
						print("> "+f[:-8])
				for f in os.listdir("tronScript"):
					if f[-8:]==".asmTron":
						print("> tronScript/"+f[:-8])
				name = input("name> ")
			else:
				name = y
			try:
				with open("saves/"+name+".asmTron", "rb") as f:
					t = f.read()
					t=zlib.decompress(t)
					magicBit,memory = pickle.loads(t)
			except Exception as e:
				try:
					with open("tronScripts/"+name+".asmTron", "rb") as f:
						t = f.read()
						t=zlib.decompress(t)
						magicBit,memory = pickle.loads(t)
				except Exception as e:
					if b64encode(str(e))[:25]=="W0Vycm5vIDJdIE5vIHN1Y2ggZ":
						print(" [!] Error: File '"+name+".asmTron' does not exist")
					else:
						print(" [!] Error: Unknown Error while reading File '"+name+"' with Errorcode: "+str(b64encode(str(e))))
			if not magicBit=="mima24_1":
				print("[!] Error: Invalid Magical Bit: "+str(magicBit))
				exit()
			print("[*] loaded "+name+".asmTron")
		elif x=="code" or x=="exec":
			if y=="":
				data=bin_input("code> ",allowAsm=True)
			else:
				data=y
			cmd = padBin(toBin(memory[pointer]))[:opBits]
			par = int(toDec(padBin(toBin(memory[pointer]))[-parBits:]))
			parse(cmd,par,verbose=True,adress=0)
		elif x=="guiDebug" or x=="gd":
			if y=="":
				hiOffset = 2**(parBits-1)
			else:
				hiOffset = 2**(int(y)-1)
			runtime = 0
			tracker = input("trackers> ").split(",")
			running = True
			lastPointer = 0
			while running:
				clear()
				print(" +"+"-"*58+"+"+"-"*52+"+"+"-"*36+"+")
				print(" |"+" "*20+" - Instructions - "+" "*20+"|"+" "*20+" - Hi-mem - "+" "*20+"|"+" "*36+"|")
				if pointer <= 20:
					starter = 0
				else:
					starter = pointer - 20
				for i in range(36):
					m = i + starter
					if pointer == m:
						outI = "> "+padBin(toBin(m),bit=parBits)+": "+padBin(toBin(memory[m]))+" ["+doc.split("\n")[toDec(padBin(toBin(memory[m]))[:opBits])+1][:parBits][5:][:3]+"]"+" <"
					elif lastPointer == m and pointer-1 != lastPointer:
						outI = "- "+padBin(toBin(m),bit=parBits)+": "+padBin(toBin(memory[m]))+" ["+doc.split("\n")[toDec(padBin(toBin(memory[m]))[:opBits])+1][:parBits][5:][:3]+"]"+" -"
					else:
						outI = "  "+padBin(toBin(m),bit=parBits)+": "+padBin(toBin(memory[m]))+" ["+doc.split("\n")[toDec(padBin(toBin(memory[m]))[:opBits])+1][:parBits][5:][:3]+"]"+"  "
					outII = "  "+padBin(toBin(hiOffset+i),bit=parBits)+": "+padBin(toBin(memory[hiOffset+i]))+"  "
					if i==1:
						outIII = "              - Akku -              "
					elif i==2:
						outIII = " "*8+padBin(toBin(registerA),bit=parBits)+" "*8
					elif i==4:
						outIII = "            - Runtime -             "
					elif i==5:
						outIII = "      "+padBin(toBin(runtime),bit=bits)+"      "
					elif i==7:
						outIII = "            - Command -             "
					elif i==8:
						temp=doc.split("\n")[toDec(padBin(toBin(memory[pointer]))[:opBits])+1][5:]
						outIII = int((36-len(temp)-1)/2)*" "+" "+temp+int((36-len(temp))/2)*" "
					elif i==10:
						outIII = "           - Parameter -            "
					elif i==11:
						outIII = "        "+padBin(toBin(memory[pointer]))[-parBits:]+"        "
					elif i==13:
						outIII = "             - Pointer -            "
					elif i==14:
						outIII = "      "+padBin(toBin(pointer))+"      "
					elif i==16:
						if len(tracker[0]):
							outIII = "            - Trackers -            "
						else:
							outIII = " "*36
					elif i>16 and i<=16+len(tracker):
						if len(tracker[0]):
							outIII = " "+str(tracker[i-16-1]) + ": " + padBin(toBin(memory[toDec(tracker[i-16-1])]))+" "
						else:
							outIII = " "*36
					else:
						outIII = " "*36
					print(" | "+outI+" | "+outII+" |"+outIII+"|")
				print(" |"+" "*58+"|"+" "*52+"|"+" "*36+"|")
				print(" +"+"-"*58+"+"+"-"*52+"+"+"-"*36+"+")

				lastPointer = pointer

				cmd = padBin(toBin(memory[pointer]))[:4]
				par = int(toDec(padBin(toBin(memory[pointer]))[-8:]))
				if parse(cmd,par,verbose=True,adress=pointer)==-1:
					pointer = 0
					break
				pointer = (pointer + 1) % 2**bits

				step=input(" <step> ")
				if step=="x":
					break

				runtime = runtime + 1
		elif x=="asm":
			end = 0
			for m in range(len(memory)):
				if padBin(toBin(memory[len(memory)-m-1]))[:opBits] != "0"*opBits:
					end = len(memory)-m
					break
			toPrint = []
			for m in range(end):
				cmd = doc.split("\n")[toDec(padBin(toBin(memory[m]))[:opBits])+1][:parBits][5:]
				if toDec(padBin(toBin(memory[m]))[opBits:]) >= toDec("1"+"0"*(parBits-1)) and toDec(padBin(toBin(memory[m]))[opBits:]) <= toDec("1"+"0"*(parBits-7)+"101000"):
					toPrint.append(padBin(toBin(m),bit=parBits)+": "+cmd+" "+colored(padBin(toBin(memory[m]))[opBits:],color[(toDec(padBin(toBin(memory[m]))[opBits:])-toDec("1"+"0"*(parBits-1)))][0],color[(toDec(padBin(toBin(memory[m]))[opBits:])-toDec("1"+"0"*(parBits-1)))][1]))
				else:
					toPrint.append(padBin(toBin(m),bit=parBits)+": "+cmd+" "+padBin(toBin(memory[m]))[opBits:])
			origPrint = copy.copy(toPrint)
			for l in range(len(toPrint)):
				cmd = origPrint[l][10:][:3]
				par = origPrint[l][14:]
				if cmd == "JMP" or cmd == "JMN":
					toPrint[toDec(par)] = toPrint[toDec(par)] + " <-+"
					if cmd=="JMP":
						toPrint[l] = toPrint[l] + " --+"
					else:
						toPrint[l] = toPrint[l] + " --?"
					if l < toDec(par):
						for i in range(len(toPrint)):
							if l < i < toDec(par):
								toPrint[i] = toPrint[i] + "   |"
							elif i==toDec(par) or i==l:
								t=1
							else:
								toPrint[i] = toPrint[i] + "    "
					else:
						for i in range(len(toPrint)):
							if toDec(par) < i < l:
								toPrint[i] = toPrint[i] + "   |"
							elif i==toDec(par) or i==l:
								t=1
							else:
								toPrint[i] = toPrint[i] + "    "
			for p in range(len(toPrint)):
				print(toPrint[p])
		elif x=="halting":
			print("[i] Will try to give you information about the divergity of your algorythm")
			print("[i] This algorythm may itself not diverge!")
			bPointer = pointer
			bRegisterA = registerA
			bMemory = memory
			runtime = 0
			stateHashes = []
			print("[*] Simulating...")
			while True:
				runtime+=1
				cmd = padBin(toBin(memory[pointer]))[:opBits]
				par = int(toDec(padBin(toBin(memory[pointer]))[-parBits:]))
				if parse(cmd,par)==-1:
					pointer = 0
					halts = True
					print("[*] Algorythm will diverge after "+str(runtime)+" steps")
					break
				pointer = (pointer + 1) % 2**parBits
				stateHash = hashlib.sha256(str(pointer)+str(registerA)+str(memory)).digest()
				if stateHash in stateHashes:
					print("[*] Algorythm will not diverge: State-Collision after "+str(runtime)+" steps")
					break
				else:
					stateHashes.append(stateHash)
			pointer = bPointer
			registerA = bRegisterA
			memory = bMemory
		elif x=="exit":
			exit()
		elif x=="":
			_=0
		else:
			print("[!] Error")