From c629ec14e7c6a9ac56f9410cb3aa2ba3951f37c4 Mon Sep 17 00:00:00 2001
From: chanakya <chanakyavasantha@gmail.com>
Date: Tue, 19 Dec 2023 19:23:13 +0530
Subject: [PATCH] update

---
 harden/firewall.py | 50 +++++++++++++++++++++++++++++-----------------
 1 file changed, 32 insertions(+), 18 deletions(-)

diff --git a/harden/firewall.py b/harden/firewall.py
index d312e03..2f95761 100644
--- a/harden/firewall.py
+++ b/harden/firewall.py
@@ -2,26 +2,40 @@
 
 def get_script(config):
     firewall_config = config["firewall"]
-    # Start with an empty script and build it up
-    script = ""
+    script = "#!/bin/bash\n\n"  # Start with a bash shebang and a newline
 
-    if firewall_config['enable']:
-        # Each file system gets its own set of commands
-        script += f"apt install ufw\n"
-        script += f"apt purge iptables-persistent\n"
-        script += f"systemctl --now enable ufw.service\n"
-        if firewall_config['configure_loopback_traffic']:
-            script += f"ufw allow in on lo\n"
-            script += "ufw allow out on lo\n"
-            script += "ufw deny in from 127.0.0.0/8\n"
-            script += "ufw deny in from ::1\n"
-        elif firewall_config['enable_default_deny']:
-            script += "ufw default deny incoming\n"
-            script += "ufw default deny outgoing\n"
-            script += "ufw default deny routed\n"
+    if firewall_config.get('enable', False):
+        script += "sudo apt install ufw -y\n"  # Install UFW
+        script += "sudo apt purge iptables-persistent -y\n"  # Purge iptables-persistent
+        script += "sudo systemctl enable --now ufw.service\n"  # Enable and start UFW
+
+        if firewall_config.get('configure_loopback_traffic', False):
+            # Configure loopback traffic rules
+            script += "sudo ufw allow in on lo\n"
+            script += "sudo ufw allow out on lo\n"
+            script += "sudo ufw deny in from 127.0.0.0/8\n"
+            script += "sudo ufw deny in from ::1\n"
+        
+        if firewall_config.get('enable_default_deny', False):
+            # Set default deny policies
+            script += "sudo ufw default deny incoming\n"
+            script += "sudo ufw default deny outgoing\n"
+            script += "sudo ufw default deny routed\n"
+
+        # Enable UFW with the applied rules
+        script += "sudo ufw --force enable\n"
+        script += "sudo ufw status verbose\n"  # Display UFW status
 
     return script
 
 if __name__ == "__main__":
-    config = config_file.init()
-    print(get_script(config))
\ No newline at end of file
+    # Example configuration
+    config = {
+        "firewall": {
+            "enable": True,
+            "configure_loopback_traffic": True,
+            "enable_default_deny": False
+        }
+    }
+    generated_script = get_script(config)
+    print(generated_script)