#!/usr/bin/python3
import os
import subprocess
import getpass
import time
import shutil
import random

def nerr(func,*args,**kwargs):
	try:
		return func(*args,**kwargs)
	except Exception:
		return

def choiceYn(question,defaultPick = None):
	qstring = "\n" + question + " y/n"
	if defaultPick != None: qstring = qstring + " [" +defaultPick+ "]"
	qstring = qstring + ": "
	while True:
		chinput = input(qstring)
		if chinput == "" and defaultPick != None:
			chinput = defaultPick
		chinput = chinput.lower()
		if chinput == "y": return True
		if chinput == "n": return False

def choice(question,options,defaultPick = None,direction = 1):
	qstring = "\n"
	for option in options:
		qstring = qstring + "- " + option + "\n"
	qstring = qstring + question
	if defaultPick != None: qstring = qstring + " [" +defaultPick+ "]"
	qstring = qstring + ": "
	
	while True:
		chinput = input(qstring)
		if chinput == "" and defaultPick != None:
			chinput = defaultPick
		if chinput in options: return chinput

def choiceInput(question,defaultPick = None):
	qstring = "\n" + question
	if defaultPick != None: qstring = qstring + " [" +defaultPick+ "]"
	qstring = qstring + ": "
	
	while True:
		chinput = input(qstring)
		if chinput == "":
			if defaultPick == None: continue
			return defaultPick
		return chinput

def choicePass(question):
	while True:
		print("")
		passw = getpass.getpass(question+ ": ")
		passwConfirm = getpass.getpass(question+ " (repeat): ")
		if passw == passwConfirm: return passw
		print("Passwords don't match, try again.")

def checkProc(proc):
	rtn = proc.wait()
	if rtn != 0: raise Exception("Process returned error: " +str(rtn))

def call(*args,**kwargs):
	proc = subprocess.Popen(*args,**kwargs)
	checkProc(proc)
	return

def callString(*args,**kwargs):
	proc = subprocess.Popen(*args,**kwargs,stdout=subprocess.PIPE)
	output = proc.stdout.read().decode("utf-8").strip("\n")
	checkProc(proc)
	return output

def callList(*args,**kwargs):
	proc = subprocess.Popen(*args,**kwargs,stdout=subprocess.PIPE)
	output = proc.stdout.read().decode("utf-8").strip("\n").split("\n")
	checkProc(proc)
	return output

def getDisks():
	return callList(["lsblk","-ndo","PATH"])

def getPartitions():
	disks = getDisks()
	partitions = callList(["lsblk","-no","PATH"])
	index = 0
	length = len(partitions)
	while index < length:
		if partitions[index] in disks:
			del partitions[index]
			length -= 1
			continue
		index += 1
	return partitions

def getDiskType(disk):
	return callString(["blkid","-o","value","-s","PTTYPE",disk])

def getPartitionType(partition):
	return callString(["blkid","-o","value","-s","TYPE",partition])

def getPartitionUUID(partition):
	return callString(["blkid","-o","value","-s","UUID",partition])

def setKbLayout(kb):
	call(["loadkeys",kb])

def resetKbLayout():
	call(["service","console-setup.sh","restart"])
	call(["udevadm","trigger","--subsystem-match=input","--action=change"])
	call(["service","keyboard-setup.sh","restart"])

def ipth(path = ""):
	return "/media/install/" +path

def mkdirp(path):
	if not os.path.isdir(path): os.makedirs(path)

def chroot(path,cmd,*args,**kwargs):
	dirs = ["dev","dev/pts","sys","proc"]
	try:
		for dir in dirs:
			call(["mount","-o","bind","/" +dir,path + "/" +dir])
		call(["chroot",path] + cmd,*args,**kwargs)
		for dir in reversed(dirs):
			call(["umount","-l",path + "/" +dir])
	except Exception as e:
		for dir in reversed(dirs):
			try:
				call(["umount","-l",path + "/" +dir])
			except:
				pass
		raise e

def randhex(length):
	chars = "0123456789ABCDEF"
	output = ""
	cLength = 0
	while cLength < length:
		output = output + chars[random.randrange(0,15)]
		cLength += 1
	return output

