Change to index file creation process

NoAuthTinGen.py

@@ -1,24 +1,48 @@
-from urllib3 import disable_warnings as disable_url_warnings
-from TinGen import UTinGen
 from argparse import ArgumentParser
-from pathlib import Path
 from binascii import unhexlify
-from TinGen.utils import create_encrypted_index
+from pathlib import Path
+from TinGen.utils import CompressionFlag
+from TinGen.utils import create_tinfoil_index
+from TinGen import UTinGen
+from urllib3 import disable_warnings as disable_url_warnings
 
 if __name__ == "__main__":
-	disable_url_warnings()
-	parser = ArgumentParser(description="Script that will allow you to easily generate an index file with Google Drive file links for use with Tinfoil without requiring authentication.")
-	parser.add_argument(nargs="*", metavar="FOLDER_ID_TO_SCAN", dest="folder_ids", help="Folder ID of Google Drive folders to scan. Can use more than 1 folder IDs at a time. FOLDERS MUST BE PUBLIC FOR SCRIPT TO WORK")
-	parser.add_argument("--index-file", metavar="INDEX_FILE_PATH", default="index.tfl", help="File Path for unencrypted index file to update.")
-	parser.add_argument("--encrypt", nargs="?", metavar="ENC_INDEX_FILE_PATH", const="enc_index.tfl", help="Use this flag is you want to encrypt the resulting index file.")
-	parser.add_argument("--public-key", metavar="PUBLIC_KEY_FILE_PATH", default="public.pem", help="File Path for Public Key to encrypt with.")
-	parser.add_argument("--vm-file", help="File Path for VM File")
-	parser.add_argument("--success", metavar="SUCCESS_MESSAGE", help="Success Message to add to index.")
-
-	args = parser.parse_args()
-
-	generator = UTinGen(index_path=args.index_file)
-	generator.index_folders(args.folder_ids, success=args.success)
-
-	if args.encrypt:
-		create_encrypted_index(generator.index_json, Path(args.encrypt), Path(args.public_key), None if not args.vm_file else Path(args.vm_file))
+    disable_url_warnings()
+    parser = ArgumentParser(description="Script that will allow you to easily generate an index file with Google Drive file links for use with Tinfoil without requiring authentication")
+    parser.add_argument(nargs="*", metavar="FOLDER_ID_TO_SCAN", dest="folder_ids", help="Folder ID of Google Drive folders to scan")
+    parser.add_argument("--index-path", default="index.tfl", help="File path for unencrypted index file to update")
+    parser.add_argument("--add-nsw-files-without-title-id", action="store_true", help="Adds files without title ID")
+    parser.add_argument("--add-non-nsw-files", action="store_true", help="Adds files without valid NSW ROM extension(NSP/NSZ/XCI/XCZ) to index")
+    parser.add_argument("--encrypt", action="store_true", help="Encrypt the resulting index file")
+    parser.add_argument("--public-key", default="public.pem", help="File path for public key to encrypt with")
+    parser.add_argument("--vm-file", help="File path for VM file")
+    parser.add_argument("--success", help="Success message to add to index")
+    compression_parser = parser.add_mutually_exclusive_group()
+    compression_parser.add_argument("--zstandard", action="store_true", help="Compresses index with Zstandard compression method")
+    compression_parser.add_argument("--zlib", action="store_true", help="Compresses index with Zlib compression method")
+    compression_parser.add_argument("--no-compress", action="store_true", help="Flag to not compress index")
+
+    args = parser.parse_args()
+
+    generator = UTinGen()
+    generator.index_generator(args.folder_ids, add_non_nsw_files=args.add_non_nsw_files, add_nsw_files_without_title_id=args.add_nsw_files_without_title_id, success=args.success)
+
+    compression_flag = CompressionFlag.ZSTD_COMPRESSION
+
+    if args.zstandard:
+        compression_flag = CompressionFlag.ZSTD_COMPRESSION
+    elif args.zlib:
+        compression_flag = CompressionFlag.ZLIB_COMPRESSION
+    elif args.no_compress:
+        compression_flag = CompressionFlag.NO_COMPRESSION
+
+    vm_file = None
+    public_key = None
+
+    if args.encrypt:
+        if args.vm_file:
+            vm_file = args.vm_file
+        if args.public_key:
+            public_key = args.public_key
+
+    create_tinfoil_index(generator.index, args.index_path, compression_flag, public_key, vm_file)

