#!/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
p = os.path.join
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)

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

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()
	else:
		opusnt.target["type"] = "offline"
		opusnt.target["path"] = arglist["path"]
	
	opusnt.tmpPath = p(sp,"tmp")
	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 regs:
			opusnt.readyFile(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)
				user = ffile.replace(userPath + os.path.sep,"",1)
				if os.path.isfile(p(ffile,"ntuser.dat")):
					regs["user-" +user] = p(ffile,"ntuser.dat")
			break
		
		printStatus("mounting registry hives...")
		for reg in regs:
			opusnt.mountReg(opusnt.regTmpPath + reg, opusnt.filePath(regs[reg]))
		
		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()