Rework GPT checking and prompts

mods/live-installer/data/opt/wdvn/installer/main

@@ -1,4 +1,5 @@
 #!/usr/bin/python3
+import sys
 import os
 import subprocess
 import getpass
@@ -104,6 +105,58 @@ def getPartitionType(partition):
 def getPartitionUUID(partition):
 	return callString(["blkid","-o","value","-s","UUID",partition])
 
+def getDiskTable(disk):
+	if getDiskType(disk) == "gpt":
+		results = callList(["fdisk","-x",disk,"-o","Device,Type-UUID"])
+	else:
+		results = callList(["fdisk","-x",disk,"-o","Device,Id"])
+	length = len(results)
+	index = length - 1
+	while index >= 0:
+		if results[index] == "": break
+		index -= 1
+	if index < 0: return {}
+	results = results[index + 2:]
+	length = len(results)
+	
+	index = 0
+	while index < length:
+		while results[index].replace("  "," ") != results[index]:
+			results[index] = results[index].replace("  "," ")
+		results[index] = results[index].split(" ")
+		index += 1
+	
+	rtn = {}
+	for part in results:
+		rtn[part[0]] = part[1]
+	
+	return rtn
+
+def getGptReport(disk):
+	rtn = {}
+	rtn["isGPT"] = False
+	rtn["espPart"] = None
+	rtn["espFormatted"] = False
+	rtn["biosPart"] = None
+	
+	if getDiskType(disk) != "gpt": return rtn
+	rtn["isGPT"] = True
+	partitions = callList(["lsblk",disk,"-no","PATH"])
+	del partitions[0]
+	diskTable = getDiskTable(disk)
+	
+	for partition in partitions:
+		if rtn["espPart"] == None:
+			if diskTable[partition] == "C12A7328-F81F-11D2-BA4B-00A0C93EC93B":
+				rtn["espPart"] = partition
+				rtn["espFormatted"] = (getPartitionType(partition) == "vfat")
+		
+		if rtn["biosPart"] == None:
+			if diskTable[partition] == "21686148-6449-6E6F-744E-656564454649":
+				rtn["biosPart"] = partition
+	
+	return rtn
+
 def setKbLayout(kb):
 	call(["loadkeys",kb])
 
@@ -198,6 +251,23 @@ def main():
 		except KeyboardInterrupt:
 			installGRUB = False
 	
+	if installGRUB:
+		gptInfo = getGptReport(grubDisk)
+		if gptInfo["isGPT"] == True:
+			if gptInfo["espPart"] == None:
+				if gptInfo["biosPart"] == None:
+					if choiceYn("WARNING: No ESP nor BIOS partition found on GPT disk " +grubDisk+ ", GRUB will not be installed. Continue?","n") == False:
+						sys.exit(1)
+					else:
+						installGRUB = False
+				else:
+					if choiceYn("WARNING: No ESP partition found. The disk might only boot on BIOS (legacy). Continue?","n") == False:
+						sys.exit(1)
+			else:
+				formatEsp = False
+				if gptInfo["espFormatted"] == False:
+					formatEsp = choiceYn("The disk's ESP partition (" +gptInfo["espPart"]+ ") does not seem to be formatted correctly. Format it?")
+	
 	print("")
 	print("The rest of this setup will continue without your input ...")
 	time.sleep(5)
@@ -296,46 +366,35 @@ def main():
 	
 	if installGRUB:
 		print("> Installing GRUB")
-		mbrInstall = True
-		diskType = getDiskType(grubDisk)
-		if diskType == "gpt":
-			print("Scanning GPT disk ...")
-			partitions = callList(["lsblk",grubDisk,"-no","PATH"])
-			
-			del partitions[0]
-			try:
-				del partitions[partitions.index(installPartition)]
-			except Exception:
-				pass
-			
-			if encrypt:
-				try:
-					del partitions[partitions.index(cryptoPartition)]
-				except Exception:
-					pass
-			
-			if len(partitions) > 0:
-				partType = getPartitionType(partitions[0])
-				if partType == "fat" or partType == "vfat":
-					print("Mounting ESP ...")
-					mkdirp(ipth("boot/efi"))
-					call(["mount",partitions[0],ipth("boot/efi")])
-					print("Installing GRUB ...")
-					mbrInstall = False
-					chroot(ipth(),["grub-install","--target=i386-efi","--uefi-secure-boot","--efi-directory=/boot/efi","--boot-directory=/boot",grubDisk])
-					chroot(ipth(),["grub-install","--target=x86_64-efi","--uefi-secure-boot","--efi-directory=/boot/efi","--boot-directory=/boot",grubDisk])
-					nerr(chroot,ipth(),["grub-install","--target=i386-pc","--boot-directory=/boot",grubDisk])
-					chroot(ipth(),["update-grub"])
-					print("Unmounting ESP ...")
-					call(["umount","-l",ipth("boot/efi")])
-					os.rmdir(ipth("boot/efi"))
 		
-		if mbrInstall:
+		if gptInfo["isGPT"]:
-			if diskType == "gpt":
+			if gptInfo["espPart"] != None:
-				print("WARNING: GPT disk unsuitable for EFI, installing BIOS payload only.")
+				if formatEsp:
-			print("Installing GRUB ...")
+					print("Formatting ESP ...")
+					call(["wipefs","-a",gptInfo["espPart"]])
+					call(["mkfs.vfat","-F","32",gptInfo["espPart"]])
+				
+				print("Mounting ESP ...")
+				mkdirp(ipth("boot/efi"))
+				call(["mount",gptInfo["espPart"],ipth("boot/efi")])
+				
+				print("Installing GRUB (EFI) ...")
+				chroot(ipth(),["grub-install","--target=i386-efi","--uefi-secure-boot","--efi-directory=/boot/efi","--boot-directory=/boot",grubDisk])
+				chroot(ipth(),["grub-install","--target=x86_64-efi","--uefi-secure-boot","--efi-directory=/boot/efi","--boot-directory=/boot",grubDisk])
+				
+				print("Unmounting ESP ...")
+				call(["umount","-l",ipth("boot/efi")])
+				os.rmdir(ipth("boot/efi"))
+			
+			if gptInfo["biosPart"] != None:
+				print("Installing GRUB (BIOS) ...")
+				chroot(ipth(),["grub-install","--target=i386-pc","--boot-directory=/boot",grubDisk])
+		else:
+			print("Installing GRUB (BIOS) ...")
 			chroot(ipth(),["grub-install","--target=i386-pc","--boot-directory=/boot",grubDisk])
-			chroot(ipth(),["update-grub"])
+		
+		print("Configuring GRUB ...")
+		chroot(ipth(),["update-grub"])
 		
 		print("> Unmounting ...")
 		call(["swapoff",ipth("swap")])