TinGen.py

@@ -1,8 +1,9 @@
 from argparse import ArgumentParser
-from TinGen import TinGen
-from pathlib import Path
 from binascii import unhexlify
-from TinGen.utils import create_encrypted_index
+from pathlib import Path
+from TinGen.utils import CompressionFlag
+from TinGen.utils import create_tinfoil_index
+from TinGen import TinGen
 
 if __name__ == "__main__":
     parser = ArgumentParser(description="Script that will allow you to generate an index file with Google Drive file links for use with Tinfoil")
@@ -13,14 +14,13 @@
     parser.add_argument("--headless", action="store_true", help="Allows to perform Google Token Authentication in headless environment")
 
     parser.add_argument("--index-file", metavar="INDEX_FILE_PATH", default="index.tfl", help="File path for index file")
-    parser.add_argument("--update-mode", action="store_true", help="Updates existing index file keeping old files, if it exists, instead of regenerating a new file")
     parser.add_argument("--share-files", action="store_true", help="Share files all files inside the index file")
     parser.add_argument("--no-recursion", dest="recursion", action="store_false", help="Scans for files only in top directory for each folder ID entered")
     parser.add_argument("--add-nsw-files-without-title-id", action="store_true", help="Adds files without title ID")
     parser.add_argument("--add-non-nsw-files", action="store_true", help="Adds files without valid NSW ROM extension(NSP/NSZ/XCI/XCZ) to index")
     parser.add_argument("--success", metavar="SUCCESS_MESSAGE", help="Adds a success message to index file to show if index is successfully read by tinfoil")
 
-    parser.add_argument("--encrypt", nargs="?", metavar="ENC_INDEX_FILE_PATH", const="enc_index.tfl", help="Encrypts the resulting index file")
+    parser.add_argument("--encrypt", action="store_true", help="Encrypts the resulting index file")
     parser.add_argument("--public-key", metavar="PUBLIC_KEY_FILE_PATH", default="public.pem", help="File Path for Public Key to encrypt with")
     parser.add_argument("--vm-file", help="File Path for VM File")
 
@@ -29,6 +29,11 @@
     parser.add_argument("--new-upload-id", action="store_true", help="Uploads the newly generated index file with a new file ID instead of replacing old one")
     parser.add_argument("--share-uploaded-index", action="store_true", help="Shares the index file that is uploaded to Google Drive")
 
+    compression_parser = parser.add_mutually_exclusive_group()
+    compression_parser.add_argument("--zstandard", action="store_true", help="Compresses index with Zstandard compression method")
+    compression_parser.add_argument("--zlib", action="store_true", help="Compresses index with Zlib compression method")
+    compression_parser.add_argument("--no-compress", action="store_true", help="Flag to not compress index")
+
     args = parser.parse_args()
     generator = TinGen(args.token, args.credentials, args.headless)
 
@@ -39,29 +44,37 @@
         print(f"Adding success message to index")
         generator.update_index_success_message(args.success)
 
-    if args.update_mode:
-        print(f"Adding files from {args.index_file} to new index")
-        generator.read_index(args.index_file)
+    compression_flag = CompressionFlag.ZSTD_COMPRESSION
+
+    if args.zstandard:
+        compression_flag = CompressionFlag.ZSTD_COMPRESSION
+    elif args.zlib:
+        compression_flag = CompressionFlag.ZLIB_COMPRESSION
+    elif args.no_compress:
+        compression_flag = CompressionFlag.NO_COMPRESSION
 
-    print(f"Writing generated index to {args.index_file}")
-    generator.write_index_to_file(args.index_file)
+    vm_file = None
+    public_key = None
+
+    if args.encrypt:
+        if args.vm_file:
+            vm_file = args.vm_file
+        if args.public_key:
+            public_key = args.public_key
+
+    print(f"Creating generated index to {args.index_file}")
+    create_tinfoil_index(generator.index, args.index_file, compression_flag, public_key, vm_file)
 
     if args.share_files:
         print(f"Sharing files in index")
         generator.share_index_files()
 
-    if args.encrypt:
-        print(f"Encrypting index to {args.encrypt}")
-        create_encrypted_index(generator.index, Path(args.encrypt), Path(args.public_key), None if not args.vm_file else Path(args.vm_file))
-
     if args.upload_folder_id:
