''' distro-installer core. '''
''' Author: Ikey Doherty <contactjfreak@gmail.com> '''

import os
import subprocess
from subprocess import Popen
import time
import shutil
import gettext
import stat

gettext.install("distro-installer", "/usr/share/locale")

class mINIFile():
        def load_from_string(self, line):
                if(line.find(":")):
                        l = line.split(":")
                        if(len(l) >= 2):
                                tmp = ":".join(l[1:]).rstrip("\r\n")
                                setattr(self, l[0], tmp)
                elif(line.find("=")):
                        l = line.split("=")
                        if(len(l) >= 2):
                                tmp = "=".join(l[1:]).rstrip("\r\n")
                                setattr(self, l[0], tmp)

        def load_from_list(self, list):
                for line in list:
                        self.load_from_string(line)
        def load_from_file(self, filename):
                try:
                        fi = open(filename, "r")
                        self.load_from_list(fi.readlines())
                        fi.close()
                except:
                        pass

        def has_key(self, key):
			return hasattr(self, key)

class SystemUser:
	''' Represents the main user '''

	def __init__(self, username=None, realname=None, password=None, sudo=False):
		''' create new SystemUser '''
		self.username = username
		self.realname = realname
		self.password = password

class InstallerEngine:
	''' This is central to the distro installer '''

	# some noice definitions..
	PROGRESS_START		= "start"
	PROGRESS_UPDATE 	= "update"
	PROGRESS_COMPLETE	= "complete"
	PROGRESS_ERROR		= "error"

	def __init__(self):
		# set up stuffhs
		self.conf_file = '/etc/distro-installer/distro.conf'
		configuration = mINIFile()
		configuration.load_from_file(self.conf_file)
		if(not configuration.has_key("Distro")):
			self.distro_name = "GNU/Linux"
		else:
			self.distro_name = configuration.Distro
		if(not configuration.has_key("Version")):
			self.distro_version = "1"
		else:
			self.distro_version = configuration.Version
		self.user = None

	def set_main_user(self, user):
		''' Set the main user to be used by the installer '''
		if(user is not None):
			self.user = user

	def get_main_user(self):
		''' Return the main user '''
		return self.user

	def format_device(self, device, filesystem):
		''' Format the given device to the specified filesystem '''
		cmd = "mkfs -t %s %s" % (filesystem, device)
		p = Popen(cmd, shell=True)
		p.wait() # this blocks
		return p.returncode

	def set_install_media(self, media=None, type=None):
		''' Sets the location of our install source '''
		self.media = media
		self.media_type = type

	def add_to_blacklist(self, blacklistee):
		''' This will add a directory or file to the blacklist, so that '''
		''' it is not copied onto the new filesystem '''
		try:
			self.blacklist.index(blacklistee)
			self.blacklist.append(blacklistee)
		except:
			# We haven't got this item yet
			pass

	def set_progress_hook(self, progresshook):
		''' Set a callback to be called on progress updates '''
		''' i.e. def my_callback(progress_type, message, current_progress, total) '''
		''' Where progress_type is any off PROGRESS_START, PROGRESS_UPDATE, PROGRESS_COMPLETE, PROGRESS_ERROR '''
		self.update_progress = progresshook

	def get_distro_name(self):
		return self.distro_name

	def get_distro_version(self):
		return self.distro_version

	def get_locale(self):
		''' Return the locale we're setting '''
		return self.locale

	def set_locale(self, newlocale):
		''' Set the locale '''
		self.locale = newlocale

	def install(self):
		''' Install this baby to disk '''
		# mount the media location.
		try:
			if(not os.path.exists("/target")):
				os.mkdir("/target")
			if(not os.path.exists("/source")):
				os.mkdir("/source")
			# find the squashfs..
			root = self.media
			root_type = self.media_type
			if(not os.path.exists(root)):
				print _("Base filesystem does not exist! Bailing")
				sys.exit(1) # change to report
			root_device = None
			# format partitions as appropriate
			for item in self.fstab.get_entries():
				if(item.mountpoint == "/"):
					root_device = item
					item.format = True
				if(item.format):
					# well now, we gets to nuke stuff.
					# report it. should grab the total count of filesystems to be formatted ..
					self.update_progress(total=4, current=1, message=_("Formatting %s to %s..." % (item.device, item.filesystem)))
					self.format_device(item.device, item.filesystem)
			# mount filesystem
			self.update_progress(total=4, current=2, message=_("Mounting %s on %s") % (root, "/source/"))
			self.do_mount(root, "/source/", root_type, options="loop")
			self.update_progress(total=4, current=3, message=_("Mounting %s on %s") % (root_device.device, r"/target/"))
			self.do_mount(root_device.device, "/target", root_device.filesystem, None)
			# walk root filesystem. we're too lazy though :P
			SOURCE = "/source/"
			DEST = "/target/"
			timestamps = list()
			our_total = -1
			our_current = 0
			os.chdir(SOURCE)
			# index the files
			for top,dirs,files in os.walk(SOURCE, topdown=False):
				our_total += len(files)
				self.update_progress(pulse=True, message=_("Indexing files to be copied.."))
			our_total += 1 # safenessness
			for top,dirs,files in os.walk(SOURCE, topdown=False):
				for f in files:
					try:
						our_current += 1
						filename = os.path.join(top, f)
						rpath = os.path.relpath(filename)
						target = os.path.join(DEST, rpath)
						dir = os.path.split(target)[0]
						# make the directories..
						if(not os.path.exists(dir)):
							os.makedirs(dir)
						if(os.path.islink(filename)):
							try:
								link_dest = os.readlink(filename)
								if(link_dest != filename):
									# oh look, a real symlink.
									os.chdir(dir) # make sure relative symlinks work..
									newfile = os.path.split(target)[1] # get the tail end
									#os.symlink(link_dest, newfile)
									os.system("cp -d \"%s\" \"%s\"" % (filename, target)) # uber lazy.. mmhm
									os.chdir(SOURCE)
								continue
							except:
								pass # who cares.. wasn't a symlink.
						if(not os.path.exists(filename)):
							continue # skip broken links
						self.update_progress(total=our_total, current=our_current, message="Copying %s" % filename)
						self.copy_file(filename, target)
					except Exception, detail:
						print detail
				for d in dirs:
					try:
						directory = os.path.join(top, d)
						filename = directory
						rpath = os.path.relpath(directory)
						target = os.path.join(DEST, rpath)
						if(os.path.islink(directory)):
							try:
								link_dest = os.readlink(filename)
								if(link_dest != filename):
									# oh look, a real symlink.
									os.chdir(dir) # make sure relative symlinks work..
									newfile = os.path.split(target)[1] # get the tail end
									#os.symlink(link_dest, newfile)
									os.system("cp -d \"%s\" \"%s\"" % (filename, target)) # uber fail.
									os.chdir(SOURCE)
								continue
							except:
								pass # who cares.. wasn't a symlink.
						if(not os.path.exists(directory)):
							continue # skip broken links
						if(not os.path.exists(target)):
							os.makedirs(target)
						# copy dir attributes
						finfo = os.stat(filename)
						owner = finfo[stat.ST_UID]
						group = finfo[stat.ST_GID]
						os.chown(target, owner, group)
						shutil.copymode(filename, target) # copy permissions.
						timestamps.append((target, finfo))
						pass
					except Exception, detail:
						print detail
			for directories in timestamps:
				directory = directories[0]
				timestamp = directories[1]
				self.update_progress(pulse=True, message=_("Restoring meta-info on %s") % directory)
				os.utime(directory, (timestamp.st_ctime, timestamp.st_mtime))

			# Steps:
			our_total = 8
			our_current = 0
			# chroot
			self.update_progress(total=our_total, current=our_current, message=_("Entering new system.."))
			child_pid = os.fork()
			if(child_pid == 0):
				# we be the child.
				os.chroot("/target/")
				# remove live user
				live_user = "mint" # make this come from config..
				our_current += 1
				self.update_progress(total=our_total, current=our_current, message=_("Removing live configuration (user)"))
				os.system("deluser %s" % live_user)
				os.system("rm -rf /home/%s" % live_user)
				# remove live-initramfs (or w/e)
				our_current += 1
				self.update_progress(total=our_total, current=our_current, message=_("Removing live configuration (packages)"))
				os.system("apt-get remove --purge live-initramfs distro-installer")
				# add new user
				our_current += 1
				self.update_progress(total=our_total, current=our_current, message=_("Adding new user to system"))
				newusers = open("/tmp/newusers.conf", "w")
				shell = "/bin/bash"
				username = "ikey"
				password = "czkgj8la"
				root_password = "czkgj8la"
				shell = "/bin/bash"
				realname = "Ikey Doherty"
				newusers.write("%s:%s:::%s,,,:/home/%s:%s\n" % (username, password, realname, username, shell))
				# add root's password
				our_current += 1
				self.update_progress(total=our_total, current=our_current, message=_("Setting root password"))
				newusers.write("root:%s:::::\n") % (root_password)
				newusers.close()
				os.system("newusers /tmp/newusers.conf")
				os.system("rm -rf /tmp/newusers.conf")
				# write the /etc/fstab
				our_current += 1
				self.update_progress(total=our_total, current=our_current, message=_("Writing filesystem mount information"))
				# make sure fstab has default /proc and /sys entries
				if(not os.path.exists("%s/etc/fstab" % "/target")):
					os.system("echo \"#### Static Filesystem Table File\" > %s/etc/fstab" % "/target")
				fstabber = open("%s/etc/fstab" % "/target", "a")
				fstabber.write("proc\t/proc\tproc\tnodev,noexec,nosuid\t0\t0\n")
				for item in self.fstab.get_entries():
					fstabber.write("%s\t%s\t%s\t%s\t%s\t%s" % (item.device, item.mount_point, item.filesystem, item.options, "0", "0"))
				fstabber.close()
				# write MBR (grub)
				our_current += 1
				self.update_progress(total=our_total, current=our_current, message=_("Installing bootloader"))
				os.system("grub-install %s" % root_device.device)
				os.system("update-grub")
				# ^ needs to be done in a chroot ._.
				# notify that we be finished now.
				our_current += 1
				self.update_progress(done=True, total=our_total, current=our_current, message=_("Done."))
				# write a file to this chroot indicating we've finished..
				fdone = open("/installed.log", "w")
				fdone.write("---INSTALLED---")
				fdone.close()
			else:
				while(True):
					# check for /installed.log in the chroot
					if(not os.path.exists(os.path.join(DEST, "installed.log"))):
						time.sleep(0.5) # sleep half a second then check again
					else:
						os.remove(os.path.join(DEST, "installed.log"))
						break

			# now unmount it
			self.do_unmount("/target")
			self.do_unmount("/source")
		except Exception,detail:
			print detail

	def do_mount(self, device, dest, type, options=None):
		''' Mount a filesystem '''
		p = None
		if(options is not None):
			p = Popen("mount -o %s -t %s %s %s" % (options, type, device, dest),shell=True)
		else:
			p = Popen("mount -t %s %s %s" % (type, device, dest),shell=True)
		p.wait()
		return p.returncode

	def do_unmount(self, mountpoint):
		''' Unmount a filesystem '''
		p = Popen("umount %s" % mountpoint, shell=True)
		p.wait()
		return p.returncode

	def copy_file(self, source, dest):
		'''
		BUF_SIZE = 16 * 1048
		input = open(source, "rb")
		dst = open(dest, "wb")
		while(True):
			read = input.read(BUF_SIZE)
			if not read:
				break
			dst.write(read)
		input.close()
		fd = dst.fileno()
		# set permissions
		shutil.copymode(source, dest) # copy permission bits
		finfo = os.stat(source)
		owner = finfo[stat.ST_UID]
		group = finfo[stat.ST_GID]
		os.fchown(fd, owner, group)
		dst.flush()
		os.fsync(fd)
		dst.close()
		finfo = os.stat(source)
		atime = finfo[stat.ST_ATIME]
		mtime = finfo[stat.ST_MTIME]
		os.utime(dest, (atime, mtime))
		'''
		os.system("cp -p \"%s\" \"%s\"" % (source, dest))

class fstab(object):
	''' This represents the filesystem table (/etc/fstab) '''
	def __init__(self):
		self.mapping = dict()

	def add_mount(self, device=None, mountpoint=None, filesystem=None, options=None):
		if(not self.mapping.has_key(device)):
			self.mapping[device] = fstab_entry(device, mountpoint, filesystem, options)

	def remove_mount(self, device):
		if(self.mapping.has_key(device)):
			del self.mapping[device]

	def get_entries(self):
		''' Return our list '''
		return self.mapping.values()

class fstab_entry(object):
	''' Represents an entry in fstab '''

	def __init__(self, device, mountpoint, filesystem, options):
		''' Creates a new fstab entry '''
		self.device = device
		self.mountpoint = mountpoint
		self.filesystem = filesystem
		self.options = options
		self.format = False