Added a PArtition() class that supports mounting and formatting. Also reworked the installation flow a bit to be a bit more clear while sacrificing some automation. Maybe I'll revert some changes and 'automatically' do certain things, but for now this shouldn't impact anyone to much

4 files changed, 130 insertions, 38 deletions

diff --git a/archinstall.py b/archinstall.py
index c8df36c2..ce108d8b 100644
--- a/archinstall.py
+++ b/archinstall.py
@@ -5,17 +5,53 @@ from helpers.disk import *
 from helpers.general import *
 from helpers.user_interaction import *
 
-class HardDrive():
-	def __init__(self, full_path:str, *args, **kwargs):
-		if not stat.S_ISBLK(os.stat(full_path).st_mode):
-			raise DiskError(f'Selected disk "{full_path}" is not a block device.')
-
-class installer():
-	def __init__(self, partition, *, profile=None, hostname='ArchInstalled'):
+class Installer():
+	def __init__(self, partition, *, profile=None, mountpoint='/mnt', hostname='ArchInstalled'):
 		self.profile = profile
 		self.hostname = hostname
+		self.mountpoint = mountpoint
 
 		self.partition = partition
 
+	def __enter__(self, *args, **kwargs):
+		self.partition.mount(self.mountpoint)
+		return self
+
+	def __exit__(self, *args, **kwargs):
+		# b''.join(sys_command(f'sync')) # No need to, since the underlaying fs() object will call sync.
+		# TODO: https://stackoverflow.com/questions/28157929/how-to-safely-handle-an-exception-inside-a-context-manager
+		if len(args) >= 2 and args[1]:
+			raise args[1]
+		return True
+
 	def minimal_installation(self):
-		pass
\ No newline at end of file
+		if (sync_mirrors := sys_command('/usr/bin/pacman -Syy')).exit_code == 0:
+			if (pacstrap := sys_command(f'/usr/bin/pacstrap {self.mountpoint} base base-devel linux linux-firmware btrfs-progs efibootmgr nano wpa_supplicant dialog')).exit_code == 0:
+				return True
+			else:
+				log(f'Could not strap in base: {pacstrap.exit_code}')
+		else:
+			log(f'Could not sync mirrors: {sync_mirrors.exit_code}')
+
+	def add_bootloader(self, partition):
+		os.makedirs(f'{self.mountpoint}/boot', exist_ok=True)
+		partition.mount(f'{self.mountpoint}/boot')
+		o = b''.join(sys_command(f'/usr/bin/arch-chroot {self.mountpoint} bootctl --no-variables --path=/boot install'))
+
+		with open('/mnt/boot/loader/loader.conf', 'w') as loader:
+			loader.write('default arch\n')
+			loader.write('timeout 5\n')
+
+		## For some reason, blkid and /dev/disk/by-uuid are not getting along well.
+		## And blkid is wrong in terms of LUKS.
+		#UUID = sys_command('blkid -s PARTUUID -o value {drive}{partition_2}'.format(**args)).decode('UTF-8').strip()
+		with open('/mnt/boot/loader/entries/arch.conf', 'w') as entry:
+			entry.write('title Arch Linux\n')
+			entry.write('linux /vmlinuz-linux\n')
+			entry.write('initrd /initramfs-linux.img\n')
+			## blkid doesn't trigger on loopback devices really well,
+			## so we'll use the old manual method until we get that sorted out.
+			# UUID = simple_command(f"blkid -s PARTUUID -o value /dev/{os.path.basename(args['drive'])}{args['partitions']['2']}").decode('UTF-8').strip()
+			# entry.write('options root=PARTUUID={UUID} rw intel_pstate=no_hwp\n'.format(UUID=UUID))
+			UUID = b''.join(sys_command(f"ls -l /dev/disk/by-uuid/ | grep {os.path.basename(partition['path'])} | awk '{{print $9}}'")).decode('UTF-8').strip()
+			entry.write(f'options cryptdevice=UUID={UUID}:luksdev root=/dev/mapper/luksdev rw intel_pstate=no_hwp\n')
\ No newline at end of file
diff --git a/helpers/disk.py b/helpers/disk.py
index 1c0a544c..e3129d37 100644
--- a/helpers/disk.py
+++ b/helpers/disk.py
@@ -2,9 +2,15 @@ import glob, re, os, json
 from collections import OrderedDict
 from helpers.general import sys_command
 from exceptions import *
+import ctypes
+import ctypes.util
+import os
 
 ROOT_DIR_PATTERN = re.compile('^.*?/devices')
 GPT = 0b00000001
+libc = ctypes.CDLL(ctypes.util.find_library('c'), use_errno=True)
+libc.mount.argtypes = (ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p, ctypes.c_ulong, ctypes.c_char_p)
+
 
 class BlockDevice():
 	def __init__(self, path, info):
@@ -32,6 +38,9 @@ class BlockDevice():
 			if not 'pkname' in self.info: raise DiskError(f'A crypt device ({self.path}) without a parent kernel device name.')
 			return f"/dev/{self.info['pkname']}"
 