-        file_to_upload = args.index_file if not args.encrypt else args.encrypt
-        print(f"Uploading {file_to_upload} to {args.upload_folder_id}")
-        generator.gdrive_service.upload_file(file_to_upload, args.upload_folder_id, args.share_uploaded_index, args.new_upload_id)
+        print(f"Uploading {args.index_file} to {args.upload_folder_id}")
+        generator.gdrive_service.upload_file(args.index_file, args.upload_folder_id, args.share_uploaded_index, args.new_upload_id)
 
     if args.upload_to_my_drive:
-        file_to_upload = args.index_file if not args.encrypt else args.encrypt
-        print(f"Uploading {file_to_upload} to \"My Drive\"")
-        generator.gdrive_service.upload_file(file_to_upload, None, args.share_uploaded_index, args.new_upload_id)
+        print(f"Uploading {args.index_file} to \"My Drive\"")
+        generator.gdrive_service.upload_file(args.index_file, None, args.share_uploaded_index, args.new_upload_id)
 
     print(f"Index Generation Complete")

TinGen/__init__.py

@@ -321,21 +321,15 @@ def get_folder_key(self, folder_id):
 			return ""
 
 class UTinGen:
-	def __init__(self, index_path="index.tfl"):
-		self.index_path = index_path
-		self.index_json = {"files": []}
+	def __init__(self):
+		self.index = {"files": []}
 		self.gdrive_service = UGdrive()
 
-	def write_index_to_file(self):
-		Path(self.index_path).parent.resolve().mkdir(parents=True, exist_ok=True)
-		with open(self.index_path, "w") as output_file:
-			json_writer(self.index_json, output_file, indent=2)
-
-	def index_folders(self, folder_ids, success=None, allow_files_without_tid=False):
+	def index_generator(self, folder_ids, add_non_nsw_files: bool, add_nsw_files_without_title_id: bool, success: str=None):
 		for folder_id in folder_ids:
 			for (file_id, file_details) in self.gdrive_service.get_files_in_folder_id(folder_id).items():
-				if allow_files_without_tid or regex_search(r"\%5B[0-9A-Fa-f]{16}\%5D", url_encode(file_details["name"], safe="")):
-					self.index_json["files"].append({"url": "gdrive:{file_id}#{file_name}".format(file_id=file_id, file_name=url_encode(file_details["name"], safe="")), "size": int(file_details["size"])})
+				if add_non_nsw_files or file_details["name"][-4:] in (".nsp", ".nsz", ".xci", ".xcz"):
+					if add_nsw_files_without_title_id or regex_search(r"\%5B[0-9A-Fa-f]{16}\%5D", url_encode(file_details["name"], safe="")):
+						self.index["files"].append({"url": "gdrive:{file_id}#{file_name}".format(file_id=file_id, file_name=url_encode(file_details["name"], safe="")), "size": int(file_details["size"])})
 		if success is not None:
-			self.index_json.update({"success": success})
-		self.write_index_to_file()
+			self.index.update({"success": success})

TinGen/utils.py

@@ -1,54 +1,77 @@
-from Crypto.Cipher.PKCS1_OAEP import new as new_pkcs1_oaep_ctx
-from Crypto.Cipher.AES import new as new_aes_ctx
+from binascii import hexlify
+from binascii import unhexlify
 from Crypto.Cipher.AES import MODE_ECB
-from Crypto.PublicKey.RSA import import_key as import_rsa_key
+from Crypto.Cipher.AES import new as new_aes_ctx
+from Crypto.Cipher.PKCS1_OAEP import new as new_pkcs1_oaep_ctx
 from Crypto.Hash import SHA256
+from Crypto.PublicKey.RSA import import_key as import_rsa_key
+from enum import IntEnum
+from json import dumps as json_serialize
 from pathlib import Path
 from random import randint
-from json import dumps as json_serialize
+from zlib import compress as zlib_compress
 from zstandard import ZstdCompressor
-from binascii import hexlify
-from binascii import unhexlify
 
-def rand_aes_key_generator() -> bytes:
-    return randint(0, 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF).to_bytes(0x10, "big")
 
-def create_encrypted_index(index_to_write: dict, out_path: Path, rsa_pub_key_path: Path, vm_path: Path):
-    if rsa_pub_key_path.is_file():
-        index_buffer = bytes(json_serialize(index_to_write).encode())
-        rsa_pub_key = import_rsa_key(open(rsa_pub_key_path).read())
-        rand_aes_key = rand_aes_key_generator()
+class CompressionFlag(IntEnum):
+    ZLIB_COMPRESSION = 0xFE
+    ZSTD_COMPRESSION = 0xFD
+    NO_COMPRESSION = 0x00
 
