From ff321b63d1d3ef05318172196aea3684240a1612 Mon Sep 17 00:00:00 2001
From: CooperWang0912 <cooperwang0912@gmail.com>
Date: Tue, 28 Jul 2026 14:41:48 +0800
Subject: [PATCH] Refactored Command Execution

---
 archinstall/applications/firewall.py  | 21 +++++++++++++++++++++
 archinstall/lib/configuration.py      |  2 +-
 archinstall/lib/models/application.py |  9 ++++++---
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/archinstall/applications/firewall.py b/archinstall/applications/firewall.py
index dadaa05b..c7be6f2a 100644
--- a/archinstall/applications/firewall.py
+++ b/archinstall/applications/firewall.py
@@ -32,6 +32,24 @@ class FirewallApp:
 			'firewalld.service',
 		]
 
+	def _allow_ufw_ssh_on_first_boot(self, install_session: Installer) -> None:
+		service_content = """[Unit]
+			Description=Allow SSH in UFW on first boot
+			After=ufw.service
+			Wants=ufw.service
+
+			[Service]
+			Type=oneshot
+			ExecStart=/usr/bin/ufw allow SSH
+			ExecStartPost=/usr/bin/systemctl disable ufw-allow-ssh.service
+
+			[Install]
+			WantedBy=multi-user.target
+		"""
+		service_path = install_session.target / 'etc/systemd/system/ufw-allow-ssh.service'
+		service_path.write_text(service_content)
+		install_session.enable_service(['ufw-allow-ssh.service'])
+
 	def install(
 		self,
 		install_session: Installer,
@@ -47,6 +65,9 @@ class FirewallApp:
 				ufw_conf = install_session.target / 'etc/ufw/ufw.conf'
 				ufw_conf.write_text(ufw_conf.read_text().replace('ENABLED=no', 'ENABLED=yes'))
 
+				if firewall_config.allow_ssh:
+					self._allow_ufw_ssh_on_first_boot(install_session)
+
 			case Firewall.FWD:
 				install_session.add_additional_packages(self.fwd_packages)
 				install_session.enable_service(self.fwd_services)
diff --git a/archinstall/lib/configuration.py b/archinstall/lib/configuration.py
index 29dabcce..c4d85ed1 100644
--- a/archinstall/lib/configuration.py
+++ b/archinstall/lib/configuration.py
@@ -30,7 +30,7 @@ async def confirm_ufw(config: ArchConfig) -> bool:
 		).show()
 
 		if result and result.get_value():
-			config.custom_commands.append('ufw allow SSH')
+			config.app_config.firewall_config.allow_ssh = True
 
 	return True
 
diff --git a/archinstall/lib/models/application.py b/archinstall/lib/models/application.py
index 1e97a784..396a10a9 100644
--- a/archinstall/lib/models/application.py
+++ b/archinstall/lib/models/application.py
@@ -40,6 +40,7 @@ class Firewall(StrEnum):
 
 class FirewallConfigSerialization(TypedDict):
 	firewall: str
+	allow_ssh: NotRequired[bool]
 
 
 class FontPackage(StrEnum):
@@ -143,16 +144,16 @@ class PrintServiceConfiguration:
 @dataclass
 class FirewallConfiguration:
 	firewall: Firewall
+	allow_ssh: bool = False
 
 	def json(self) -> FirewallConfigSerialization:
-		return {
-			'firewall': self.firewall.value,
-		}
+		return {'firewall': self.firewall.value, 'allow_ssh': self.allow_ssh}
 
 	@classmethod
 	def parse_arg(cls, arg: dict[str, Any]) -> Self:
 		return cls(
 			Firewall(arg['firewall']),
+			allow_ssh=arg.get('allow_ssh', False),
 		)
 
 
@@ -285,6 +286,8 @@ class ApplicationConfiguration(SubConfig):
 
 		if self.firewall_config:
 			out.append(tr('Firewall "{}"').format(self.firewall_config.firewall))
+			if self.firewall_config.allow_ssh:
+				out.append({tr('SSH allowed')})
 
 		if self.fonts_config and self.fonts_config.fonts:
 			fonts = ', '.join(f.value for f in self.fonts_config.fonts)