Add UFW firewall configuration to ConfigureCommand with SSH lockout prevention

src/application/command_handlers/configure/handler.rs

@@ -8,7 +8,8 @@ use super::errors::ConfigureCommandHandlerError;
 use crate::adapters::ansible::AnsibleClient;
 use crate::application::command_handlers::common::StepResult;
 use crate::application::steps::{
-    ConfigureSecurityUpdatesStep, InstallDockerComposeStep, InstallDockerStep,
+    ConfigureFirewallStep, ConfigureSecurityUpdatesStep, InstallDockerComposeStep,
+    InstallDockerStep,
 };
 use crate::domain::environment::repository::{EnvironmentRepository, TypedEnvironmentRepository};
 use crate::domain::environment::state::{ConfigureFailureContext, ConfigureStep};
@@ -24,6 +25,7 @@ use crate::shared::error::Traceable;
 /// 1. Install Docker
 /// 2. Install Docker Compose
 /// 3. Configure automatic security updates
+/// 4. Configure UFW firewall
 ///
 /// # State Management
 ///
@@ -161,6 +163,27 @@ impl ConfigureCommandHandler {
             .execute()
             .map_err(|e| (e.into(), current_step))?;
 
+        let current_step = ConfigureStep::ConfigureFirewall;
+        // Allow tests or CI to explicitly skip the firewall configuration step
+        // (useful for container-based test runs where iptables/ufw require
+        // elevated kernel capabilities not available in unprivileged containers).
+        let skip_firewall = std::env::var("TORRUST_TD_SKIP_FIREWALL_IN_CONTAINER")
+            .map(|v| v == "true")
+            .unwrap_or(false);
+
+        if skip_firewall {
+            info!(
+                command = "configure",
+                step = "configure_firewall",
+                status = "skipped",
+                "Skipping UFW firewall configuration due to TORRUST_TD_SKIP_FIREWALL_IN_CONTAINER"
+            );
+        } else {
+            ConfigureFirewallStep::new(Arc::clone(&self.ansible_client))
+                .execute()
+                .map_err(|e| (e.into(), current_step))?;
+        }
+
         // Transition to Configured state
         let configured = environment.clone().configured();
 

src/application/steps/mod.rs

@@ -36,7 +36,7 @@ pub use rendering::{
     RenderAnsibleTemplatesError, RenderAnsibleTemplatesStep, RenderOpenTofuTemplatesStep,
 };
 pub use software::{InstallDockerComposeStep, InstallDockerStep};
-pub use system::{ConfigureSecurityUpdatesStep, WaitForCloudInitStep};
+pub use system::{ConfigureFirewallStep, ConfigureSecurityUpdatesStep, WaitForCloudInitStep};
 pub use validation::{
     ValidateCloudInitCompletionStep, ValidateDockerComposeInstallationStep,
     ValidateDockerInstallationStep,

src/application/steps/system/configure_firewall.rs

@@ -0,0 +1,139 @@
+//! UFW firewall configuration step
+//!
+//! This module provides the `ConfigureFirewallStep` which handles configuration
+//! of UFW (Uncomplicated Firewall) on remote hosts via Ansible playbooks.
+//! This step ensures that the firewall is configured with restrictive default
+//! policies while maintaining SSH access to prevent lockout.
+//!
+//! ## Key Features
+//!
+//! - Configures UFW with restrictive default policies (deny incoming, allow outgoing)
+//! - Preserves SSH access on the configured port
+//! - Uses Tera template for dynamic SSH port resolution
+//! - Comprehensive SSH lockout prevention measures
+//! - Verification steps to ensure firewall is active and SSH is accessible
+//!
+//! ## Configuration Process
+//!
+//! The step executes the "configure-firewall" Ansible playbook which handles:
+//! - UFW installation and setup
+//! - Reset UFW to clean state
+//! - Set restrictive default policies
+//! - Allow SSH access BEFORE enabling firewall (critical for preventing lockout)
+//! - Enable UFW firewall
+//! - Verify firewall status and SSH access
+//!
+//! ## SSH Lockout Prevention
+//!
+//! This is a **high-risk operation** that could result in SSH lockout if not
+//! handled correctly. Safety measures include:
+//!
+//! 1. **Correct Sequencing**: SSH rules are added BEFORE enabling firewall
+//! 2. **Dual SSH Protection**: Both port-specific and service-name rules
+//! 3. **Port Configuration**: Uses actual SSH port from user configuration
+//! 4. **Verification Steps**: Ansible tasks verify SSH access is preserved
+//! 5. **Comprehensive Logging**: Detailed logging of each firewall step
+
+use std::sync::Arc;
+use tracing::{info, instrument, warn};
+
+use crate::adapters::ansible::AnsibleClient;
+use crate::shared::command::CommandError;
+
+/// Step that configures UFW firewall on a remote host via Ansible
+///
+/// This step configures a restrictive UFW firewall policy while ensuring
+/// SSH access is maintained. The SSH port is resolved during template rendering
+/// and embedded in the final Ansible playbook. The configuration follows the
+/// principle of "allow SSH BEFORE enabling firewall" to prevent lockout.
+pub struct ConfigureFirewallStep {
+    ansible_client: Arc<AnsibleClient>,
+}
+
+impl ConfigureFirewallStep {
+    /// Create a new firewall configuration step
+    ///
+    /// # Arguments
+    ///
+    /// * `ansible_client` - Ansible client for running playbooks
+    ///
+    /// # Note
+    ///
+    /// SSH port configuration is resolved during template rendering phase,
+    /// not at step execution time. The rendered playbook contains the
+    /// resolved SSH port value.
+    #[must_use]
+    pub fn new(ansible_client: Arc<AnsibleClient>) -> Self {
+        Self { ansible_client }
+    }
+
+    /// Execute the firewall configuration
+    ///
+    /// # Safety
+    ///
+    /// This method is designed to prevent SSH lockout by:
+    /// 1. Resetting UFW to clean state
+    /// 2. Allowing SSH access BEFORE enabling firewall
+    /// 3. Using the correct SSH port from user configuration
+    ///
+    /// The SSH port is resolved during template rendering and embedded in the
+    /// playbook, so this method executes a playbook with pre-configured values.
+    ///
+    /// # Errors
+    ///
+    /// Returns `CommandError` if:
+    /// - Ansible playbook execution fails
+    /// - UFW commands fail
+    /// - SSH rules cannot be applied
+    /// - Firewall verification fails
+    #[instrument(
+        name = "configure_firewall",
+        skip_all,
+        fields(step_type = "system", component = "firewall", method = "ansible")
+    )]
+    pub fn execute(&self) -> Result<(), CommandError> {
+        warn!(
+            step = "configure_firewall",
+            action = "configure_ufw",
+            "Configuring UFW firewall - CRITICAL: SSH access will be restricted to configured port"
+        );
+
+        // Run Ansible playbook (SSH port already resolved during template rendering)
+        match self.ansible_client.run_playbook("configure-firewall") {
+            Ok(_) => {
+                info!(
+                    step = "configure_firewall",
+                    status = "success",
+                    "UFW firewall configured successfully with SSH access preserved"
+                );
+                Ok(())
+            }
+            Err(e) => {
+                // Propagate errors to the caller. Tests that run in container environments
+                // should explicitly opt-out of running this step (for example via an
+                // environment variable) instead of relying on runtime error detection.
+                Err(e)
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::path::PathBuf;
+    use std::sync::Arc;
+
+    use super::*;
+
+    #[test]
+    fn it_should_create_configure_firewall_step() {
+        let ansible_client = Arc::new(AnsibleClient::new(PathBuf::from("test_inventory.yml")));
+        let step = ConfigureFirewallStep::new(ansible_client);
+
+        // Test that the step can be created successfully
+        assert_eq!(
+            std::mem::size_of_val(&step),
+            std::mem::size_of::<Arc<AnsibleClient>>()
+        );
+    }
+}

src/application/steps/system/mod.rs

@@ -7,16 +7,18 @@
  * Current steps:
  * - Cloud-init completion waiting
  * - Automatic security updates configuration
+ * - UFW firewall configuration
  *
  * Future steps may include:
  * - User account setup and management
- * - Firewall configuration
  * - Log rotation configuration
  * - System service management
  */
 
+pub mod configure_firewall;
 pub mod configure_security_updates;
 pub mod wait_cloud_init;
 
+pub use configure_firewall::ConfigureFirewallStep;
 pub use configure_security_updates::ConfigureSecurityUpdatesStep;
 pub use wait_cloud_init::WaitForCloudInitStep;

src/bin/e2e_config_tests.rs

@@ -112,6 +112,10 @@ struct CliArgs {
 pub async fn main() -> Result<()> {
     let cli = CliArgs::parse();
 
+    // Set environment variable to skip firewall configuration in container-based tests
+    // UFW/iptables requires kernel capabilities not available in unprivileged containers
+    std::env::set_var("TORRUST_TD_SKIP_FIREWALL_IN_CONTAINER", "true");
+
     // Initialize logging with production log location for E2E tests using the builder pattern
     LoggingBuilder::new(std::path::Path::new("./data/logs"))
         .with_format(cli.log_format.clone())

src/domain/environment/state/configure_failed.rs

@@ -47,6 +47,8 @@ pub enum ConfigureStep {
     InstallDockerCompose,
     /// Configuring automatic security updates
     ConfigureSecurityUpdates,
+    /// Configuring UFW firewall
+    ConfigureFirewall,
 }
 
 /// Error state - Application configuration failed