utage: encryption/decryption tool

adx/adx.py

@@ -8,6 +8,8 @@ def extract_loop_data_from_dir(dir_in, fout):
 	if fout is None:
 		fout = "BgmLoop.json"
 	files = glob.glob(os.path.join(dir_in, "**/*.adx"), recursive=True)
+	if len(files) == 0:
+		raise FileNotFoundError("No valid files found in directory: " + dir_in)
 	extract_loop_data_from_files(files, fout)
 
 def extract_loop_data_from_files(files_in, fout):

divatool.py

@@ -2,7 +2,10 @@
 
 import argparse
 import sys
+from os import path
+
 from adx import adx
+from utage import utage
 
 def main():
 	parser = argparse.ArgumentParser(description="Unpack and process XDU data")
@@ -21,11 +24,43 @@ def main():
 	                             help="Input directory containing ADX files",
 	                             type=str)
 
+	# utage subcommand
+	parser_utage = subparsers.add_parser("utage", help="UTAGE stuff")
+	parser_utage_subparsers = parser_utage.add_subparsers(metavar="<command>", title="subcommand", dest="subcommand")
+	# utage crypt
+	parser_utage_crypt = parser_utage_subparsers.add_parser("crypt", help="Encrypt/decrypt utage tsv files")
+	parser_utage_crypt.add_argument("--encrypt", "-e",
+	                                help="Encrypt (default: Decrypt)",
+					dest="encrypt", action="store_true", default=False)
+	parser_utage_crypt.add_argument("--no-compression", "-n",
+	                                help="Do not compress/decompress. Only applies to tsv, png will never be compressed",
+	                                dest="ncomp", action="store_true", default=False)
+	parser_utage_crypt.add_argument("--key", "-k",
+	                                help="Encryption key (Default: SampleSecretKey)",
+	                                dest="key", type=str, default="SampleSecretKey")
+	parser_utage_crypt.add_argument("--hex", "-x",
+	                                help="KEY is hexadecimal (Default: False)",
+	                                dest="hex", action="store_true", default=False)
+	parser_utage_crypt.add_argument("TARGET",
+	                                help="Input", type=str)
+
 	args = parser.parse_args()
 
 	if args.command == "adx":
 		if args.subcommand == "loop":
 			adx.extract_loop_data_from_dir(args.ADX_DIR, args.JSON_OUT)
+	elif args.command == "utage":
+		if args.subcommand == "crypt":
+			if not args.hex:
+				key = bytearray(args.key, "utf-8")
+			else:
+				key = bytearray.fromhex(args.key)
+			if path.isfile(args.TARGET):
+				utage.crypt_file(args.TARGET, key, args.encrypt, args.ncomp)
+			elif path.isdir(args.TARGET):
+				utage.crypt_dir(args.TARGET, key, args.encrypt, args.ncomp)
+			else:
+				raise FileNotFoundError("Could not find {}".format(args.TARGET))
 
 	return 0
 

utage/crypt.py

