Update import_haproxy_waf.py

import_haproxy_waf.py

@@ -2,134 +2,161 @@ import os
 import subprocess
 import logging
 from pathlib import Path
+import shutil
+import filecmp
+import time
 
-# Configure logging
+# --- Configuration ---
-logging.basicConfig(
+LOG_LEVEL = logging.INFO  # DEBUG, INFO, WARNING, ERROR
-    level=logging.INFO,
+WAF_DIR = Path(os.getenv("WAF_DIR", "waf_patterns/haproxy")).resolve()
-    format="%(asctime)s - %(levelname)s - %(message)s",
+HAPROXY_WAF_DIR = Path(os.getenv("HAPROXY_WAF_DIR", "/etc/haproxy/waf/")).resolve()
-    handlers=[logging.StreamHandler()],
+HAPROXY_CONF = Path(os.getenv("HAPROXY_CONF", "/etc/haproxy/haproxy.cfg")).resolve()
-)
+BACKUP_DIR = Path(os.getenv("BACKUP_DIR", "/etc/haproxy/waf_backup/")).resolve()
-
-# Constants (configurable via environment variables)
-WAF_DIR = Path(os.getenv("WAF_DIR", "waf_patterns/haproxy")).resolve()  # Source directory for WAF files
-HAPROXY_WAF_DIR = Path(os.getenv("HAPROXY_WAF_DIR", "/etc/haproxy/waf/")).resolve()  # Target directory
-HAPROXY_CONF = Path(os.getenv("HAPROXY_CONF", "/etc/haproxy/haproxy.cfg")).resolve()  # HAProxy config file
 
 # HAProxy WAF configuration snippet
 WAF_CONFIG_SNIPPET = """
-# WAF and Bot Protection
+# WAF and Bot Protection (Generated by import_haproxy_waf.py)
 frontend http-in
     bind *:80
-    default_backend web_backend
+    mode http
-    acl bad_bot hdr_sub(User-Agent) -i waf/bots.acl
+    option httplog
-    acl waf_attack path_reg waf/waf.acl
+    # WAF and Bot Protection ACLs and Rules
-    http-request deny if bad_bot
+    # Include generated ACL files
-    http-request deny if waf_attack
+    include /etc/haproxy/waf/*.acl
 """