+	#	if not stat.S_ISBLK(os.stat(full_path).st_mode):
+	#		raise DiskError(f'Selected disk "{full_path}" is not a block device.')
+
 	@property
 	def partitions(self):
 		o = b''.join(sys_command(f'partprobe {self.path}'))
@@ -50,11 +59,7 @@ class BlockDevice():
 			root_path = f"/dev/{r['blockdevices'][0]['name']}"
 			for part in r['blockdevices'][0]['children']:
 				part_id = part['name'][len(os.path.basename(self.path)):]
-				parts[part_id] = {
-					'size' : part['size'],
-					'id' : part_id,
-					'path' : root_path + part_id
-				}
+				parts[part_id] = Partition(root_path + part_id, part_id=part_id, size=part['size'])
 
 		return {k: parts[k] for k in sorted(parts)}
 
@@ -71,6 +76,45 @@ class BlockDevice():
 			raise KeyError(f'{self} does not contain information: "{key}"')
 		return self.info[key]
 
+class Partition():
+	def __init__(self, path, part_id=None, size=-1):
+		if not part_id: part_id = os.path.basename(path)
+		self.path = path
+		self.part_id = part_id
+		self.mountpoint = None
+		self.filesystem = None # TODO: Autodetect if we're reusing a partition
+		self.size = size # TODO: Refresh?
+
+	def __repr__(self, *args, **kwargs):
+		return f'Partition({self.path})'
+
+	def format(self, filesystem):
+		if filesystem == 'btrfs':
+			o = b''.join(sys_command(f'/usr/bin/mkfs.btrfs -f {self.path}'))
+			if not b'UUID' in o:
+				return False
+			self.filesystem = 'btrfs'
+		elif filesystem == 'fat32':
+			o = b''.join(sys_command(f'/usr/bin/mkfs.vfat -F32 {self.path}'))
+			if (b'mkfs.fat' not in o and b'mkfs.vfat' not in o) or b'command not found' in o:
+				return None
+			return True
+		else:
+			raise DiskError(f'Fileformat {filesystem} is not yet implemented.')
+		return True
+
+	def mount(self, target, fs=None, options=''):
+		if not fs:
+			if not self.filesystem: raise DiskError('Need to format (or define) the filesystem before mounting.')
+			fs = self.filesystem
+		# TODO: Move this to the BlockDevice or something.
+		ret = libc.mount(self.path.encode(), target.encode(), fs.encode(), 0, options.encode())
+		if ret < 0:
+			errno = ctypes.get_errno()
+			raise OSError(errno, f"Error mounting {self.path} ({fs}) on {target} with options '{options}': {os.strerror(errno)}")
+		self.mountpoint = target
+
+
 class luks2():
 	def __init__(self, filesystem):
 		self.filesystem = filesystem
@@ -80,7 +124,7 @@ class luks2():
 
 	def __exit__(self, *args, **kwargs):
 		# TODO: https://stackoverflow.com/questions/28157929/how-to-safely-handle-an-exception-inside-a-context-manager
-		if len(args):
+		if len(args) >= 2 and args[1]:
 			raise args[1]
 		return True
 
@@ -91,13 +135,13 @@ class luks2():
 		with open(key_file, 'wb') as fh:
 			fh.write(password)
 
-		o = b''.join(sys_command(f'/usr/bin/cryptsetup -q -v --type luks2 --pbkdf argon2i --hash {hash_type} --key-size {key_size} --iter-time {iter_time} --key-file {os.path.abspath(key_file)} --use-urandom luksFormat {partition["path"]}'))
+		o = b''.join(sys_command(f'/usr/bin/cryptsetup -q -v --type luks2 --pbkdf argon2i --hash {hash_type} --key-size {key_size} --iter-time {iter_time} --key-file {os.path.abspath(key_file)} --use-urandom luksFormat {partition.path}'))
 		if not b'Command successful.' in o:
-			raise DiskError(f'Could not encrypt volume "{partition["path"]}": {o}')
+			raise DiskError(f'Could not encrypt volume "{partition.path}": {o}')
 	
 		return key_file
 
-	def mount(self, partition, mountpoint, key_file):
+	def unlock(self, partition, mountpoint, key_file):
 		"""
 		Mounts a lukts2 compatible partition to a certain mountpoint.
 		Keyfile must be specified as there's no way to interact with the pw-prompt atm.
@@ -106,8 +150,9 @@ class luks2():
 		:type mountpoint: str
 		"""
 		if '/' in mountpoint: os.path.basename(mountpoint) # TODO: Raise exception instead?
