Update import_haproxy_waf.py

import_haproxy_waf.py

@@ -1,63 +1,126 @@
 import os
 import subprocess
 import logging
+from pathlib import Path
 
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s",
+    handlers=[logging.StreamHandler()],
+)
 
-# Paths
-WAF_DIR = "waf_patterns/haproxy"
-HAPROXY_WAF_DIR = "/etc/haproxy/waf/"
-HAPROXY_CONF = "/etc/haproxy/haproxy.cfg"
+# Constants (configurable via environment variables)
+WAF_DIR = Path(os.getenv("WAF_DIR", "waf_patterns/haproxy"))  # Source directory for WAF files
+HAPROXY_WAF_DIR = Path(os.getenv("HAPROXY_WAF_DIR", "/etc/haproxy/waf/"))  # Target directory
+HAPROXY_CONF = Path(os.getenv("HAPROXY_CONF", "/etc/haproxy/haproxy.cfg"))  # HAProxy config file
+
+# HAProxy WAF configuration snippet
+WAF_CONFIG_SNIPPET = """
+# WAF and Bot Protection
+frontend http-in
+    bind *:80
+    default_backend web_backend
+    acl bad_bot hdr_sub(User-Agent) -i waf/bots.acl
+    acl waf_attack path_reg waf/waf.acl
+    http-request deny if bad_bot
+    http-request deny if waf_attack
+"""
 
-INCLUDE_STATEMENT = "    acl bad_bot hdr_sub(User-Agent) -i waf/bots.acl"
 
 def copy_waf_files():
+    """
+    Copy HAProxy WAF ACL files to the target directory.
+    """
     logging.info("Copying HAProxy WAF patterns...")
 
-    # Ensure the target directory exists
-    os.makedirs(HAPROXY_WAF_DIR, exist_ok=True)
+    try:
+        # Ensure the target directory exists
+        HAPROXY_WAF_DIR.mkdir(parents=True, exist_ok=True)
+        logging.info(f"[+] Created or verified directory: {HAPROXY_WAF_DIR}")
+
+        # Copy ACL files
+        for file in ["bots.acl", "waf.acl"]:
+            src_path = WAF_DIR / file
+            dst_path = HAPROXY_WAF_DIR / file
+
+            if not src_path.exists():
+                logging.warning(f"[!] {file} not found in {WAF_DIR}")
+                continue
+
+            try:
+                subprocess.run(["cp", str(src_path), str(dst_path)], check=True)
+                logging.info(f"[+] {file} copied to {HAPROXY_WAF_DIR}")
+            except subprocess.CalledProcessError as e:
+                logging.error(f"[!] Failed to copy {file}: {e}")
+                raise
+    except Exception as e:
+        logging.error(f"[!] Error copying WAF files: {e}")
+        raise
 
-    # Copy acl files
-    for file in ["bots.acl", "waf.acl"]:
-        src_path = os.path.join(WAF_DIR, file)
-        dst_path = os.path.join(HAPROXY_WAF_DIR, file)
-        
-        if os.path.exists(src_path):
-            subprocess.run(["cp", src_path, dst_path], check=True)
-            logging.info(f"[+] {file} copied to {HAPROXY_WAF_DIR}")
-        else:
-            logging.warning(f"[!] {file} not found in {WAF_DIR}")
 
 def update_haproxy_conf():
+    """
+    Ensure the WAF configuration snippet is included in haproxy.cfg.
+    """
     logging.info("Ensuring WAF patterns are included in haproxy.cfg...")
 
-    with open(HAPROXY_CONF, "r") as f:
-        config = f.read()
+    try:
+        # Read the current configuration
+        with open(HAPROXY_CONF, "r") as f:
+            config = f.read()
+
+        # Append WAF configuration snippet if not present
+        if WAF_CONFIG_SNIPPET.strip() not in config:
+            logging.info("Adding WAF rules to haproxy.cfg...")
+            with open(HAPROXY_CONF, "a") as f:
+                f.write(WAF_CONFIG_SNIPPET)
+            logging.info("[+] WAF rules added to haproxy.cfg.")
+        else:
+            logging.info("WAF patterns already included in haproxy.cfg.")
+    except Exception as e:
+        logging.error(f"[!] Error updating HAProxy configuration: {e}")
+        raise
 
-    if INCLUDE_STATEMENT not in config:
-        logging.info("Adding WAF rules to haproxy.cfg...")
-        with open(HAPROXY_CONF, "a") as f:
-            f.write(
-                f"\n# WAF and Bot Protection\n"
-                f"frontend http-in\n"
-                f"    bind *:80\n"
-                f"    default_backend web_backend\n"
-                f"    acl bad_bot hdr_sub(User-Agent) -i waf/bots.acl\n"
-                f"    acl waf_attack path_reg waf/waf.acl\n"
-                f"    http-request deny if bad_bot\n"
-                f"    http-request deny if waf_attack\n"
-            )
-    else:
-        logging.info("WAF patterns already included in haproxy.cfg.")
 
 def reload_haproxy():
+    """
+    Reload HAProxy to apply the new WAF rules.
+    """
     logging.info("Testing HAProxy configuration...")
-    subprocess.run(["haproxy", "-c", "-f", HAPROXY_CONF], check=True)
-    logging.info("Reloading HAProxy to apply new WAF rules...")
-    subprocess.run(["systemctl", "reload", "haproxy"], check=True)
+
+    try:
+        # Test HAProxy configuration
+        subprocess.run(["haproxy", "-c", "-f", str(HAPROXY_CONF)], check=True)
+        logging.info("[+] HAProxy configuration test passed.")
+
+        # Reload HAProxy
+        subprocess.run(["systemctl", "reload", "haproxy"], check=True)
+        logging.info("[+] HAProxy reloaded successfully.")
+    except subprocess.CalledProcessError as e:
+        logging.error(f"[!] HAProxy configuration test failed: {e}")
+        raise
+    except FileNotFoundError:
+        logging.error("[!] 'haproxy' or 'systemctl' command not found. Are you on a supported system?")
+        raise
+    except Exception as e:
+        logging.error(f"[!] Error reloading HAProxy: {e}")
+        raise
+
+
+def main():
+    """
+    Main function to execute the script.
+    """
+    try:
+        copy_waf_files()
+        update_haproxy_conf()
+        reload_haproxy()
+        logging.info("[✔] HAProxy configured with latest WAF rules.")
+    except Exception as e:
+        logging.critical(f"[!] Script failed: {e}")
+        exit(1)
+
 
 if __name__ == "__main__":
-    copy_waf_files()
-    update_haproxy_conf()
-    reload_haproxy()
-    logging.info("[✔] HAProxy configured with latest WAF rules.")
+    main()