Fix 2307 - Custom mirrors (#2350)

* Fix 2307 - Custom mirrors

* Update

4 files changed, 61 insertions, 49 deletions

diff --git a/archinstall/lib/installer.py b/archinstall/lib/installer.py
index 443e2b90..8a2dde5f 100644
--- a/archinstall/lib/installer.py
+++ b/archinstall/lib/installer.py
@@ -15,7 +15,7 @@ from .hardware import SysInfo
 from .locale import LocaleConfiguration
 from .locale import verify_keyboard_layout, verify_x11_keyboard_layout
 from .luks import Luks2
-from .mirrors import use_mirrors, MirrorConfiguration, add_custom_mirrors
+from .mirrors import MirrorConfiguration
 from .models.bootloader import Bootloader
 from .models.network_configuration import Nic
 from .models.users import User
@@ -318,17 +318,44 @@ class Installer:
 	def post_install_check(self, *args: str, **kwargs: str) -> List[str]:
 		return [step for step, flag in self.helper_flags.items() if flag is False]
 
-	def set_mirrors(self, mirror_config: MirrorConfiguration):
+	def set_mirrors(self, mirror_config: MirrorConfiguration, on_target: bool = False):
+		"""
+		Set the mirror configuration for the installation.
+
+		:param mirror_config: The mirror configuration to use.
+		:type mirror_config: MirrorConfiguration
+
+		:on_target: Whether to set the mirrors on the target system or the live system.
+		:param on_target: bool
+		"""
+		debug('Setting mirrors')
+
 		for plugin in plugins.values():
 			if hasattr(plugin, 'on_mirrors'):
 				if result := plugin.on_mirrors(mirror_config):
 					mirror_config = result
 
-		destination = f'{self.target}/etc/pacman.d/mirrorlist'
-		if mirror_config.mirror_regions:
-			use_mirrors(mirror_config.mirror_regions, destination)
-		if mirror_config.custom_mirrors:
-			add_custom_mirrors(mirror_config.custom_mirrors)
+		if on_target:
+			local_pacman_conf = Path(f'{self.target}/etc/pacman.conf')
+			local_mirrorlist_conf = Path(f'{self.target}/etc/pacman.d/mirrorlist')
+		else:
+			local_pacman_conf = Path('/etc/pacman.conf')
+			local_mirrorlist_conf = Path('/etc/pacman.d/mirrorlist')
+
+		mirrorlist_config = mirror_config.mirrorlist_config()
+		pacman_config = mirror_config.pacman_config()
+
+		if pacman_config:
+			debug(f'Pacman config: {pacman_config}')
+
+			with local_pacman_conf.open('a') as fp:
+				fp.write(pacman_config)
+
+		if mirrorlist_config:
+			debug(f'Mirrorlist: {mirrorlist_config}')
+
+			with local_mirrorlist_conf.open('a') as fp:
+				fp.write(mirrorlist_config)
 
 	def genfstab(self, flags: str = '-pU'):
 		fstab_path = self.target / "etc" / "fstab"
diff --git a/archinstall/lib/mirrors.py b/archinstall/lib/mirrors.py
index 74cdd0aa..61f3c568 100644
--- a/archinstall/lib/mirrors.py
+++ b/archinstall/lib/mirrors.py
@@ -5,7 +5,7 @@ from typing import Dict, Any, List, Optional, TYPE_CHECKING
 
 from .menu import AbstractSubMenu, Selector, MenuSelectionType, Menu, ListManager, TextInput
 from .networking import fetch_data_from_url
-from .output import info, warn, FormattedOutput
+from .output import warn, FormattedOutput
 from .storage import storage
 
 if TYPE_CHECKING:
@@ -77,6 +77,28 @@ class MirrorConfiguration:
 			'custom_mirrors': [c.json() for c in self.custom_mirrors]
 		}
 
+	def mirrorlist_config(self) -> str:
+		config = ''
+
+		for region, mirrors in self.mirror_regions.items():
+			for mirror in mirrors:
+				config += f'\n\n## {region}\nServer = {mirror}\n'
+
+		for cm in self.custom_mirrors:
+			config += f'\n\n## {cm.name}\nServer = {cm.url}\n'
+
+		return config
+
+	def pacman_config(self) -> str:
+		config = ''
+
+		for mirror in self.custom_mirrors:
+			config += f'\n\n[{mirror.name}]\n'
+			config += f'SigLevel = {mirror.sign_check.value} {mirror.sign_option.value}\n'
+			config += f'Server = {mirror.url}\n'
+
+		return config
+
 	@classmethod
 	def parse_args(cls, args: Dict[str, Any]) -> 'MirrorConfiguration':
 		config = MirrorConfiguration()
@@ -273,33 +295,6 @@ def select_custom_mirror(prompt: str = '', preset: List[CustomMirror] = []):
 	return custom_mirrors
 
 