-		sys_command(f'/usr/bin/cryptsetup open {partition["path"]} {mountpoint} --key-file {os.path.abspath(key_file)} --type luks2')
-		return os.path.islink(f'/dev/mapper/{mountpoint}')
+		sys_command(f'/usr/bin/cryptsetup open {partition.path} {mountpoint} --key-file {os.path.abspath(key_file)} --type luks2')
+		if os.path.islink(f'/dev/mapper/{mountpoint}'):
+			return Partition(f'/dev/mapper/{mountpoint}')
 
 	def close(self, mountpoint):
 		sys_command(f'cryptsetup close /dev/mapper/{mountpoint}')
diff --git a/helpers/general.py b/helpers/general.py
index 036d2d68..32814ddc 100644
--- a/helpers/general.py
+++ b/helpers/general.py
@@ -16,6 +16,7 @@ class sys_command():#Thread):
 	def __init__(self, cmd, callback=None, start_callback=None, *args, **kwargs):
 		if not 'worker_id' in kwargs: kwargs['worker_id'] = gen_uid()
 		if not 'emulate' in kwargs: kwargs['emulate'] = False
+		if not 'surpress_errors' in kwargs: kwargs['surpress_errors'] = False
 		if kwargs['emulate']:
 			log(f"Starting command '{cmd}' in emulation mode.")
 		self.raw_cmd = cmd
@@ -170,11 +171,9 @@ class sys_command():#Thread):
 		if 'ignore_errors' in self.kwargs:
 			self.exit_code = 0
 
-		if self.exit_code != 0:
+		if self.exit_code != 0 and not self.kwargs['surpress_errors']:
 			log(f"'{self.raw_cmd}' did not exit gracefully, exit code {self.exit_code}.", origin='spawn', level=3)
 			log(self.trace_log.decode('UTF-8'), origin='spawn', level=3)
-		#else:
-			#log(f"{self.cmd[0]} exit nicely.", origin='spawn', level=5)
 
 		self.ended = time.time()
 		with open(f'{self.cwd}/trace.log', 'wb') as fh:
diff --git a/installer.py b/installer.py
index 91a8139a..faeb7edc 100644
--- a/installer.py
+++ b/installer.py
@@ -1,28 +1,40 @@
 import archinstall, getpass
 
-archinstall.sys_command(f'cryptsetup close /dev/mapper/luksloop')
+## dd if=/dev/zero of=test.img bs=1G count=4
+## losetup -fP test.img
+archinstall.sys_command(f'umount -R /mnt', surpress_errors=True)
+archinstall.sys_command(f'cryptsetup close /dev/mapper/luksloop', surpress_errors=True)
 
-#selected_hdd = archinstall.select_disk(archinstall.all_disks())
-selected_hdd = archinstall.all_disks()['/dev/loop0']
+#harddrive = archinstall.select_disk(archinstall.all_disks())
+harddrive = archinstall.all_disks()['/dev/loop0']
 disk_password = '1234' # getpass.getpass(prompt='Disk password (won\'t echo): ')
 
-with archinstall.Filesystem(selected_hdd, archinstall.GPT) as fs:
+with archinstall.Filesystem(harddrive, archinstall.GPT) as fs:
+	print(f'Formatting {harddrive}')
 	fs.use_entire_disk('luks2')
 	with archinstall.luks2(fs) as crypt:
-		if selected_hdd.partition[1]['size'] == '512M':
+		if harddrive.partition[1].size == '512M':
 			raise OSError('Trying to encrypt the boot partition for petes sake..')
 
-		key_file = crypt.encrypt(selected_hdd.partition[1], password=disk_password, key_size=512, hash_type='sha512', iter_time=10000, key_file='./pwfile')
-		crypt.mount(selected_hdd.partition[1], 'luksloop', key_file)
-	exit(1)
-	with archinstall.installer(root_partition, hostname='testmachine') as installation:
-		if installation.minimal_installation():
-			installation.add_bootloader()
+		print(f'Encrypting {harddrive.partition[1]}')
+		key_file = crypt.encrypt(harddrive.partition[1], password=disk_password, key_size=512, hash_type='sha512', iter_time=10000, key_file='./pwfile')
 
-			installation.add_additional_packages(['nano', 'wget', 'git'])
-			installation.install_profile('desktop')
+		unlocked_device = crypt.unlock(harddrive.partition[1], 'luksloop', key_file)
+		
+		print('Formatting partitions.')
+		harddrive.partition[0].format('fat32')
+		unlocked_device.format('btrfs')
+		
+		with archinstall.Installer(unlocked_device, hostname='testmachine') as installation:
+			print('Installing minimal installation to disk.')
+			if installation.minimal_installation():
+				print('Adding bootloader.')
+				installation.add_bootloader(harddrive.partition[0])
 
-			installation.user_create('anton', 'test')
-			installation.user_set_pw('root', 'toor')
+				installation.add_additional_packages(['nano', 'wget', 'git'])
+				installation.install_profile('desktop')
 
-			installation.add_AUR_support()
\ No newline at end of file
+				installation.user_create('anton', 'test')
+				installation.user_set_pw('root', 'toor')
+
+				installation.add_AUR_support()
\ No newline at end of file