+# --- Logging Setup ---
+logging.basicConfig(level=LOG_LEVEL, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
 
 
 def copy_waf_files():
-    """
+    """Copies WAF files, handling existing files, creating backups."""
-    Copy HAProxy WAF ACL files to the target directory.
+    logger.info("Copying HAProxy WAF patterns...")
 
-    Raises:
+    HAPROXY_WAF_DIR.mkdir(parents=True, exist_ok=True)
-        Exception: If there is an error copying files.
+    BACKUP_DIR.mkdir(parents=True, exist_ok=True)  # Ensure backup dir exists
-    """
-    logging.info("Copying HAProxy WAF patterns...")
 
-    try:
+    for acl_file in WAF_DIR.glob("*.acl"):  # Find all .acl files
-        # Ensure the target directory exists
+        dst_path = HAPROXY_WAF_DIR / acl_file.name
-        HAPROXY_WAF_DIR.mkdir(parents=True, exist_ok=True)
-        logging.info(f"[+] Created or verified directory: {HAPROXY_WAF_DIR}")
 
-        # Copy ACL files
+        try:
-        for file in ["bots.acl", "waf.acl"]:
+            if dst_path.exists():
-            src_path = WAF_DIR / file
+                # Compare and backup if different
-            dst_path = HAPROXY_WAF_DIR / file
+                if filecmp.cmp(acl_file, dst_path, shallow=False):
+                    logger.info(f"Skipping {acl_file.name} (identical file exists).")
+                    continue
+                # Different file exists: backup
+                backup_path = BACKUP_DIR / f"{dst_path.name}.{int(time.time())}"
+                logger.warning(f"Existing {dst_path.name} differs. Backing up to {backup_path}")
+                shutil.copy2(dst_path, backup_path)  # Backup old file
 
-            if not src_path.exists():
+            # Copy the (new or updated) file
-                logging.warning(f"[!] {file} not found in {WAF_DIR}")
+            shutil.copy2(acl_file, dst_path)
-                continue
+            logger.info(f"Copied {acl_file.name} to {dst_path}")
 
-            try:
+        except OSError as e:
-                subprocess.run(["cp", str(src_path), str(dst_path)], check=True)
+            logger.error(f"Error copying {acl_file.name}: {e}")
-                logging.info(f"[+] {file} copied to {HAPROXY_WAF_DIR}")
+            raise
-            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
 
 
 def update_haproxy_conf():
-    """
+    """Ensures the include statement is in haproxy.cfg, avoiding duplicates."""
-    Ensure the WAF configuration snippet is included in haproxy.cfg.
+    logger.info("Checking HAProxy configuration for WAF include...")
-
-    Raises:
-        Exception: If there is an error updating the HAProxy configuration.
-    """
-    logging.info("Ensuring WAF patterns are included in haproxy.cfg...")
 
     try:
-        # Read the current configuration
         with open(HAPROXY_CONF, "r") as f:
-            config = f.read()
+            config_lines = f.readlines()
 
-        # Append WAF configuration snippet if not present
+        # Check if the *exact* snippet is already present.  We'll check for the
-        if WAF_CONFIG_SNIPPET.strip() not in config:
+        # key parts of the snippet to be more robust.
-            logging.info("Adding WAF rules to haproxy.cfg...")
+        snippet_present = False
-            with open(HAPROXY_CONF, "a") as f:
+        for line in config_lines:
-                f.write(WAF_CONFIG_SNIPPET)
+            if "include /etc/haproxy/waf/*.acl" in line:
-            logging.info("[+] WAF rules added to haproxy.cfg.")
+                snippet_present = True
+                break
+
+        if not snippet_present:
+             # Find the 'frontend http-in' section
+            frontend_start = -1
+            for i, line in enumerate(config_lines):
+                if line.strip().startswith("frontend http-in"):
+                    frontend_start = i
+                    break
+
+            if frontend_start == -1:
+                logger.warning("No 'frontend http-in' section found. Appending to end of file.")
+                with open(HAPROXY_CONF, "a") as f:
+                    f.write(f"\n{WAF_CONFIG_SNIPPET}\n")
+                logger.info(f"Added WAF configuration snippet to {HAPROXY_CONF}")
+            else:
+                # Find the end of the 'frontend http-in' section
+                frontend_end = -1
+                for i in range(frontend_start + 1, len(config_lines)):
+                    if line.strip() == "" or  not line.startswith("    "): # Check it is part of the config
+                         frontend_end = i
+                         break
+
+
+                if frontend_end == -1:
+                    frontend_end = len(config_lines)  # End of file
+
+                # Insert the include statement *within* the frontend section.
+                config_lines.insert(frontend_end, "    include /etc/haproxy/waf/*.acl\n")
+
+                # Write the modified configuration back to the file
+                with open(HAPROXY_CONF, "w") as f:
+                    f.writelines(config_lines)
+                logger.info(f"Added WAF include to 'frontend http-in' section in {HAPROXY_CONF}")
         else:
-            logging.info("WAF patterns already included in haproxy.cfg.")
+            logger.info("WAF include statement already present.")
-    except Exception as e:
+
-        logging.error(f"[!] Error updating HAProxy configuration: {e}")
+    except FileNotFoundError:
+        logger.error(f"HAProxy configuration file not found: {HAPROXY_CONF}")
+        raise
+    except OSError as e:
+        logger.error(f"Error updating HAProxy configuration: {e}")
         raise
 
 
 def reload_haproxy():
-    """
+    """Tests the HAProxy configuration and reloads if valid."""
-    Reload HAProxy to apply the new WAF rules.
+    logger.info("Reloading HAProxy...")
-
-    Raises:
-        Exception: If there is an error reloading HAProxy.
-    """
-    logging.info("Testing HAProxy configuration...")
 
     try:
-        # Test HAProxy configuration
+        # Test configuration
-        subprocess.run(["haproxy", "-c", "-f", str(HAPROXY_CONF)], check=True)
+        result = subprocess.run(["haproxy", "-c", "-f", str(HAPROXY_CONF)],
-        logging.info("[+] HAProxy configuration test passed.")
+                                capture_output=True, text=True, check=True)
+        logger.info(f"HAProxy configuration test successful:\n{result.stdout}")
+
 
         # Reload HAProxy
-        subprocess.run(["systemctl", "reload", "haproxy"], check=True)
+        result = subprocess.run(["systemctl", "reload", "haproxy"],
-        logging.info("[+] HAProxy reloaded successfully.")
+                                capture_output=True, text=True, check=True)
+        logger.info("HAProxy reloaded.")
+
+
     except subprocess.CalledProcessError as e:
-        logging.error(f"[!] HAProxy configuration test failed: {e}")
+        logger.error(f"HAProxy command failed: {e.cmd} - Return code: {e.returncode}")
-        raise
+        logger.error(f"Stdout: {e.stdout}")
+        logger.error(f"Stderr: {e.stderr}")
+        raise  # Re-raise to signal failure
     except FileNotFoundError:
-        logging.error("[!] 'haproxy' or 'systemctl' command not found. Are you on a supported system?")
+        logger.error("'haproxy' or 'systemctl' command not found. Is HAProxy/systemd installed?")
-        raise
-    except Exception as e:
-        logging.error(f"[!] Error reloading HAProxy: {e}")
         raise
 
 
 def main():
-    """
+    """Main function."""
-    Main function to execute the script.
-    """
     try:
         copy_waf_files()
         update_haproxy_conf()
         reload_haproxy()
-        logging.info("[✔] HAProxy configured with latest WAF rules.")
+        logger.info("HAProxy WAF configuration updated successfully.")
     except Exception as e:
-        logging.critical(f"[!] Script failed: {e}")
+        logger.critical(f"Script failed: {e}")
         exit(1)
 
 
 if __name__ == "__main__":
-    main()
+    main()