#!/usr/bin/env python3
import sys

oldexcepthook = sys.excepthook
def newexcepthook(type,value,traceback):
	oldexcepthook(type,value,traceback)
	input("Press ENTER to quit.")
sys.excepthook = newexcepthook

import os
def p(*args):
	args = list(args)
	index = 0
	length = len(args)
	while index < length:
		arg = args[index]
		arg = arg.replace("/",os.path.sep)
		arg = arg.replace("\\",os.path.sep)
		if index != 0: arg = arg.lstrip(os.path.sep)
		arg = arg.rstrip(os.path.sep)
		while os.path.sep + os.path.sep in arg:
			arg = arg.replace(os.path.sep + os.path.sep,os.path.sep)
		
		if arg == "":
			del args[index]
			length -= 1
			continue
		
		args[index] = arg
		index += 1
	
	return os.path.sep.join(args)
pUp = os.path.dirname
s = False
if getattr(sys, 'frozen', False) and hasattr(sys, '_MEIPASS'):
	s = os.path.realpath(sys.executable)
else:
	s = os.path.realpath(__file__)
sp = pUp(s)
sys.path = [sp] + sys.path

# script start
version = (0,0,0,"alpha")
import api.opusnt as opusnt
import json
import subprocess
import winreg

class coloramaFallback:
	class Fore:
		BLACK = ""
		RED = ""
		GREEN = ""
		YELLOW = ""
		BLUE = ""
		MAGENTA = ""
		CYAN = ""
		WHITE = ""
		RESET = ""
	
	class Back:
		BLACK = ""
		RED = ""
		GREEN = ""
		YELLOW = ""
		BLUE = ""
		MAGENTA = ""
		CYAN = ""
		WHITE = ""
		RESET = ""
	
	class Style:
		DIM = ""
		NORMAL = ""
		BRIGHT = ""
		RESET_ALL = ""

def coloramaInit():
	global colorama
	coloramaSuccess = False
	if not "-nocolor" in sys.argv:
		try:
			import colorama
			colorama.init()
			coloramaSuccess = True
		except Exception as e:
			print("Could not import/init colorama: " +str(e),file=sys.stderr)
			print("Colors deactivated.\n",file=sys.stderr)
	
	if not coloramaSuccess:
		colorama = coloramaFallback
coloramaInit()

def runCode(str, lcs = False, description = "loose-code"):
	if lcs == False: lcs = {}
	code = compile(str,"code","exec")
	exec(code,globals(),lcs)
	return lcs

def runScript(sf, lcs = False):
	if lcs == False: lcs = {}
	with open(sf) as script:
		runCode(script.read(),lcs,sf)
		return lcs

def prettyJson(ls):
	return json.dumps(ls,indent=2)

def printStatus(st,color = colorama.Fore.YELLOW):
	print(color + colorama.Style.BRIGHT + st + colorama.Style.RESET_ALL)

