#!/usr/bin/env python3
import sys
import os
import shutil
diskSize = "1536M"
swapSize = "256M"
outputImage = os.path.abspath(sys.argv[1])
mountPoint = os.path.abspath(sys.argv[2].rstrip("/"))

mountPoints = [ # List of mount points to dismount when quitting
	os.path.join(mountPoint,"proc"),
	os.path.join(mountPoint,"sys"),
	os.path.join(mountPoint,"dev","pts"),
	os.path.join(mountPoint,"dev"),
	mountPoint
]

preReqs = [ # List of required commands to check for
	"fallocate",
	"losetup",
	"fdisk",
	"losetup",
	"mkswap",
	"mkfs.ext4",
	"mount",
	"blkid",
	"debootstrap",
	"chmod",
	"chroot",
	"umount"
]

loopName = False

class processError(Exception): pass

import subprocess
def procWait(proc):
	rtn = proc.wait()
	if rtn != 0: raise processError

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

def callo(*args,**kwargs):
	data = b""
	proc = subprocess.Popen(*args,**kwargs,stdout=subprocess.PIPE)
	while True:
		b = proc.stdout.read()
		if b == b"": break
		data += b
	procWait(proc)
	return data.decode("utf-8")

def cleanUp():
	global diskSize
	global swapSize
	global outputImage
	global mountPoint
	global loopName
	
	for mp in mountPoints:
		print("Dismounting " +mp+ " ...")
		try:
			call(["umount",mp])
		except: pass
	
	try:
		os.rmdir(mountPoint)
	except Exception as e:
		print(e)

	print("Detaching loop device ...")
	try:
		call(["losetup","-d",loopName])
	except:
		pass

def main():
	global diskSize
	global swapSize
	global outputImage
	global mountPoint
	global loopName
	
	print("Checking prerequisites ...")
	call(["sh","-c","command"])
	for req in preReqs:
		try:
			call(["sh","-c","command -v " + req],stdout=subprocess.DEVNULL)
		except:
			print("Missing command: " +req)
			raise

	print("Creating disk image ...")
	if os.path.isfile(outputImage): os.remove(outputImage)
	call(["fallocate","-l",diskSize,outputImage])

	print("Creating loop device ...")
	call(["losetup","-fP",outputImage])

	print("Setting up loop device ...")
	table = callo(["losetup","--list"]).replace("\r","").split("\n")
	length = len(table)
	index = 0
	while index < length:
		while "  " in table[index]:
			table[index] = table[index].replace("  "," ")
		if table[index] == "":
			del table[index]
			length -= 1
			continue
		table[index] = table[index].split(" ")
		index += 1
	nameIndex = table[0].index("NAME")
	backfileIndex = table[0].index("BACK-FILE")
	
	index = 1
	while index < length:
		if table[index][backfileIndex] == outputImage:
			loopName = table[index][nameIndex]
			break
		index += 1

	if loopName == False:
		print("Loop device not found.")
		raise processError

	print("Partitioning ...")
	proc = subprocess.Popen(["fdisk",loopName],stdin=subprocess.PIPE)
	proc.communicate(("""\
n
e


+""" +swapSize+ """
t
swap
n
p



t

linux
w

\
""").encode("utf-8"))
	procWait(proc)

	print("Making swap ...")
	call(["mkswap",loopName + "p1"])

	print("Formatting / (ext4) ...")
	call(["mkfs.ext4",loopName+ "p2"])

	print("Mounting / ...")
	if not os.path.isdir(mountPoint): os.makedirs(mountPoint)
	call(["mount",loopName + "p2",mountPoint])

	print("Installing Devuan Chimaera ...")
	tries = 0
	while tries < 10:
		try:
			call(["debootstrap","--arch=i386","chimaera",mountPoint,"http://deb.devuan.org/merged"])
			break
		except:
			tries += 1
			continue

	if tries >= 10:
		raise processError
	
	if os.path.isdir("skel"):
		print("Adding skel ...")
		if not os.path.isdir(os.path.join(mountPoint,"etc","skel")):
			os.makedirs(os.path.join(mountPoint,"etc","skel"))
		for root,dirs,files in os.walk("skel"):
			for file in dirs:
				ffile = os.path.join(root,file)
				lfile = ffile.replace("skel" + os.path.sep,"",1)
				ofile = os.path.join(mountPoint,"etc","skel",lfile)
				if not os.path.isdir(ofile): os.makedirs(ofile)
			
			for file in files:
				ffile = os.path.join(root,file)
				lfile = ffile.replace("skel" + os.path.sep,"",1)
				ofile = os.path.join(mountPoint,"etc","skel",lfile)
				if os.path.isfile(ofile): os.remove(ofile)
				shutil.copyfile(ffile,ofile) 
	
	print("Setting up chroot ...")
	call(["mount","-o","bind","/dev",os.path.join(mountPoint,"dev")])
	call(["mount","--bind","/dev/pts",os.path.join(mountPoint,"dev","pts")])
	call(["mount","-o","bind","/sys",os.path.join(mountPoint,"sys")])
	call(["mount","-t","proc","/proc",os.path.join(mountPoint,"proc")])
	with open("setup","r",encoding="utf-8") as fhIn:
		with open(os.path.join(mountPoint,"setup"),"w",encoding="utf-8") as fhOut:
			for line in fhIn:
				fhOut.write(line.replace("$LOOPNAME",loopName))
	call(["chmod","+x",os.path.join(mountPoint,"setup")])
	
	print("Adding devices to fstab ...")
	idSwap = callo(["blkid","-o","value","-s","UUID",loopName+ "p1"]).strip(" \t\r\n")
	idRoot = callo(["blkid","-o","value","-s","UUID",loopName+ "p2"]).strip(" \t\r\n")
	with open(os.path.join(mountPoint,"etc","fstab"),"a+",encoding="utf-8") as fh:
		fh.write("\nUUID=" +idSwap+ " none swap sw 0 0")
		fh.write("\nUUID=" +idRoot+ " / ext4 errors=remount-ro 0 1")

	print("Running setup ...")
	call(["chroot",mountPoint,"/setup"])
	
	print("Injecting OOBE ...")
	os.rename(os.path.join(mountPoint,"bin","login"),os.path.join(mountPoint,"bin","login.bak"))
	shutil.copyfile("oobe",os.path.join(mountPoint,"bin","login"))
	call(["chmod","+x",os.path.join(mountPoint,"bin","login")])
	
	cleanUp()

try:
	main()
except:
	print("\n --- AN ERROR OCCURED, CLEANING UP ... --- ")
	cleanUp()
	print("")
	raise