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 LDA memory > A
0010 ADD A + memspace
0011 SUB A - memspace
0100 STA A > memory
0101 OUT A > OUT
0110 JMP Jump to adress
0111 LDI par > A
1000 JIF Jump if greater
1001 JIE Jump if equal
1010 LDP memory[A] > A
1011 INP IN > A
1100 JMA Jump to A
1101 --- NOP equivalent
1110 STP memory[par] > memory[A]
1111 HLT Halt
"""

asmMap = ["NOP","LDA","ADD","SUB","STA","OUT","JMP","LDI","JIF","JIE","---","INP","JMA","HLO","HST","HLT"]

bits = 12

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**8):
	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):
	raw = raw_input(display)
	if raw=="x":
		return -1
	raw = raw.replace(" ","")
	ret = toDec(raw)
	if ret>2**bit:
		print("[!] input exceted maximal size for "+str(bit)+"-bit integer!")
		return bin_input(display)
	if ret==-1:
		print("[!] "+raw+" is not a valid binary number!")
		return bin_input(display)
	return ret

def parse(cmd,par,verbose=False,adress=0):
	global registerA
	global memory
	global out
	global pointer
	v=verbose
	cmd16 = cmd
	cmd = cmd[-4:]
	if cmd=="0000": #nop
		if v: print padBin(toBin(adress),bit=8)+": NOP"
	elif cmd=="0001": #lda
		registerA = memory[par]
		if v: print padBin(toBin(adress),bit=8)+": LDA memory["+padBin(toBin(par),bit=8)+"] > Register A"
	elif cmd=="0010": #add
		registerA = (registerA + memory[par]) % 2**8
		if v: print padBin(toBin(adress),bit=8)+": ADD Register A + memory["+padBin(toBin(par),bit=8)+"]"
	elif cmd=="0011": #sub
		registerA = (registerA - memory[par]) % 2**8
		if registerA<0:
			registerA = (2**bits+registerA) % 2**8
		if v: print padBin(toBin(adress),bit=8)+": SUB Register A - memory["+padBin(toBin(par),bit=8)+"]"
	elif cmd=="0100": #sta
		memory[par] = registerA
		if v: print padBin(toBin(adress),bit=8)+": STA Register A > memory["+padBin(toBin(par),bit=8)+"]"
	elif cmd=="0101": #out
		out = registerA
		print("[OUT] "+padBin(toBin(out),bit=8))
		if v: print padBin(toBin(adress),bit=8)+": OUT Register A > out"
	elif cmd=="0110": #jmp
		pointer = par - 1
		if v: print padBin(toBin(adress),bit=8)+": JMP "+padBin(toBin(par),bit=8)
	elif cmd=="0111": #ldi
		registerA = par
		if v: print padBin(toBin(adress),bit=8)+": LDI "+padBin(toBin(par),bit=8)+" > Register A"
	elif cmd=="1000": #jif
		if registerA>=(2**8)/2: #10000000
			pointer = par - 1
			if v: print padBin(toBin(adress),bit=8)+": JIF {TRUE} > "+padBin(toBin(par),bit=8)
		else:
			if v: print padBin(toBin(adress),bit=8)+": JIF {FALSE} > "+padBin(toBin(par),bit=8)
	elif cmd=="1001": #fie
		if registerA==(2**8)/2: #10000000
			pointer = par - 1
			if v: print padBin(toBin(adress),bit=8)+": JIE {TRUE} > "+padBin(toBin(par),bit=8)
		else:
			if v: print padBin(toBin(adress),bit=8)+": JIE {FALSE} > "+padBin(toBin(par),bit=8)
	elif cmd=="1010": #ldp
		registerA = memory[registerA]
		if v: print padBin(toBin(adress),bit=8)+": LDP memory["+padBin(toBin(registerA),bit=8)+"] > Register A"
	elif cmd=="1011": #inp
		if v: print padBin(toBin(adress),bit=8)+": INP input > Register A"
		registerA = bin_input("[INPUT]> ")
	elif cmd=="1100": #jma
		pointer = registerA - 1
		if v: print padBin(toBin(adress),bit=8)+": JMA Jump to Register A"
	elif cmd=="1101": #---
		if v: print "Got unimplemented op-code: 1101"
	elif cmd=="1110": #STP
		memory[par] = memory[registerA]
		if v: print padBin(toBin(adress),bit=8)+": STP memory["+padBin(toBin(registerA),bit=8)+"] > memory["+padBin(toBin(par),bit=8)+"]"
	elif cmd=="1111": #hlt
		return -1
	return True

if __name__=="__main__":
  print("[asmTron by CyberCodeDragon]")
  while True:
	x = ""
	x=raw_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("highClear [hc]    - clear highMem,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")
		print("asm               - shellcode visualizer (assembly / jump-arrows)")
		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=8)+": "+padBin(toBin(memory[y+m])))
	elif x=="run" or x=="u":
		while True:
			cmd = padBin(toBin(memory[pointer]))[:4]
			par = int(toDec(padBin(toBin(memory[pointer]))[-8:]))
			if parse(cmd,par)==-1:
				pointer = 0
				break
			pointer = (pointer + 1) % 2**bits
	elif x=="verboseRun" or x=="vu":
		while True:
			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
	elif x=="edit" or x=="e":
		adress=bin_input("offset> ",bit=8)
		while True:
			value=bin_input(padBin(toBin(adress),bit=8)+": "+padBin(toBin(memory[adress]))+" > ")
			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**bits):
			memory.append(0)
	elif x=="highClear" or x=="highReset" or x=="hr" or x=="hc":
		pointer = 0
		registerA = 0
		out = 0
		for i in range((2**bits)/2):
			memory[(2**bits)/2+i] = 0
	elif x=="step" or x=="s":
		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
			print("[*] Programm halted")
		pointer = (pointer + 1) % (2**bits)
	elif x=="debug":
		while True:
			cmd = padBin(toBin(memory[pointer]))[:4]
			par = int(toDec(padBin(toBin(memory[pointer]))[-8:]))
			if parse(cmd,par,verbose=True,adress=pointer)==-1:
				pointer = -1
				print("[*] Programm halted")
				break
			pointer = (pointer + 1) % (2**bits)
	elif x=="pointer":
		print("POINTER: "+padBin(toBin(pointer),bit=8))
	elif x=="setPointer":
		if y=="":
			pointer = bin_input("pointer> ",bit=8)
		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("."):
				if f[-8:]==".asmTron":
					print("> "+f[:-8])
			name = raw_input("name> ")
		else:
			name = y
		with open(name+".asmTron", 'wb') as f:
			t=pickle.dumps([12,memory])
			t=zlib.compress(t)
			f.write(t)
	elif x=="load" or x=="l":
		if y=="":
			print("[*] Avaible saves:")
			for f in os.listdir("."):
				if f[-8:]==".asmTron":
					print("> "+f[:-8])
			for f in os.listdir("tronScript"):
				if f[-8:]==".asmTron":
					print("> tronScript/"+f[:-8])
			name = raw_input("name> ")
		else:
			name = y
		try:
			with open(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==12:
			print("[!] Error: Invalid Magical Bit: "+str(magicBit))
			exit()
		print("[*] loaded "+name+".asmTron")
	elif x=="code" or x=="exec":
		if y=="":
			data=bin_input("code> ")
		else:
			data=y
		cmd = padBin(toBin(memory[pointer]))[:4]
		par = int(toDec(padBin(toBin(memory[pointer]))[-8:]))
		parse(cmd,par,verbose=True,adress=0)
	elif x=="guiDebug" or x=="gd":
		runtime = 0
		tracker = raw_input("trackers> ").split(",")
		running = True
		lastPointer = 0
		while running:
			clear()
			print(" +"+"-"*34+"+"+"-"*28+"+"+"-"*24+"+")
			print(" |"+" "*8+" - Instructions - "+" "*8+"|"+" "*8+" - Hi-mem - "+" "*8+"|"+" "*24+"|")
			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=8)+": "+padBin(toBin(memory[m]))+" ["+doc.split("\n")[toDec(padBin(toBin(memory[m]))[:4])+1][:8][5:]+"]"+" <"
				elif lastPointer == m and pointer-1 != lastPointer:
					outI = "- "+padBin(toBin(m),bit=8)+": "+padBin(toBin(memory[m]))+" ["+doc.split("\n")[toDec(padBin(toBin(memory[m]))[:4])+1][:8][5:]+"]"+" -"
				else:
					outI = "  "+padBin(toBin(m),bit=8)+": "+padBin(toBin(memory[m]))+" ["+doc.split("\n")[toDec(padBin(toBin(memory[m]))[:4])+1][:8][5:]+"]"+"  "
				outII = "  "+padBin(toBin(128+i),bit=8)+": "+padBin(toBin(memory[128+i]))+"  "
				if i==1:
					outIII = "     - Register A -     "
				elif i==2:
					outIII = " "*8+padBin(toBin(registerA),bit=8)+" "*8
				elif i==4:
					outIII = "      - Runtime -       "
				elif i==5:
					outIII = "      "+padBin(toBin(runtime),bit=12)+"      "
				elif i==7:
					outIII = "      - Command -       "
				elif i==8:
					temp=doc.split("\n")[toDec(padBin(toBin(memory[pointer]))[:4])+1][5:]
					outIII = int((24-len(temp))/2-1)*" "+" "+temp+int((24-len(temp))/2)*" "
				elif i==10:
					outIII = "     - Parameter -      "
				elif i==11:
					outIII = "        "+padBin(toBin(memory[pointer]))[-8:]+"        "
				elif i==13:
					outIII = "    - OUT-Register -    "
				elif i==14:
					outIII = "        "+padBin(toBin(out),bit=8)+"        "
				elif i==16:
					outIII = "       - Pointer -      "
				elif i==17:
					outIII = "      "+padBin(toBin(pointer))+"      "
				elif i==19:
					if len(tracker[0]):
						outIII = "      - Trackers -      "
					else:
						outIII = " "*24
				elif i>19 and i<=19+len(tracker):
					if len(tracker[0]):
						outIII = " "+str(tracker[i-19-1]) + ": " + padBin(toBin(memory[toDec(tracker[i-19-1])]))+" "
					else:
						outIII = " "*24
				else:
					outIII = " "*24
				print(" | "+outI+" | "+outII+" |"+outIII+"|")
			print(" |"+" "*34+"|"+" "*28+"|"+" "*24+"|")
			print(" +"+"-"*34+"+"+"-"*28+"+"+"-"*24+"+")

			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=raw_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]))[:4] != "0000":
				end = len(memory)-m
				break
		toPrint = []
		for m in range(end):
			cmd = doc.split("\n")[toDec(padBin(toBin(memory[m]))[:4])+1][:8][5:]
			if toDec(padBin(toBin(memory[m]))[4:]) >= toDec("10000000") and toDec(padBin(toBin(memory[m]))[4:]) <= toDec("10101000"):
				toPrint.append(padBin(toBin(m),bit=8)+": "+cmd+" "+colored(padBin(toBin(memory[m]))[4:],color[(toDec(padBin(toBin(memory[m]))[4:])-toDec("10000000"))][0],color[(toDec(padBin(toBin(memory[m]))[4:])-toDec("10000000"))][1]))
			else:
				toPrint.append(padBin(toBin(m),bit=8)+": "+cmd+" "+padBin(toBin(memory[m]))[4:])
		origPrint = copy.copy(toPrint)
		for l in range(len(toPrint)):
			cmd = origPrint[l][10:][:3]
			par = origPrint[l][14:]
			if cmd == "JMP" or cmd == "JIF" or cmd == "JIE":
				toPrint[toDec(par)] = toPrint[toDec(par)] + " <-+"
				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]))[:4]
			par = int(toDec(padBin(toBin(memory[pointer]))[-8:]))
			if parse(cmd,par)==-1:
				pointer = 0
				halts = True
				print("[*] Algorythm will diverge after "+str(runtime)+" steps")
				break
			pointer = (pointer + 1) % 2**bits
			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")