Added a set_timezone() and fixed set_locale() in the Installer() class. Also added a mirrors.py helper to rudimentary set mirror data on the installer host

3 files changed, 59 insertions, 4 deletions

diff --git a/archinstall/__init__.py b/archinstall/__init__.py
index 9cf7faec..7042322c 100644
--- a/archinstall/__init__.py
+++ b/archinstall/__init__.py
@@ -4,4 +4,5 @@ from .lib.user_interaction import *
 from .lib.exceptions import *
 from .lib.installer import *
 from .lib.profiles import *
-from .lib.luks import *
\ No newline at end of file
+from .lib.luks import *
+from .lib.mirrors import *
diff --git a/archinstall/lib/installer.py b/archinstall/lib/installer.py
index 7f9aba71..90849bf2 100644
--- a/archinstall/lib/installer.py
+++ b/archinstall/lib/installer.py
@@ -70,17 +70,26 @@ class Installer():
 			raise RequirementError(f'Could not generate fstab, strapping in packages most likely failed (disk out of space?)\n{o}')
 		return True
 
-	def set_hostname(self, hostname=None):
+	def set_hostname(self, hostname=None, *args, **kwargs):
 		if not hostname: hostname = self.hostname
 		with open(f'{self.mountpoint}/etc/hostname', 'w') as fh:
 			fh.write(self.hostname + '\n')
 
-	def set_locale(self, locale, encoding='UTF-8'):
+	def set_locale(self, locale, encoding='UTF-8', *args, **kwargs):
+		if not len(locale): return True
+
 		with open(f'{self.mountpoint}/etc/locale.gen', 'a') as fh:
 			fh.write(f'{locale} {encoding}\n')
 		with open(f'{self.mountpoint}/etc/locale.conf', 'w') as fh:
 			fh.write(f'LANG={locale}\n')
-		sys_command(f'/usr/bin/arch-chroot {self.mountpoint} locale-gen')
+
+		return True if sys_command(f'/usr/bin/arch-chroot {self.mountpoint} locale-gen').exit_code == 0 else False
+
+	def set_timezone(self, zone, *args, **kwargs):
+		if not len(zone): return True
+
+		o = b''.join(sys_command(f'/usr/bin/arch-chroot {self.mountpoint} ln -s /usr/share/zoneinfo/{zone} /etc/localtime'))
+		return True
 
 	def minimal_installation(self):
 		self.pacstrap('base base-devel linux linux-firmware btrfs-progs efibootmgr nano'.split(' '))
diff --git a/archinstall/lib/mirrors.py b/archinstall/lib/mirrors.py
new file mode 100644
index 00000000..63eb4678
--- /dev/null
+++ b/archinstall/lib/mirrors.py
@@ -0,0 +1,45 @@
+from .exceptions import *
+from .general import *
+
+def filter_mirrors_by_region(regions, *args, **kwargs):
+	"""
+	This function will change the active mirrors on the live medium by
+	filtering which regions are active based on `regions`.
+
+	:param region: A series of country codes separated by `,`. For instance `SE,US` for sweden and United States.
+	:type region: str
+	"""
+	region_list = []
+	for region in regions.split(','):
+		region_list.append(f'country={region}')
+	o = b''.join(sys_command((f"/usr/bin/wget 'https://www.archlinux.org/mirrorlist/?{'&'.join(region_list)}&protocol=https&ip_version=4&ip_version=6&use_mirror_status=on' -O /root/mirrorlist")))
+	o = b''.join(sys_command(("/usr/bin/sed -i 's/#Server/Server/' /root/mirrorlist")))
+	o = b''.join(sys_command(("/usr/bin/mv /root/mirrorlist /etc/pacman.d/")))
+	
+	return True
+
+def insert_mirrors(mirrors, *args, **kwargs):
+	"""
+	This function will insert a given mirror-list at the top of `/etc/pacman.d/mirrorlist`.
+	It will not flush any other mirrors, just insert new ones.
+
+	:param mirrors: A dictionary of `{'url' : 'country', 'url2' : 'country'}`
+	:type mirrors: dict
+	"""
+	original_mirrorlist = ''
+	with open('/etc/pacman.d/mirrorlist', 'r') as original:
+		original_mirrorlist = original.read()
+
+	with open('/etc/pacman.d/mirrorlist', 'w') as new_mirrorlist:
+		for mirror, country in mirrors.items():
+			new_mirrorlist.write(f'## {country}\n')
+			new_mirrorlist.write(f'Server = {mirror}\n')
+		new_mirrorlist.write('\n')
+		new_mirrorlist.write(original_mirrorlist)
+
+	return True
+
+def re_rank_mirrors(top=10, *positionals, **kwargs):
+	if sys_command((f'/usr/bin/rankmirrors -n {top} /etc/pacman.d/mirrorlist > /etc/pacman.d/mirrorlist')).exit_code == 0:
+		return True
+	return False
\ No newline at end of file