import sys
from chardet.universaldetector import UniversalDetector


def parse_args(INFO):

	MENU = { \

			"OBF":"Obfuscated algorithm file", \
			"OF":"Output file (Not implemented)"

		}

	for NUM in range(1,len(sys.argv)):
		CURR = sys.argv[NUM]

		if CURR[:2] == "--":
			if CURR[2:] == "help":
				print_dict(MENU)
				sys.exit()

		KEY = CURR.split("=")[0]

		if KEY in INFO.keys():
			INFO[KEY] = CURR[(len(KEY)+1):]

	return INFO

def parse_obfs(OBFS,IFILE):
	for LINE in IFILE:
		SFILE = LINE.split(":")[0].split(",")
		LO = LINE[(len(LINE.split(":")[0])+1):]
		KEY = LO.replace(" ","").split("=")[0]
		PCS = LO.replace(" ","").replace("\n","").split("=")[1].split("+")
		OBFS[KEY] = { \
					"SFILE":SFILE, \
					"PCS":[], \
					"ALG": {}, \
					"FULLSTR":"", \
					"FILE":""
				}

		if "Array" in LINE:
			OBFS[KEY]["PCS"].append("ARR")

			for PC in LINE.replace(" ","").replace(")","").replace("\n","").split("(")[1].split(","):
				OBFS[KEY]["PCS"].append(PC)

		else:
			for PC in LO.replace(" ","").replace("\n","").split("=")[1].split("+"):
				OBFS[KEY]["PCS"].append(PC)

	return OBFS

def parse_files(OBFS):

	for KEY,DATA in OBFS.items():

		DET = UniversalDetector()
		DET.reset()

		for FNAME in DATA["SFILE"]:

			with open(FNAME,"rb") as FILE:
				for LINE in FILE:
					DET.feed(LINE)
					if DET.done:
						break
			DET.close()

			with open(FNAME,"r",encoding=DET.result["encoding"]) as FILE:
				for LINE in FILE:
					DATA["FILE"] += LINE

	return OBFS

def begin_deobf(OBFS):
	for KEY,DATA in OBFS.items():

		FILE = DATA["FILE"].split("\n")

		for LINE in FILE:

			BEGIN = LINE.replace(" ","").split("=")[0]

			if BEGIN in DATA["PCS"]:
				TYPE = LINE.replace(" ","").split("=")[1].split("(")[0]

				if TYPE == "Mid":
					OSTR = LINE.replace(" ","").replace("\n","").replace(")","").split("(")[1].split(",")[0]
					ENTRY = int(LINE.replace(" ","").replace("\n","").replace(")","").split("(")[1].split(",")[1])
					LEN = int(LINE.replace(" ","").replace("\n","").replace(")","").split("(")[1].split(",")[2])
					DATA["ALG"][BEGIN] = { \

									"TYPE":TYPE, \
									"OSTR":OSTR, \
									"ENTRY":ENTRY, \
									"LEN":LEN, \
									"STR":""

									}

			elif all(VAL in LINE for VAL in (KEY,"(",")")) and not LINE.startswith(KEY) and not LINE.startswith("Dim") and not LINE.startswith("Function"):

				NAME = LINE.replace(" ","").split("(")[0].split(",")[1]
				
				DATA["ALG"]["FUNC"] = { \

										"NAME":NAME, \
										"CONT":[]

									}

				

	return OBFS

def construct_deobf(OBFS):

	for KEY,DATA in OBFS.items():
		FILE = DATA["FILE"].split("\n")

		if "ARR" in DATA["PCS"]:
			IS_FUNC = False
			for CHUNK in FILE:
				if "Function "+DATA["ALG"]["FUNC"]["NAME"] in CHUNK and not IS_FUNC:
					DATA["FULLSTR"] += "FUNCTION CONTAINING ALGORITHM\n"
					DATA["FULLSTR"] += CHUNK+"\n"
					DATA["ALG"]["FUNC"]["CONT"].append(CHUNK)
					IS_FUNC = True
				elif IS_FUNC and "End Function" in CHUNK:
					IS_FUNC = False
				elif IS_FUNC: 
					DATA["ALG"]["FUNC"]["CONT"].append(CHUNK)
					DATA["FULLSTR"] += CHUNK+"\n"
				elif DATA["ALG"]["FUNC"]["NAME"] in CHUNK:
					DATA["FULLSTR"] += "LINE CONTAINING VARIABLE\n"				
					DATA["FULLSTR"] += CHUNK+"\n\n"


		else:

			for PC1 in DATA["ALG"].keys():

				PC2 = DATA["ALG"][PC1]["OSTR"]

				for CHUNK in FILE:
					if CHUNK.replace(" ","").startswith(PC2):

						if DATA["ALG"][PC1]["TYPE"] == "Mid":
							ENTRY = DATA["ALG"][PC1]["ENTRY"]
							LEN = DATA["ALG"][PC1]["LEN"]

							FOUND_STR = CHUNK.replace('"','').split("=")[1][ENTRY:(ENTRY+LEN)]
							DATA["ALG"][PC1]["STR"] = FOUND_STR

	
			for PC in DATA["PCS"]:
				if PC in OBFS.keys() or PC.startswith("Chr("):
					DATA["FULLSTR"] += "<<"+PC+">>    "
					continue
				DATA["FULLSTR"] += DATA["ALG"][PC]["STR"]

	return OBFS

def print_results(OBFS):
	for KEY,DATA in OBFS.items():

		print("      Obfuscated Variable:     "+KEY)
		print("              Source File:     "+str(DATA["SFILE"]))
		print("     De-obfuscated String\n------------------------------------------------")

		print(OBFS[KEY]["FULLSTR"])
		print("\n\n")

def print_dict(INFO):
	for KEY,DATA in INFO.items():
		print("  "+KEY+" "*(10-len(KEY))+"		"+str(DATA))

def main():

	files = {}

	cliArgs = { \

			"IF":"", \
			"OF":""

		 }

	cliArgs = parse_args(cliArgs)

	if cliArgs["OF"]:
		outFile = open(cliArgs["OF"], "w")

	if cliArgs["IF"]:
		obfFile = open(cliArgs["IF"], "r")

		obfStrs = {}

		obfStrs = parse_obfs(obfStrs,obfFile)

		obfStrs = parse_files(obfStrs)

		obfStrs = begin_deobf(obfStrs)

		obfStrs = construct_deobf(obfStrs)

		if "TEST" in sys.argv:
			sys.exit()

		print_results(obfStrs)

if __name__ == "__main__":
	main()
