Merge pull request #393 from archlinux/torxed-plugins

Adding in plugin support and on_<event> calls in installation.

3 files changed, 165 insertions, 2 deletions

diff --git a/archinstall/__init__.py b/archinstall/__init__.py
index 948dc070..fcd9a706 100644
--- a/archinstall/__init__.py
+++ b/archinstall/__init__.py
@@ -24,6 +24,7 @@ from .lib.user_interaction import *
 parser = ArgumentParser()
 
 __version__ = "2.3.0.dev0"
+storage['__version__'] = __version__
 
 
 def initialize_arguments():
@@ -60,7 +61,10 @@ def initialize_arguments():
 
 
 arguments = initialize_arguments()
+from .lib.plugins import plugins, load_plugin # This initiates the plugin loading ceremony
 
+if arguments.get('plugin', None):
+	load_plugin(arguments['plugin'])
 
 # TODO: Learn the dark arts of argparse... (I summon thee dark spawn of cPython)
 
diff --git a/archinstall/lib/installer.py b/archinstall/lib/installer.py
index 0044532f..c5e77b7e 100644
--- a/archinstall/lib/installer.py
+++ b/archinstall/lib/installer.py
@@ -3,6 +3,7 @@ from .hardware import *
 from .locale_helpers import verify_keyboard_layout, verify_x11_keyboard_layout
 from .mirrors import *
 from .storage import storage
+from .plugins import plugins
 from .user_interaction import *
 
 # Any package that the Installer() is responsible for (optional and the default ones)
@@ -132,6 +133,12 @@ class Installer:
 	def pacstrap(self, *packages, **kwargs):
 		if type(packages[0]) in (list, tuple):
 			packages = packages[0]
+
+		for plugin in plugins.values():
+			if hasattr(plugin, 'on_pacstrap'):
+				if (result := plugin.on_pacstrap(packages)):
+					packages = result
+
 		self.log(f'Installing packages: {packages}', level=logging.INFO)
 
 		if (sync_mirrors := SysCommand('/usr/bin/pacman -Syy')).exit_code == 0:
@@ -143,6 +150,11 @@ class Installer:
 			self.log(f'Could not sync mirrors: {sync_mirrors.exit_code}', level=logging.INFO)
 
 	def set_mirrors(self, mirrors):
+		for plugin in plugins.values():
+			if hasattr(plugin, 'on_mirrors'):
+				if result := plugin.on_mirrors(mirrors):
+					mirrors = result
+
 		return use_mirrors(mirrors, destination=f'{self.target}/etc/pacman.d/mirrorlist')
 
 	def genfstab(self, flags='-pU'):
@@ -154,6 +166,10 @@ class Installer:
 		if not os.path.isfile(f'{self.target}/etc/fstab'):
 			raise RequirementError(f'Could not generate fstab, strapping in packages most likely failed (disk out of space?)\n{fstab}')
 
+		for plugin in plugins.values():
+			if hasattr(plugin, 'on_genfstab'):
+				plugin.on_genfstab(self)
+
 		return True
 
 	def set_hostname(self, hostname: str, *args, **kwargs):
@@ -177,6 +193,11 @@ class Installer:
 		if not len(zone):
 			return True  # Redundant
 
+		for plugin in plugins.values():
+			if hasattr(plugin, 'on_timezone'):
+				if result := plugin.on_timezone(zone):
+					zone = result
+
 		if (pathlib.Path("/usr") / "share" / "zoneinfo" / zone).exists():
 			(pathlib.Path(self.target) / "etc" / "localtime").unlink(missing_ok=True)
 			SysCommand(f'/usr/bin/arch-chroot {self.target} ln -s /usr/share/zoneinfo/{zone} /etc/localtime')
@@ -200,6 +221,10 @@ class Installer:
 			if (output := self.arch_chroot(f'systemctl enable {service}')).exit_code != 0:
 				raise ServiceException(f"Unable to start service {service}: {output}")
 
+			for plugin in plugins.values():
+				if hasattr(plugin, 'on_service'):
+					plugin.on_service(service)
+
 	def run_command(self, cmd, *args, **kwargs):
 		return SysCommand(f'/usr/bin/arch-chroot {self.target} {cmd}')
 
@@ -229,6 +254,11 @@ class Installer:
 
 			conf = Networkd(Match={"Name": nic}, Network=network)
 
+		for plugin in plugins.values():
+			if hasattr(plugin, 'on_configure_nic'):
+				if (new_conf := plugin.on_configure_nic(nic, dhcp, ip, gateway, dns)):
+					conf = new_conf
+
 		with open(f"{self.target}/etc/systemd/network/10-{nic}.network", "a") as netconf:
 			netconf.write(str(conf))
 
@@ -292,6 +322,12 @@ class Installer:
 		return False
 
 	def mkinitcpio(self, *flags):
+		for plugin in plugins.values():
+			if hasattr(plugin, 'on_mkinitcpio'):
+				# Allow plugins to override the usage of mkinitcpio altogether.
+				if plugin.on_mkinitcpio(self):
+					return True
+
 		with open(f'{self.target}/etc/mkinitcpio.conf', 'w') as mkinit:
 			mkinit.write(f"MODULES=({' '.join(self.MODULES)})\n")
 			mkinit.write(f"BINARIES=({' '.join(self.BINARIES)})\n")
@@ -366,9 +402,20 @@ class Installer:
 			self.log(f"Running post-installation hook: {function}", level=logging.INFO)
 			function(self)
 
+		for plugin in plugins.values():
+			if hasattr(plugin, 'on_install'):
+				plugin.on_install(self)
+
 		return True
 
 	def add_bootloader(self, bootloader='systemd-bootctl'):