-def add_custom_mirrors(mirrors: List[CustomMirror]):
-	"""
-	This will append custom mirror definitions in pacman.conf
-
-	:param mirrors: A list of custom mirrors
-	:type mirrors: List[CustomMirror]
-	"""
-	with open('/etc/pacman.conf', 'a') as pacman:
-		for mirror in mirrors:
-			pacman.write(f"\n\n[{mirror.name}]\n")
-			pacman.write(f"SigLevel = {mirror.sign_check.value} {mirror.sign_option.value}\n")
-			pacman.write(f"Server = {mirror.url}\n")
-
-
-def use_mirrors(
-	regions: Dict[str, List[str]],
-	destination: str = '/etc/pacman.d/mirrorlist'
-):
-	with open(destination, 'w') as fp:
-		for region, mirrors in regions.items():
-			for mirror in mirrors:
-				fp.write(f'## {region}\n')
-				fp.write(f'Server = {mirror}\n')
-
-	info(f'A new package mirror-list has been created: {destination}')
-
-
 def _parse_mirror_list(mirrorlist: str) -> Dict[str, List[str]]:
 	file_content = mirrorlist.split('\n')
 	file_content = list(filter(lambda x: x, file_content))  # filter out empty lines
diff --git a/archinstall/scripts/guided.py b/archinstall/scripts/guided.py
index 44b0ae17..f56ce5b4 100644
--- a/archinstall/scripts/guided.py
+++ b/archinstall/scripts/guided.py
@@ -10,7 +10,6 @@ from archinstall.lib.global_menu import GlobalMenu
 from archinstall.lib.configuration import ConfigurationOutput
 from archinstall.lib.installer import Installer
 from archinstall.lib.menu import Menu
-from archinstall.lib.mirrors import use_mirrors, add_custom_mirrors
 from archinstall.lib.models import AudioConfiguration
 from archinstall.lib.models.bootloader import Bootloader
 from archinstall.lib.models.network_configuration import NetworkConfiguration
@@ -132,12 +131,8 @@ def perform_installation(mountpoint: Path):
 				# generate encryption key files for the mounted luks devices
 				installation.generate_key_files()
 
-		# Set mirrors used by pacstrap (outside of installation)
 		if mirror_config := archinstall.arguments.get('mirror_config', None):
-			if mirror_config.mirror_regions:
-				use_mirrors(mirror_config.mirror_regions)
-			if mirror_config.custom_mirrors:
-				add_custom_mirrors(mirror_config.custom_mirrors)
+			installation.set_mirrors(mirror_config, on_target=False)
 
 		installation.minimal_installation(
 			testing=enable_testing,
@@ -148,7 +143,7 @@ def perform_installation(mountpoint: Path):
 		)
 
 		if mirror_config := archinstall.arguments.get('mirror_config', None):
-			installation.set_mirrors(mirror_config)  # Set the mirrors in the installation medium
+			installation.set_mirrors(mirror_config, on_target=True)
 
 		if archinstall.arguments.get('swap'):
 			installation.setup_swap('zram')
diff --git a/archinstall/scripts/swiss.py b/archinstall/scripts/swiss.py
index 1db304a8..8a5488bc 100644
--- a/archinstall/scripts/swiss.py
+++ b/archinstall/scripts/swiss.py
@@ -4,7 +4,6 @@ from typing import TYPE_CHECKING, Any, Dict, Optional
 
 import archinstall
 from archinstall import SysInfo, info, debug
-from archinstall.lib import mirrors
 from archinstall.lib import models
 from archinstall.lib import disk
 from archinstall.lib import locale
@@ -188,12 +187,8 @@ def perform_installation(mountpoint: Path, exec_mode: ExecutionMode):
 					# generate encryption key files for the mounted luks devices
 					installation.generate_key_files()
 
-			# Set mirrors used by pacstrap (outside of installation)
 			if mirror_config := archinstall.arguments.get('mirror_config', None):
-				if mirror_config.mirror_regions:
-					mirrors.use_mirrors(mirror_config.mirror_regions)
-				if mirror_config.custom_mirrors:
-					mirrors.add_custom_mirrors(mirror_config.custom_mirrors)
+				installation.set_mirrors(mirror_config)
 
 			installation.minimal_installation(
 				testing=enable_testing,
@@ -203,7 +198,7 @@ def perform_installation(mountpoint: Path, exec_mode: ExecutionMode):
 			)
 
 			if mirror_config := archinstall.arguments.get('mirror_config', None):
-				installation.set_mirrors(mirror_config)  # Set the mirrors in the installation medium
+				installation.set_mirrors(mirror_config, on_target=True)
 
 			if archinstall.arguments.get('swap'):
 				installation.setup_swap('zram')