Added support for zip file extraction

README.md

@@ -10,12 +10,12 @@ This program will work with the following extensions:
 .chd > .m3u creation
 ```
 
-**You need the file "CueMaker.py and the links.cfg" everything else will be fetched from my "Content" repo.**
+You can also tell it to extract zip files, create the necessary .cue and .m3u files, and recompress them.
 
-Using this program is really simple. I'm using the argparse library, so you can pass the -h flag for usage instructions. Alternatively, you can see the output of the command here:
+Using this program is really simple. You can use the -h flag for usage instructions. Alternatively, you can see the output of the command here:
 
 ```
-usage: cuemaker [-h] [-r] [-g] [-m] {playstation,playstation2,saturn,3do} directory
+usage: cuemaker [-h] [-r] [-g] [-m] [-z] {playstation,playstation2,saturn,3do} directory
 
 Original .cue file fetcher for game roms and .m3u creator.
 
@@ -29,15 +29,26 @@ optional arguments:
   -r, --recursive       search sub-folders
   -g, --generic         create generic .cue files if originals can't be found
   -m, --m3u             create .m3u files for multiple disc games
+  -z, --zips            extracts zip files containing roms, processes them and re-zips them in a
+                        single one (-r needs to be activated)
 ```
+
+This is an example command:
+
+```
+$ cuemaker -zmr playstation .
+```
+
+This command will extract all zip files that finds in the current directory, search sub-folders for roms, create .m3u files if the games are multi-disc, and recompress anything that was previously compressed, treating them all as playstation games (meaning it will try to match them agains the playstation games database).
+
 ## Formatting
 
 You will need to follow certaing formatting on the roms name's in order to ensure the correct functionality of the program. The list of roms is ordered alphabetically when the program is executed, for this reason you need to make sure that Track and Disc files are in the correct order. To make it easy, make sure your roms comply with these points:
 
 * If the file is a Track file, it must be enclosed in parenthesis and must have a space beforehand: " (Track 2)"
 * If the game has multiple discs, you need to state it the same way as tracks: " (Disc 3)"
 * If the game has multiple discs and multiple track files, the Disc indicator **must come before the Track indicator.**
-* Preferably, put the Disc and Track indicators after the game's name, you may specify region anywhere else.
+* The Disc and Track indicators must be after the game's name, you may specify region anywhere else.
 
 This is a correctly formatted list of roms:
 

cuemaker.py

@@ -5,7 +5,7 @@
 #
 #This program is free software: you can redistribute it and/or modify
 #it under the terms of the GNU General Public License as published by
-#the Free Software Foundation, either version 3 of the License, or
+#the Free Softwareation, either version 3 of the License, or
 #(at your option) any later version.
 #
 #This program is distributed in the hope that it will be useful,
@@ -24,19 +24,22 @@
 ##################################
 
 import os, glob, hashlib, argparse, configparser, platform, ssl
+from zipfile import ZipFile, ZIP_DEFLATED
 from urllib.request import urlopen
 from urllib.parse import quote
+from shutil import rmtree
 
 
 # create SSL certificates
 if (not os.environ.get('PYTHONHTTPSVERIFY', '') and getattr(ssl, '_create_unverified_context', None)):
 	ssl._create_default_https_context = ssl._create_unverified_context
 
-#Stats to display at the end
+# stats to display at the end
 m3uWriteCounter = 0
 cueCreatedCounter = 0
+zipCreatedCounter = 0
 
-#Variables used for multi-track roms
+# variables used for multi-track roms
 currentGameCue = ""
 currentGameCuePath = ""
 currentGameTrackNumber = 0
@@ -59,7 +62,7 @@ def createGenericCue(cuePath, fileName):
 	cueCreatedCounter += 1
 
 def fetchCue(entryName):
-	#Creates links given an entry name
+	# creates links given an entry name
 	if args.system == "playstation":
 		link = "{}{}.cue".format(configParser.get('Links', 'psxCueBase'), entryName[:-4]).replace(" ", "%20")
 		cueText = urlopen(link).read().decode("UTF-8")
@@ -82,19 +85,32 @@ def getTrackNumber(entryName):
 	return str(trackNumber).zfill(2)
 
 def replaceCueFileName(cueText, fileName):
-	#Replaces the "FILE" entry in the cue with the provided file name
+	# replaces the "FILE" entry in the cue with the provided file name
 	nameIndex = cueText.find("\"")
 	lastIndex = cueText.find("\"", nameIndex + 1)
 	cueText = cueText.replace(cueText[nameIndex + 1:lastIndex], fileName, 1)
 	return cueText
 
+def zipFiles(folderPath):
+	global zipCreatedCounter
+	print("Compressing files in \"\u001b[1;33m" + folderPath + "\u001b[0m\"...")
+	with ZipFile(folderPath + ".zip", 'a') as currentZip:
+		files = []
+		files.extend(glob.glob(folderPath + "/*"))
+		for file in files:
+			if not os.path.basename(file) in currentZip.namelist():
+				currentZip.write(file, arcname=os.path.basename(file), compress_type=ZIP_DEFLATED)
+		zipCreatedCounter += 1
+	print("Deleting \"\u001b[1;33m" + folderPath + "\u001b[0m\"...\n--------")
+	rmtree(folderPath)
+
 def generateCue(file):
 	global currentGameCue
 	global currentGameCuePath
 	global cueCreatedCounter
 	global currentGameTrackNumber
 
-	#Ignore .chd files
+	# ignore .chd files
 	if file.rfind(".chd") != -1:
 		cueFiles.append(file)
 		return None
@@ -108,7 +124,7 @@ def generateCue(file):
 
 	cuePath = fileDirectory + fileName[0:len(fileName) - 4] + ".cue"
 	print("Found file: \"\u001b[1;33m" + fileName + "\u001b[0m\"")
-	#Only attempt to create cues if there's none
+	# only attempt to create cues if there's none
 	if not os.path.isfile(cuePath):
 		print("\u001b[1;32mGenerating SHA-1 hash...\u001b[0m\n")
 		fileHash = getSha1(file)
@@ -130,54 +146,53 @@ def generateCue(file):
 					print("\u001b[1;31mCouldn't find database entry.\u001b[0m\n--------")
 					return None
 
-		#Create an entry name based by a fixed offset from the hash
+		# create an entry name based by a fixed offset from the hash
 		print("\u001b[1;33mHash \u001b[1;32m" + fileHash + "\u001b[1;33m matches database!")
 		lineEnd = hashFile.find("\n", foundEntry)
 		entryName = hashFile[int(foundEntry + len(fileHash) + 2):int(lineEnd)]
 
-		#If the file has multiple tracks, use special conditions
+		# if the file has multiple tracks, use special conditions
 		if entryName.rfind("Track ") != -1:
 
-			#Modify entry name to get the appropiate link
+			# modify entry name to get the appropiate link
 			trackNumber = getTrackNumber(entryName)
 			entryName = entryName.replace(entryName[entryName.lower().index(" (track"):], ".cue")
 
 			if int(trackNumber) == 1:
-				#Fetch the cue and save its entirety into the global variable "currentGameCue" for later use
+				# fetch the cue and save its entirety into the global variable "currentGameCue" for later use
 				try:
 					cueText = fetchCue(entryName)
 				except Exception:
 					print("\u001b[1;31mCouldn't find original cue on GitHub.\u001b[0m\n--------")
 					return None
 
-				#Generate text to write to cue
+				# generate text to write to cue
 				currentGameCue = cueText
 				trackIndex = currentGameCue.find("TRACK ")
 				nextFileIndex = cueText.find("FILE \"", trackIndex)
 				trackCueText = cueText[:nextFileIndex]
 				trackCueText = replaceCueFileName(trackCueText, fileName)
 
-				#Generate path for cue
+				# generate path for cue
 				cuePathTrackIndex = cuePath.lower().rfind(" (track")
 				cuePathParenthesisIndex = cuePath.find(")", cuePathTrackIndex)
 				cuePath = cuePath.replace(cuePath[cuePathTrackIndex:cuePathParenthesisIndex+1], "")
 				cue = open(cuePath, "a+")
-				cueFiles.append(cuePath)
 
-				#Set data for future track files
+				# set data for future track files
 				currentGameCuePath = cuePath
 				currentGameTrackNumber = int(trackNumber) + 1
 				cueCreatedCounter += 1
 				cueFiles.append(cuePath)
 			else:
-				#If track number is other than 1, use the cue from the previous file with "track 1"
+				# if track number is other than 1, use the cue from the previous file with "track 1"
 
-				#If there's no set "current cue" this means there was no "track 1" rom
+				# if there's no set "current cue" this means there was no "track 1" rom
 				if currentGameCue == "":
 					print("\u001b[1;31mERROR: ROM doesn't have a TRACK 1\u001b[0m\n--------")
 					return None
 
-				#Generate text to write to cue
+				# generate text to write to cue
 				trackIndex = currentGameCue.find("TRACK " + str(currentGameTrackNumber).zfill(2))
 				fileIndex = currentGameCue.rfind("FILE \"", 0, trackIndex)
 				nextFileIndex = currentGameCue.find("FILE \"", trackIndex)
@@ -193,9 +208,9 @@ def generateCue(file):
 				print("\u001b[1;31mTrack file already present.\u001b[0m")
 			cue.close()
 
-		#For single file roms
+		# for single file roms
 		else:
-			#Try to fetch the cue, if it fails default to generic
+			# try to fetch the cue, if it fails default to generic
 			try:
 				cueText = fetchCue(entryName)
 			except Exception:
@@ -233,11 +248,9 @@ def createM3u(cueFiles):
 
 		if fileName.lower().find(" (disc") != -1 or fileName.lower().find("_(disc") != -1:
 			print("Found file: \"\u001b[1;33m" + fileName + "\u001b[0m\"")
-			print(fileName)
 			discWordIndex = fileName.lower().index(" (disc")
 			discEndIndex = fileName.find(")", discWordIndex) + 1
 			m3uFilePath = fileDirectory + fileName.replace(fileName[discWordIndex:discEndIndex], "").replace(fileName[-4:], ".m3u")
-			print(m3uFilePath)
 			m3u = open(m3uFilePath, "a+")
 			m3u.seek(0)
 			if m3u.read().find(fileName) == -1:
@@ -254,7 +267,7 @@ def createM3u(cueFiles):
 
 	configParser = configparser.RawConfigParser()
 
-	#Windows uses local folder for the links.cfg
+	# windows uses local folder for the links.cfg
 	if platform.system() == "Windows":
 		configFilePath = "links.cfg"
 	else:
@@ -268,12 +281,14 @@ def createM3u(cueFiles):
 	parser.add_argument("-r", "--recursive", action="store_true", help="search sub-folders")
 	parser.add_argument("-g", "--generic", action="store_true", help="create generic .cue files if originals can't be found")
 	parser.add_argument("-m", "--m3u", action="store_true", help="create .m3u files for multiple disc games")
+	parser.add_argument("-z", "--zips", action="store_true", help="extracts zip files containing roms, processes them and re-zips them in a single one (-r needs to be activated)")
 	args = parser.parse_args()
 
 	types = ("*.bin", "*.img", "*.chd")
 
 	matchingFiles = []
 	cueFiles = []
+	zipsFolders = []
 
 	if args.system == "playstation":
 		try:
@@ -330,6 +345,34 @@ def createM3u(cueFiles):
 		print("\u001b[0;31mInvalid directory\u001b[0m")
 		exit()
 
+	while(True):
+		if args.zips:
+			if not args.recursive:
+				print("\u001b[0;31mError: Recursive needs to be on to unzip files. Ommiting.\u001b[0m")
+				break
+
+			matchingZips = []
+			matchingZips.extend(glob.glob("**/*.zip", recursive=True))
+			matchingZips = sorted(matchingZips)
+
+			print("\n\u001b[1;32mExtracting .zip...\u001b[0m\n")
+
+			for zip in matchingZips:
+				with ZipFile(zip, 'r') as currentZip:
+					try:
+						discWordIndex = zip.lower().index(" (disc")
+						discEndIndex = zip.find(")", discWordIndex) + 1
+						tmpZipFolder = zip.replace(zip[discWordIndex:discEndIndex], "")
+					except ValueError:
+						tmpZipFolder = zip
+
+					print("Extracting: \"\u001b[1;33m" + zip + "...\u001b[0m\"\n--------")
+					tmpZipFolder = tmpZipFolder[:-4]
+					currentZip.extractall(tmpZipFolder)
+					if not tmpZipFolder in zipsFolders:
+						zipsFolders.append(tmpZipFolder)
+		break
+
 	if args.recursive:
 		for files in types:
 			matchingFiles.extend(glob.glob("**/" + files, recursive=True))
@@ -339,12 +382,17 @@ def createM3u(cueFiles):
 			matchingFiles.extend(glob.glob(files))
 			matchingFiles = sorted(matchingFiles)
 
-	print("\u001b[1;32mCreating .cue...\u001b[0m\n")
+	print("\n\u001b[1;32mCreating .cue...\u001b[0m\n")
 
 	for file in matchingFiles:
 		generateCue(file)
 
 	if args.m3u:
 		createM3u(cueFiles)
 
-	print("\n\n\033[32mFinished!\nCreated \u001b[35m" + str(cueCreatedCounter) + "\u001b[32m .cue and wrote \u001b[35m" + str(m3uWriteCounter) + "\u001b[32m lines to .m3u!\u001b[0m")
+	if args.zips:
+		print("\n\u001b[1;32mCompressing and deleting temporary folders...\u001b[0m\n")
+		for folder in zipsFolders:
+			zipFiles(folder)
+
+	print("\n\n\033[32mFinished!\nCreated \u001b[35m" + str(cueCreatedCounter) + "\u001b[32m .cue, \u001b[35m" + str(zipCreatedCounter) + "\u001b[32m .zip and wrote \u001b[35m" + str(m3uWriteCounter) + "\u001b[32m lines to .m3u!\u001b[0m")