def main():
	while choiceYn("Partition disk?"):
		try:
			disk = choice("Disk",getDisks())
			call(["fdisk",disk])
		except KeyboardInterrupt:
			pass
	
	installPartition = choice("OS Partition",getPartitions())
	formatPartition = choiceYn("Format partition?","y")
	if formatPartition:
		fileSystems = []
		for path in os.environ["PATH"].split(":"):
			if not os.path.isdir(path): continue
			for root,dirs,files in os.walk(path):
				for file in files:
					ffile = os.path.join(root,file)
					lfile = ffile.replace(path + os.path.sep,"",1)
					if lfile.startswith("mkfs."): fileSystems.append(lfile.replace("mkfs.","",1))
		
		fileSystems.sort()
		defaultChoice = None
		if "ext4" in fileSystems: defaultChoice = "ext4"
		formatPartition = choice("File system",fileSystems,defaultChoice)
	else:
		formatPartition = None
	
	if formatPartition != None:
		encrypt = choiceYn("Encrypt data?")
		if encrypt:
			encryptionPasswords = []
			setKbLayout("us")
			try:
				encryptionPasswords.append(choicePass("Encryption password"))
			except KeyboardInterrupt:
				encrypt = False
			resetKbLayout()
		
		if encrypt:
			while choiceYn("Add another password?"):
				setKbLayout("us")
				try:
					encryptionPasswords.append(choicePass("Encryption password"))
				except KeyboardInterrupt:
					resetKbLayout()
					break
				resetKbLayout()
	
	installGRUB = choiceYn("Install GRUB?","y")
	if installGRUB:
		try:
			grubDisk = choice("GRUB disk",getDisks())
		except KeyboardInterrupt:
			installGRUB = False
	
	print("")
	print("The rest of this setup will continue without your input ...")
	time.sleep(5)
	print("")
	
	if formatPartition != None:
		print("> Wiping partition signature ...")
		call(["wipefs","-a",installPartition])
	
	if encrypt:
		print("> Encrypting partition ...")
		call(["apt","-y","install","cryptsetup"],stdout=subprocess.DEVNULL)
		cryptPartition = installPartition
		
		mainPass = encryptionPasswords.pop(0)
		proc = subprocess.Popen(["cryptsetup","luksFormat","--type","luks1",cryptPartition],stdin=subprocess.PIPE)
		proc.stdin.write((mainPass + "\n").encode("utf-8"))
		proc.stdin.flush()
		checkProc(proc)
		
		for passw in encryptionPasswords:
			print("Adding another password ...")			
			proc = subprocess.Popen(["cryptsetup","luksAddKey",cryptPartition],stdin=subprocess.PIPE)
			proc.stdin.write((mainPass + "\n" + passw + "\n").encode("utf-8"))
			proc.stdin.flush()
			checkProc(proc)
		
		print("Creating block device ...")
		proc = subprocess.Popen(["cryptsetup","luksOpen",cryptPartition,"system"],stdin=subprocess.PIPE)
		proc.stdin.write((mainPass + "\n").encode("utf-8"))
		proc.stdin.flush()
		checkProc(proc)
		installPartition = "/dev/mapper/system"
	
	if formatPartition != None:
		print("> Formatting partition ...")
		call(["mkfs." +formatPartition,installPartition])
	
	print("> Grabbing partition UUIDs ...")
	installPartitionUUID = getPartitionUUID(installPartition)
	if encrypt: cryptPartitionUUID = getPartitionUUID(cryptPartition)
	
	print("> Mounting main partition ...")
	mkdirp(ipth())
	call(["mount",installPartition,ipth()])
	
	print("> Adding files")
	print("Creating swap ...")
	call(["dd","if=/dev/zero","of=" +ipth("swap"),"bs=1M","count=512","status=progress"])
	call(["chmod","600",ipth("swap")])
	call(["mkswap",ipth("swap")])
	call(["swapon",ipth("swap")])
	
	print("Unpacking OS ...")
	call(["unsquashfs","-f","-d",ipth(),"/lib/live/mount/medium/live/filesystem.squashfs"])
	
	print("Setting hostname ...")
	hostname = randhex(8)
	fh = open(ipth("etc/hostname"),"w")
	fh.write(hostname + "\n")
	fh.close()
	
	if encrypt:
		print("Writing crypto grub config ...")
		fh = open(ipth("etc/default/grub.d/cryptodisk.cfg"),"w")
		fh.write("GRUB_ENABLE_CRYPTODISK=y\n")
		fh.write('GRUB_CMDLINE_LINUX="$GRUB_CMDLINE_LINUX cryptdevice=UUID=' +cryptPartitionUUID+ ' root=UUID=' +installPartitionUUID+ '"\n')
		fh.close()
		
		print("Writing crypttab ...")
		fh = open(ipth("etc/crypttab"),"w")
		fh.write('system UUID=' +cryptPartitionUUID+ ' none luks\n')
		fh.close()
	
	print("Writing fstab ...")
	fh = open(ipth("etc/fstab"),"w")
	fh.write('UUID=' +installPartitionUUID+ ' / ' +getPartitionType(installPartition)+ ' defaults 0 1\n')
	fh.write('/swap none swap sw 0 0\n')
	fh.close()
	
	print("Copying keyboard settings ...")
	shutil.copyfile("/etc/default/keyboard",ipth("etc/default/keyboard"))
	
	print("> Removing live-specific packages ...")
	chroot(ipth(),["apt","-y","remove","live-boot","live-boot-initramfs-tools"])
	chroot(ipth(),["apt","-y","autoremove"])
	os.remove(ipth("bin/login"))
	os.rename(ipth("bin/login.oobe"),ipth("bin/login"))
	
	if encrypt:
		print("> Installing encryption-specific packages ...")
		chroot(ipth(),["apt","-y","install","cryptsetup-initramfs"])
	else:
		print("> Updating initramfs ...")
		chroot(ipth(),["update-initramfs","-u"])
	
	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 diskType == "gpt":
				print("WARNING: GPT disk unsuitable for EFI, installing BIOS payload only.")
			print("Installing GRUB ...")
			chroot(ipth(),["grub-install","--target=i386-pc","--boot-directory=/boot",grubDisk])
			chroot(ipth(),["update-grub"])
		
		print("> Unmounting ...")
		call(["swapoff",ipth("swap")])
		call(["umount",ipth()])
		os.rmdir(ipth())
		if encrypt: call(["cryptsetup","luksClose","system"])
		
		print("Success. Press ENTER to quit setup.")
		input()
	
main()