def main():
	arglist = {}
	for arg in sys.argv[1:]:
		splitarg = arg.split("=",1)
		if len(splitarg) < 2: splitarg.append("")
		
		while len(splitarg[0]) > 0 and splitarg[0][0] in [" ","	"]: splitarg[0] = splitarg[0][1:]
		while len(splitarg[0]) > 0 and splitarg[0][-1] in [" ","	"]: splitarg[0] = splitarg[0][:-1]
		while len(splitarg[1]) > 0 and splitarg[1][0] in [" ","	"]: splitarg[1] = splitarg[1][1:]
		while len(splitarg[1]) > 0 and splitarg[1][-1] in [" ","	"]: splitarg[1] = splitarg[1][:-1]
		
		arglist[splitarg[0]] = splitarg[1]
	
	printStatus("\narguments:")
	print(prettyJson(arglist)+ "\n")
	
	if "wim" in arglist:
		opusnt.target["type"] = "wim"
		opusnt.target["path"] = arglist["wim"]
		if arglist["index"] != "*":
			opusnt.target["index"] = int(arglist["index"])
			opusnt.target["maxIndex"] = int(arglist["index"])
		else:
			opusnt.target["index"] = 1
			opusnt.getMaxIndex()
	elif "online" in arglist:
		opusnt.target["type"] = "online"
		opusnt.target["path"] = os.environ["SystemDrive"]
		opusnt.target["regmap"] = {}
	else:
		opusnt.target["type"] = "offline"
		opusnt.target["path"] = arglist["path"]
	
	opusnt.tmpPath = p(sp,"tmp." +str(os.getpid()))
	if "tmp" in arglist: opusnt.tmpPath = arglist["tmp"]
	if os.path.isdir(opusnt.tmpPath):
		subprocess.run(["rmdir","/s","/q",opusnt.tmpPath],check=True,shell=True)
	
	modpath = p(sp,"mods")
	if "modpath" in arglist: modpath = arglist["modpath"]
	
	while opusnt.target["index"] <= opusnt.target["maxIndex"]:
		os.makedirs(opusnt.tmpPath)
		
		regs = {
			"software": p("Windows","System32","config","SOFTWARE"),
			"system": p("Windows","System32","config","SYSTEM"),
			"user-SYSTEM": p("Windows","System32","config","DEFAULT"),
			"user-SYSTEM2": p("Windows","System32","config","systemprofile","ntuser.dat")
		}
		
		printStatus("\nextracting registry hives (" +str(opusnt.target["index"])+ "/" +str(opusnt.target["maxIndex"])+ ")...")
		for reg in list(regs.keys()):
			try:
				opusnt.readyFile(regs[reg])
			except:
				print(colorama.Fore.YELLOW + colorama.Style.BRIGHT + "warning: " +colorama.Style.RESET_ALL+ " extracting " +regs[reg]+ " failed, skipping.")
				del regs[reg]
		
		userPath = "Users"
		opusnt.readyFile(userPath)
		userPath = opusnt.filePath(userPath)
		for root,dirs,files in os.walk(userPath):
			for file in dirs:
				ffile = p(userPath,file)
				if opusnt.target["type"] in ["offline","online"]:
					tfile = ffile.replace(opusnt.target["path"] + os.path.sep,"",1)
				else:
					tfile = ffile.replace(opusnt.tmpPath + os.path.sep,"",1)
				user = ffile.replace(userPath + os.path.sep,"",1)
				if os.path.isfile(p(ffile,"ntuser.dat")):
					regs["user-" +user] = p(tfile,"ntuser.dat")
			break
		
		printStatus("mounting registry hives...")
		regNames = []
		for reg in regs:
			regNames.append(reg)
		for reg in regNames:
			opusnt.mountReg(opusnt.regTmpPath + reg, opusnt.filePath(regs[reg]),regs,reg)
		regNames = None
		
		if opusnt.target["type"] == "online":
			opusnt.target["regmap"][opusnt.regTmpPath + "software"] = "HKEY_LOCAL_MACHINE\\SOFTWARE"
			opusnt.target["regmap"][opusnt.regTmpPath + "system"] = "HKEY_LOCAL_MACHINE\\SYSTEM"
			for key in opusnt.regQueryKeys("HKEY_USERS"):
				if key.endswith("_Classes"): continue
				regs["user-mounted_" +key.split("\\")[1]] = ""
				opusnt.target["regmap"][opusnt.regTmpPath + "user-mounted_" +key.split("\\")[1]] = key
		
		opusnt.getVersion(winreg.HKEY_LOCAL_MACHINE,(opusnt.regTmpPath + "software").split("\\",1)[-1])
		printStatus("\nwindows version:")
		print(prettyJson(opusnt.target["version"]) + "\n")
		
		modlist = opusnt.getModlist(modpath)
		for mod in modlist:
			lmod = mod.replace(modpath + os.sep,"",1)
			print(colorama.Fore.YELLOW + colorama.Style.BRIGHT+ "applying mod: " + colorama.Style.RESET_ALL + lmod,end="")
			
			# process conditional scripts
			skip = False
			lmodSplit = lmod.split(os.sep)
			currentCondSplit = ""
			for split in lmodSplit:
				if currentCondSplit == "":
					currentCondSplit = split
				else:
					currentCondSplit = p(currentCondSplit,split)
				
				condFile = p(modpath,currentCondSplit,"condition.py")
				if os.path.isfile(condFile):
					lcs = locals()
					runScript(condFile,lcs)
					if "skip" in lcs: skip = lcs["skip"]
			
			if skip: print(colorama.Fore.RED + colorama.Style.BRIGHT + " - skipping (conditional)" +colorama.Style.RESET_ALL); continue
			print("")
			
			# run modscript
			dataPass = True
			modscript = p(mod,"modscript.py")
			if os.path.isfile(modscript):
				lcs = locals()
				runScript(modscript,lcs)
				if "dataPass" in lcs: dataPass = lcs["dataPass"]
			
			# add files
			if dataPass == False: continue
			if not os.path.isdir(p(mod,"data")): continue
			opusnt.addFiles(p(mod,"data"))
		
		printStatus("unmounting registry hives...")
		for reg in regs:
			opusnt.unmountReg(opusnt.regTmpPath + reg)
		
		if opusnt.target["type"] == "wim":
			printStatus("updating wim...")
			opusnt.addFiles(opusnt.tmpPath)
		
		printStatus("deleting temporary folder...")
		subprocess.run(["rmdir","/s","/q",opusnt.tmpPath],check=True,shell=True)
		opusnt.target["index"] += 1

if __name__ == "__main__":
	main()