+		for plugin in plugins.values():
+			if hasattr(plugin, 'on_add_bootloader'):
+				# Allow plugins to override the boot-loader handling.
+				# This allows for bot configuring and installing bootloaders.
+				if plugin.on_add_bootloader(self):
+					return True
+
 		boot_partition = None
 		root_partition = None
 		for partition in self.partitions:
@@ -487,8 +534,19 @@ class Installer:
 	def user_create(self, user: str, password=None, groups=None, sudo=False):
 		if groups is None:
 			groups = []
-		self.log(f'Creating user {user}', level=logging.INFO)
-		o = b''.join(SysCommand(f'/usr/bin/arch-chroot {self.target} useradd -m -G wheel {user}'))
+
+		# This plugin hook allows for the plugin to handle the creation of the user.
+		# Password and Group management is still handled by user_create()
+		handled_by_plugin = False
+		for plugin in plugins.values():
+			if hasattr(plugin, 'on_user_create'):
+				if result := plugin.on_user_create(user):
+					handled_by_plugin = result
+
+		if not handled_by_plugin:
+			self.log(f'Creating user {user}', level=logging.INFO)
+			o = b''.join(SysCommand(f'/usr/bin/arch-chroot {self.target} useradd -m -G wheel {user}'))
+
 		if password:
 			self.user_set_pw(user, password)
 
diff --git a/archinstall/lib/plugins.py b/archinstall/lib/plugins.py
new file mode 100644
index 00000000..a61be30b
--- /dev/null
+++ b/archinstall/lib/plugins.py
@@ -0,0 +1,101 @@
+import hashlib
+import importlib
+import logging
+import os
+import sys
+import pathlib
+import urllib.parse
+import urllib.request
+from importlib import metadata
+
+from .output import log
+from .storage import storage
+
+plugins = {}
+
+# 1: List archinstall.plugin definitions
+# 2: Load the plugin entrypoint
+# 3: Initiate the plugin and store it as .name in plugins
+for plugin_definition in metadata.entry_points()['archinstall.plugin']:
+	plugin_entrypoint = plugin_definition.load()
+	try:
+		plugins[plugin_definition.name] = plugin_entrypoint()
+	except Exception as err:
+		log(err, level=logging.ERROR)
+		log(f"The above error was detected when loading the plugin: {plugin_definition}", fg="red", level=logging.ERROR)
+
+
+# The following functions and core are support structures for load_plugin()
+def localize_path(profile_path :str) -> str:
+	if (url := urllib.parse.urlparse(profile_path)).scheme and url.scheme in ('https', 'http'):
+		converted_path = f"/tmp/{os.path.basename(profile_path).replace('.py', '')}_{hashlib.md5(os.urandom(12)).hexdigest()}.py"
+
+		with open(converted_path, "w") as temp_file:
+			temp_file.write(urllib.request.urlopen(url.geturl()).read().decode('utf-8'))
+
+		return converted_path
+	else:
+		return profile_path
+
+
+def import_via_path(path :str, namespace=None): # -> module (not sure how to write that in type definitions)
+	if not namespace:
+		namespace = os.path.basename(path)
+
+		if namespace == '__init__.py':
+			path = pathlib.PurePath(path)
+			namespace = path.parent.name
+
+	try:
+		spec = importlib.util.spec_from_file_location(namespace, path)
+		imported = importlib.util.module_from_spec(spec)
+		sys.modules[namespace] = imported
+		spec.loader.exec_module(sys.modules[namespace])
+
+		return namespace
+	except Exception as err:
+		log(err, level=logging.ERROR)
+		log(f"The above error was detected when loading the plugin: {path}", fg="red", level=logging.ERROR)
+
+		try:
+			del(sys.modules[namespace])
+		except:
+			pass
+
+def find_nth(haystack, needle, n):
+	start = haystack.find(needle)
+	while start >= 0 and n > 1:
+		start = haystack.find(needle, start+len(needle))
+		n -= 1
+	return start
+
+def load_plugin(path :str): # -> module (not sure how to write that in type definitions)
+	parsed_url = urllib.parse.urlparse(path)
+
+	# The Profile was not a direct match on a remote URL
+	if not parsed_url.scheme:
+		# Path was not found in any known examples, check if it's an absolute path
+		if os.path.isfile(path):
+			namespace = import_via_path(path)
+	elif parsed_url.scheme in ('https', 'http'):
+		namespace = import_via_path(localize_path(path))
+
+	if namespace in sys.modules:
+		# Version dependency via __archinstall__version__ variable (if present) in the plugin
+		# Any errors in version inconsistency will be handled through normal error handling if not defined.
+		if hasattr(sys.modules[namespace], '__archinstall__version__'):
+			archinstall_major_and_minor_version = float(storage['__version__'][:find_nth(storage['__version__'], '.', 2)])
+
+			if sys.modules[namespace].__archinstall__version__ < archinstall_major_and_minor_version:
+				log(f"Plugin {sys.modules[namespace]} does not support the current Archinstall version.", fg="red", level=logging.ERROR)
+
+		# Locate the plugin entry-point called Plugin()
+		# This in accordance with the entry_points() from setup.cfg above
+		if hasattr(sys.modules[namespace], 'Plugin'):
+			try:
+				plugins[namespace] = sys.modules[namespace].Plugin()
+			except Exception as err:
+				log(err, level=logging.ERROR)
+				log(f"The above error was detected when initiating the plugin: {path}", fg="red", level=logging.ERROR)
+		else:
+			log(f"Plugin '{path}' is missing a valid entry-point or is corrupt.", fg="yellow", level=logging.WARNING)
\ No newline at end of file