@@ -0,0 +1,159 @@
+def decompress(data):
+	osize = 0
+	isize = len(data)
+	osize_expected = int.from_bytes(data[:4], byteorder='little', signed=False)
+	odata = bytearray(osize_expected)
+	odata[:isize] = data
+	i = 4
+	while i < isize:
+		if (data[i] & 128) != 0:
+			num3 = data[i] & 15
+			num3 += 3
+			num4 = (data[i] & 112) << 4 | data[i+1]
+			num4 += 1
+			for j in range(0, num3):
+				odata[osize + j] = odata[osize - num4 + j]
+			i += 1
+		else:
+			num3 = data[i] + 1
+			for j in range(0, num3):
+				odata[osize + j] = data[i + 1 + j]
+			i += num3
+		osize += num3
+		i += 1
+	return odata
+
+def compress(data):
+	num = len(data)
+	bytes2 = num.to_bytes(4, byteorder='little')
+	anum3 = num + num/128 + 1
+	array = bytearray(int(anum3))
+	num2 = 0
+	num3 = 0
+	num4 = 0
+	index = Index()
+	while num3 < num:
+		num5 = 0
+		num6 = 0
+		num7 = min(18, num - num3)
+		num8 = index.getFirst(data[num3])
+		while not index.isEnd(num8):
+			node = index.getNode(num8)
+			mPos = node.mPos
+			i = 1
+			while i < num7:
+				if data[mPos+i] != data[num3+i]:
+					break
+				i = i+1
+			if num5 < i:
+				num6 = mPos
+				num5 = i
+				if num5 == num7:
+					break
+			num8 = node.mNext
+		if num5 >= 3:
+			for j in range(num5):
+				num9 = num3 + j - 2048
+				if num9 >= 0:
+					index.remove(data[num9], num9)
+				index.add(data[num3+j], num3+j)
+			if num4 < num3:
+				array[num2] = (num3 - num4 - 1)
+				num2 = num2+1
+				for j in range(num4, num3):
+					array[num2] = data[j]
+					num2 = num2+1
+			num10 = num5 - 3
+			num11 = num3 - num6 - 1
+			num12 = 0x80 | num10
+			num12 |= (num11 & 0x700) >> 4
+			array[num2] = num12
+			array[num2+1] = (num11&0xff)
+			num2 = num2+2
+			num3 = num3 + num5
+			num4 = num3
+		else:
+			num9 = num3 - 2048
+			if num9 >= 0:
+				index.remove(data[num9],num9)
+			index.add(data[num3], num3)
+			num3 = num3+1
+			if num3 - num4 == 128:
+				array[num2] = num3 - num4 - 1
+				num2 = num2+1
+				for j in range(num4, num3):
+					array[num2] = data[j]
+					num2 = num2+1
+				num4 = num3
+	if num4 < num3:
+		array[num2] = num3 - num4 - 1
+		num2 = num2+1
+		for j in range(num4, num3):
+			array[num2] = data[j]
+			num2 = num2+1
+	osize = num2
+	array2 = bytearray(osize + 4)
+	array2[:4] = bytes2[:4]
+	array2[4:] = array[:osize]
+	return array2
+
+class Node:
+	mNext = 0
+	mPrev = 0
+	mPos = 0
+
+class Index:
+	mNodes = []
+	mStack = []
+	mStackPos = 0
+
+	def __init__(self):
+		for i in range(2304):
+			x = Node()
+			self.mNodes.append(x)
+		for i in range(2048, 2304):
+			self.mNodes[i].mNext = i
+			self.mNodes[i].mPrev = i
+		for i in range(2048):
+			self.mStack.append(i)
+		self.mStackPos = 2048
+
+	def getFirst(self, c):
+		return self.mNodes[2048+c].mNext
+
+	def getNode(self, i):
+		return self.mNodes[i]
+
+	def add(self, c, pos):
+		self.mStackPos = self.mStackPos - 1
+		num = self.mStack[self.mStackPos]
+		node = self.mNodes[num]
+		node2 = self.mNodes[2048+c]
+		node.mNext = node2.mNext
+		node.mPrev = 2048 + c
+		node.mPos = pos
+		self.mNodes[node2.mNext].mPrev = num
+		node2.mNext = num
+
+	def remove(self, c, pos):
+		mPrev = self.mNodes[2048+c].mPrev
+		node = self.mNodes[mPrev]
+		self.mStack[self.mStackPos] = self.mNodes[node.mPrev].mNext
+		self.mStackPos = self.mStackPos + 1
+		self.mNodes[node.mPrev].mNext = node.mNext
+		self.mNodes[node.mNext].mPrev = node.mPrev
+
+	def isEnd(self, idx):
+		return idx >= 2048
+
+def xor_crypt(data, key):
+	odata = bytearray(len(data))
+	odata[:] = data
+	i = 0
+	for x in data:
+		if (x != 0):
+			m = key[i % len(key)]
+			if (x != m):
+				odata[i] = x ^ m
+		i += 1
+	return odata

utage/utage.py

@@ -0,0 +1,46 @@
+import glob
+import io
+import traceback
+
+from functools import partial
+from multiprocessing.pool import Pool
+from os import path
+
+from . import crypt
+
+def crypt_dir(dir_in, key, encrypt=False, no_compress=False):
+	if not encrypt:
+		files = glob.glob(path.join(dir_in, "**/*.utage"), recursive=True)
+	else:
+		files = glob.glob(path.join(dir_in, "**/*.tsv"), recursive=True)
+
+	if len(files) == 0:
+		raise FileNotFoundError("No valid files found in directory: " + dir_in)
+
+	pcrypt = partial(crypt_file, key=key, encrypt=encrypt, no_compress=no_compress)
+	with Pool() as p:
+		p.map(pcrypt, files)
+
+def crypt_file(file_in, key, encrypt=False, no_compress=False):
+	with open(file_in, "rb") as inf:
+		in_data = inf.read()
+	try:
+		if encrypt:
+			if not file_in.endswith(".png") and not file_in.endswith(".tsv"):
+				raise ValueError("Invalid File Type for {}".format(file_in))
+			if not file_in.endswith(".png") and not no_compress:
+				enc_data = crypt.compress(in_data)
+			enc_data = crypt.xor_crypt(enc_data, key)
+			with io.open(file_in + ".utage", "wb") as output:
+				output.write(enc_data)
+		else:
+			if not file_in.endswith(".utage"):
+				raise ValueError("Invalid File Type for {}".format(file_in))
+			dec_data = crypt.xor_crypt(in_data, key)
+			if not file_in.endswith(".png.utage") and not no_compress:
+				dec_data = crypt.decompress(dec_data)
+			with io.open(file_in.replace(".utage", ""), "wb") as output:
+				output.write(dec_data)
+	except Exception as e:
+		traceback.print_exc()
+		print("Error processing file: " + file_in)