opus-nt/api/opusnt.py

#!/usr/bin/env python3
dummyMode = False
if __name__ == "__main__":
	print("this is supposed to be imported")
	exit()

print("opus-nt api says hello world")
import os
import subprocess
import winreg
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)

target = {
	"type": False, # "wim", "offline" - whether to target a wim file or an offline installed windows
	"path": False, # the path to the wim or windows installation
	"index": 1, # the index of the wim to target, if a wim is selected
	"maxIndex": 1
}

tmpPath = "tmp"
regTmpPath = "HKEY_LOCAL_MACHINE\\opus-nt_"

def runReg(cmd,check = True):
	if target["type"] == "online":
		cmd[1] = applyRegmap(cmd[1])
		if cmd[0] == "unload":
			if not cmd[1].startswith(regTmpPath): return
	cmd = ["reg"] + cmd
	try:
		if dummyMode:
			print(cmd)
		else:
			subprocess.run(cmd,stdout=subprocess.DEVNULL,check=True)
	except Exception as e:
		if check == True: raise e

def getModlist(path):
	modList = []
	for root,dirs,files in os.walk(path):
		for file in dirs:
			ffile = p(root,file)
			lfile = ffile.replace(path + os.path.sep,"",1)
			if lfile[0] == "-": continue
			if lfile[0] == "[" and lfile[-1] == "]":
				modList = modList + sorted(getModlist(ffile))
				continue
			
			modList.append(ffile)
		break
	
	return modList

def addFiles(path):
	if dummyMode: return
	if target["type"] == "offline":
		for root,dirs,files in os.walk(path):
			for file in dirs:
				ffile = p(root,file)
				lfile = ffile.replace(path + os.path.sep,"",1)
				ofile = p(target["path"],lfile)
				if not os.path.isdir(ofile): os.makedirs(ofile)
			
			for file in files:
				ffile = p(root,file)
				lfile = ffile.replace(path + os.path.sep,"",1)
				ofile = p(target["path"],lfile)
				if os.path.isfile(ofile): os.remove(ofile)
				fhandle = open(ffile,"rb")
				ofhandle = open(ofile,"wb")
				bytes = fhandle.read(1000000)
				while bytes:
					ofhandle.write(bytes)
					bytes = fhandle.read(1000000)
				
				ofhandle.close()
				fhandle.close()
	
	if target["type"] == "wim":
		subprocess.run([
			"wimlib-imagex","update",
			target["path"],str(target["index"]),
			"--command=add '" +path+ "' '\\'"
		],check=True,stdout=subprocess.DEVNULL)

def readyFile(path):
	if target["type"] in ["offline","online"]: return
	
	if target["type"] == "wim":
		subprocess.run([
			"wimlib-imagex","extract",
			target["path"],str(target["index"]),
			path,"--dest-dir=" +tmpPath,
			"--preserve-dir-structure"
		],check=True,stdout=subprocess.DEVNULL)

def renameFile(path,newpath):
	if dummyMode: return
	if target["type"] == "offline":
		os.rename(filePath(path),filePath(newpath))
	
	if target["type"] == "wim":
		subprocess.run([
			"wimlib-imagex","update",
			target["path"],str(target["index"]),
			"--command=rename '" +path+ "' '" +newpath+ "'"
		],check=True,stdout=subprocess.DEVNULL)

def removeFile(path):
	if dummyMode: return
	if target["type"] == "offline":
		os.remove(filePath(path))
	
	if target["type"] == "wim":
		subprocess.run([
			"wimlib-imagex","update",
			target["path"],str(target["index"]),
			"--command=delete '" +path+ "'"
		],check=True)

def fileExists(path):
	if target["type"] == "offline":
		return os.path.isfile(filePath(path))
	
	# This is super super super dirty and dumb, fix this somehow
	if target["type"] == "wim":
		trashpath = p(tmpPath,"opus-nt_trash")
		if not os.path.isdir(trashpath): os.makedirs(trashpath)
		try:
			subprocess.run([
				"wimlib-imagex","extract",
				target["path"],str(target["index"]),
				path,"--dest-dir=" +trashpath,
				"--preserve-dir-structure"
			],check=True,stdout=subprocess.DEVNULL)
		except:
			subprocess.run(["rmdir","/s","/q",trashpath],check=True,shell=True)
			return False
		subprocess.run(["rmdir","/s","/q",trashpath],check=True,shell=True)
		return True

def filePath(path):
	if target["type"] in ["offline","online"]: return p(target["path"],path)
	if target["type"] == "wim": return p(tmpPath,path)

def mountReg(path,regFile,regList,reg):
	if target["type"] == "online":
		try:
			runReg(["load",path,regFile])
		except Exception:
			if reg.startswith("user-"):
				del regList[reg]
		return
	runReg(["load",path,regFile])

def unmountReg(path):
	runReg(["unload",path])

def getVersion(hkey,SOFTWARE):
	versionData = {}
	if target["type"] == "online": SOFTWARE = "software"
	key = winreg.OpenKey(hkey,SOFTWARE + "\\Microsoft\\Windows NT\\CurrentVersion")
	
	isTen = False
	try:
		winreg.QueryValueEx(key,"CurrentMajorVersionNumber")[0]
		isTen = True
	except:
		pass
	
	if isTen == False:
		version = winreg.QueryValueEx(key,"CurrentVersion")[0].split(".")
		versionData["versionMajor"] = int(version[0])
		versionData["versionMinor"] = int(version[1])
	
	if isTen == True:
		versionData["versionMajor"] = int(winreg.QueryValueEx(key,"CurrentMajorVersionNumber")[0])
		versionData["versionMinor"] = int(winreg.QueryValueEx(key,"CurrentMinorVersionNumber")[0])
	
	versionData["edition"] = winreg.QueryValueEx(key,"EditionID")[0]
	versionData["build"] = int(winreg.QueryValueEx(key,"CurrentBuildNumber")[0])
	buildLabEx = winreg.QueryValueEx(key,"BuildLabEx")[0]
	if "x86" in buildLabEx:
		versionData["architecture"] = "x86"
	elif "amd64" in buildLabEx:
		versionData["architecture"] = "amd64"
	else: # this sucks
		versionData["architecture"] = "ia64"
	
	target["version"] = versionData

def getMaxIndex():
	subprocess.run([
		"wimlib-imagex","info",
		target["path"],"1"
	],check=True,stdout=subprocess.DEVNULL)
	
	ind = 2
	while True:
		try:
			subprocess.run([
				"wimlib-imagex","info",
				target["path"],str(ind)
			],check=True,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL)
			ind += 1
		except Exception as e:
			target["maxIndex"] = ind - 1
			return

def regQueryKeys(rp):
	if target["type"] == "online": rp = applyRegmap(rp)
	output = subprocess.check_output(["reg","query",rp,"/f","*","/k"]).decode("utf-8","ignore")
	output = output.split("\r\n")
	output = output[1:-2]
	return output

def applyRegmap(rp):
	for path in target["regmap"]:
		if rp.startswith(path):
			return rp.replace(path,target["regmap"][path],1)
	return rp