-        aes_ctx = new_aes_ctx(rand_aes_key, MODE_ECB)
-        pkcs1_oaep_ctx = new_pkcs1_oaep_ctx(rsa_pub_key, hashAlgo=SHA256, label=b"")
 
-        to_compress_buffer = b""
+def create_tinfoil_index(index_to_write: dict, out_path: Path, compression_flag: int, rsa_pub_key_path: Path=None, vm_path: Path=None):
+    to_compress_buffer = b""
 
-        if vm_path is not None and vm_path.is_file():
-            to_compress_buffer += b"\x13\x37\xB0\x0B"
-            vm_buffer = b""
+    if vm_path is not None and vm_path.is_file():
+        to_compress_buffer += b"\x13\x37\xB0\x0B"
+        vm_buffer = b""
 
-            with open(vm_path, "rb") as vm_stream:
-                vm_buffer += vm_stream.read()
+        with open(vm_path, "rb") as vm_stream:
+            vm_buffer += vm_stream.read()
 
-            to_compress_buffer += len(vm_buffer).to_bytes(4, "little")
-            to_compress_buffer += vm_buffer
+        to_compress_buffer += len(vm_buffer).to_bytes(4, "little")
+        to_compress_buffer += vm_buffer
 
-        to_compress_buffer += index_buffer
+    to_compress_buffer += bytes(json_serialize(index_to_write).encode())
 
-        index_compressed_buffer = ZstdCompressor(level=22).compress(to_compress_buffer)
+    to_write_buffer = b""
+    session_key = b""
 
-        session_key = pkcs1_oaep_ctx.encrypt(rand_aes_key)
-        encrypted_index = aes_ctx.encrypt(index_compressed_buffer + (b"\x00" * (0x10 - (len(index_compressed_buffer) % 0x10))))
+    if compression_flag == CompressionFlag.ZSTD_COMPRESSION:
+        to_write_buffer += ZstdCompressor(level=22).compress(to_compress_buffer)
 
-        Path(out_path.parent).mkdir(parents=True, exist_ok=True)
+    elif compression_flag == CompressionFlag.ZLIB_COMPRESSION:
+        to_write_buffer += zlib_compress(to_compress_buffer, 9)
 
-        with open(out_path, "wb") as out_stream:
-            out_stream.write(b"TINFOIL")
-            out_stream.write(b"\xFD") # Compression indicator flag - 0xF0 | 0x0D for zstandard, 0xF0 | 0x0E for zlib, 0xF0 | 0x00 for no compression
-            out_stream.write(session_key)
-            out_stream.write(len(index_compressed_buffer).to_bytes(8, "little"))
-            out_stream.write(encrypted_index)
+    elif compression_flag == CompressionFlag.NO_COMPRESSION:
+        to_write_buffer += to_compress_buffer
 
     else:
-        print(f"{rsa_pub_key_path} does not exist. Cannot encrypt without key.")
+        raise NotImplementedError("Compression method supplied is not implemented yet.")
+
+    to_write_buffer += (b"\x00" * (0x10 - (len(to_write_buffer) % 0x10)))
+
+    if rsa_pub_key_path.is_file():
+        def rand_aes_key_generator() -> bytes:
+            return randint(0, 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF).to_bytes(0x10, "big")
+
+        rsa_pub_key = import_rsa_key(open(rsa_pub_key_path).read())
+        rand_aes_key = rand_aes_key_generator()
+
+        pkcs1_oaep_ctx = new_pkcs1_oaep_ctx(rsa_pub_key, hashAlgo=SHA256, label=b"")
+        aes_ctx = new_aes_ctx(rand_aes_key, MODE_ECB)
+
+        session_key += pkcs1_oaep_ctx.encrypt(rand_aes_key)
+        to_write_buffer = aes_ctx.encrypt(to_write_buffer)
+
+    else:
+        session_key += b"\x00" * 0xFF
+
+    Path(out_path.parent).mkdir(parents=True, exist_ok=True)
+
+    with open(out_path, "wb") as out_stream:
+        out_stream.write(b"TINFOIL")
+        out_stream.write(bytes(compression_flag))
+        out_stream.write(session_key)
+        out_stream.write(len(to_write_buffer).to_bytes(8, "little"))
+        out_stream.write(to_write_buffer)