feat : ip address validate

Author: mukeshdhadhariya

src/firewall/blocking.py

@@ -1,19 +1,22 @@
-"""IP blocking and unblocking functionality with cross-platform support"""
+"""IP blocking and unblocking functionality with IPv4 + IPv6 + CIDR + cross-platform support"""
 
 import subprocess
 import threading
 import platform
+import tempfile
+import shutil
+import ipaddress
 from datetime import datetime, timedelta
+from pathlib import Path
 from typing import Dict, Set, List
 from colorama import Fore, Style
 from utils.logger import get_logger
 from utils.system import get_platform_firewall_command
-import ipaddress 
 
 
 class IPBlocker:
     """Handles IP blocking operations across different platforms"""
-    
+
     def __init__(self, block_duration: int, whitelist: Set[str]):
         self.block_duration = block_duration
         self.whitelist = whitelist
@@ -23,35 +26,35 @@ def __init__(self, block_duration: int, whitelist: Set[str]):
         self.platform = platform.system().lower()
         self.firewall_cmd = get_platform_firewall_command()
 
-    # Added helper for IP validation
-    def _validate_ip(self, ip: str) -> str:
-        """
-        Validate and normalize an IP or CIDR (IPv4 or IPv6).
-        Raises ValueError if invalid.
-        """
-        try:
-            network = ipaddress.ip_network(ip, strict=False)
-            return str(network)
-        except ValueError as e:
-            self.logger.error(f"Invalid IP/network '{ip}': {e}")
-            raise ValueError(f"Invalid IP/network: {ip}")
+        # macOS setup for pfctl
+        if self.platform == "darwin":
+            self._init_macos_firewall()
 
-    def block_ip(self, ip: str, reason: str) -> bool:
-        """Block an IP address using the appropriate firewall system"""
+    def _is_valid_ip(self, ip: str) -> bool:
+        """Validate IPv4, IPv6, and CIDR ranges."""
         try:
-            ip = self._validate_ip(ip)  # Validate IP first
+            ipaddress.ip_network(ip, strict=False)
+            return True
         except ValueError:
             return False
 
+    def _is_whitelisted(self, ip: str) -> bool:
+        return ip in self.whitelist
+
+    def block_ip(self, ip: str, reason: str) -> bool:
+        """Block an IP or CIDR range."""
+        if not self._is_valid_ip(ip):
+            self.logger.warning(f"Invalid IP/CIDR rejected: {ip}")
+            return False
+
         if self._is_whitelisted(ip):
             self.logger.info(f"IP {ip} is whitelisted, not blocking")
             return False
-            
+
         with self.lock:
             if ip not in self.blocked_ips:
                 try:
                     success = self._execute_block_command(ip)
-                    
                     if success:
                         self.blocked_ips[ip] = datetime.now()
                         self.logger.warning(f"🚫 BLOCKED IP: {ip} - Reason: {reason}")
@@ -60,26 +63,23 @@ def block_ip(self, ip: str, reason: str) -> bool:
                     else:
                         self.logger.error(f"Failed to block IP {ip}")
                         return False
-                        
                 except Exception as e:
                     self.logger.error(f"Exception while blocking IP {ip}: {e}")
                     return False
             else:
                 self.logger.debug(f"IP {ip} already blocked")
                 return True
-    
+
     def unblock_ip(self, ip: str) -> bool:
-        """Manually unblock a specific IP address"""
-        try:
-            ip = self._validate_ip(ip)  # Validate IP before unblocking
-        except ValueError:
+        """Unblock IPv4/IPv6 address or CIDR."""
+        if not self._is_valid_ip(ip):
+            self.logger.warning(f"Invalid IP/CIDR for unblock: {ip}")
             return False
 
         with self.lock:
             if ip in self.blocked_ips:
                 try:
                     success = self._execute_unblock_command(ip)
-                    
                     if success:
                         del self.blocked_ips[ip]
                         self.logger.info(f"✅ UNBLOCKED IP: {ip}")
@@ -88,7 +88,6 @@ def unblock_ip(self, ip: str) -> bool:
                     else:
                         self.logger.error(f"Failed to unblock IP {ip}")
                         return False
-                        
                 except Exception as e:
                     self.logger.error(f"Exception while unblocking IP {ip}: {e}")
                     return False
@@ -97,149 +96,160 @@ def unblock_ip(self, ip: str) -> bool:
                 return False
 
     def unblock_expired_ips(self) -> List[str]:
-        """Unblock IPs that have exceeded the block duration"""
         current_time = datetime.now()
         block_duration = timedelta(seconds=self.block_duration)
         unblocked_ips = []
-        
+
         with self.lock:
-            expired_ips = [
-                ip for ip, block_time in self.blocked_ips.items()
-                if current_time - block_time > block_duration
+            expired = [
+                ip for ip, t in self.blocked_ips.items()
+                if current_time - t > block_duration
             ]
-            
-            for ip in expired_ips:
+
+            for ip in expired:
                 if self.unblock_ip(ip):
                     unblocked_ips.append(ip)
-        
+
         return unblocked_ips
-    
-    def _is_whitelisted(self, ip: str) -> bool:
-        """Check if IP is in whitelist"""
-        return ip in self.whitelist
-    
+ 
+
     def _execute_block_command(self, ip: str) -> bool:
-        """Execute platform-specific block command"""
         try:
-            if self.platform == 'linux':
+            if self.platform == "linux":
                 return self._block_ip_linux(ip)
-            elif self.platform == 'darwin':  # macOS
+            elif self.platform == "darwin":
                 return self._block_ip_macos(ip)
-            elif self.platform == 'windows':
+            elif self.platform == "windows":
                 return self._block_ip_windows(ip)
             else:
-                self.logger.error(f"Blocking not implemented for platform: {self.platform}")
+                self.logger.error(f"Blocking not implemented for: {self.platform}")
                 return False
         except Exception as e:
             self.logger.error(f"Platform-specific blocking failed: {e}")
             return False
-    
+
     def _execute_unblock_command(self, ip: str) -> bool:
-        """Execute platform-specific unblock command"""
         try:
-            if self.platform == 'linux':
+            if self.platform == "linux":
                 return self._unblock_ip_linux(ip)
-            elif self.platform == 'darwin':  # macOS
+            elif self.platform == "darwin":
                 return self._unblock_ip_macos(ip)
-            elif self.platform == 'windows':
+            elif self.platform == "windows":
                 return self._unblock_ip_windows(ip)
             else:
-                self.logger.error(f"Unblocking not implemented for platform: {self.platform}")
+                self.logger.error(f"Unblocking not implemented for: {self.platform}")
                 return False
         except Exception as e:
             self.logger.error(f"Platform-specific unblocking failed: {e}")
             return False
-    
+
+
     def _block_ip_linux(self, ip: str) -> bool:
-        """Block IP using iptables on Linux"""
-        cmd = ['sudo', 'iptables', '-A', 'INPUT', '-s', ip, '-j', 'DROP']
-        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
+        network = ipaddress.ip_network(ip, strict=False)
+
+        if network.version == 4:
+            cmd = ["sudo", "iptables", "-A", "INPUT", "-s", ip, "-j", "DROP"]
+        else:
+            cmd = ["sudo", "ip6tables", "-A", "INPUT", "-s", ip, "-j", "DROP"]
+
+        result = subprocess.run(cmd, capture_output=True, text=True)
         return result.returncode == 0
-    
+
     def _unblock_ip_linux(self, ip: str) -> bool:
-        """Unblock IP using iptables on Linux"""
-        cmd = ['sudo', 'iptables', '-D', 'INPUT', '-s', ip, '-j', 'DROP']
+        network = ipaddress.ip_network(ip, strict=False)
+
+        if network.version == 4:
+            cmd = ["sudo", "iptables", "-D", "INPUT", "-s", ip, "-j", "DROP"]
+        else:
+            cmd = ["sudo", "ip6tables", "-D", "INPUT", "-s", ip, "-j", "DROP"]
+
         result = subprocess.run(cmd, capture_output=True, text=True)
         return result.returncode == 0
-    
+
+
+    def _init_macos_firewall(self):
+        """Initialize pfctl and create rules."""
+        try:
+            pfctl_path = shutil.which("pfctl")
+            if not pfctl_path:
+                self.logger.warning("pfctl not found on macOS.")
+                return
+
+            # Create/verify table
+            subprocess.run(["sudo", "pfctl", "-t", "blocked_ips", "-T", "show"],
+                           capture_output=True)
+
+            self._reload_macos_rules()
+            self.logger.info("macOS firewall initialized.")
+        except Exception as e:
+            self.logger.error(f"Failed to init macOS firewall: {e}")
+
+    def _reload_macos_rules(self):
+        """Write pfctl rule file and reload firewall"""
+        try:
+            pf_conf = """
+table <blocked_ips> persist
+block drop in quick from <blocked_ips> to any
+block drop out quick from any to <blocked_ips>
+pass in all
+pass out all
+"""
+            temp_path = Path(tempfile.gettempdir()) / "simple_fw.conf"
+            temp_path.write_text(pf_conf)
+
+            subprocess.run(["sudo", "pfctl", "-f", str(temp_path)],
+                           capture_output=True)
+        except Exception as e:
+            self.logger.error(f"Failed to reload pfctl rules: {e}")
+
     def _block_ip_macos(self, ip: str) -> bool:
-        """Block IP using pfctl on macOS"""
-        cmd1 = ['sudo', 'pfctl', '-t', 'blocked_ips', '-T', 'add', ip]
-        result1 = subprocess.run(cmd1, capture_output=True, text=True)
-        cmd2 = ['sudo', 'pfctl', '-e']
-        subprocess.run(cmd2, capture_output=True, text=True)
-        return result1.returncode == 0
-    
+        cmd = ["sudo", "pfctl", "-t", "blocked_ips", "-T", "add", ip]
+        res = subprocess.run(cmd, capture_output=True, text=True)
+        return res.returncode == 0 or "already" in res.stderr.lower()
+
     def _unblock_ip_macos(self, ip: str) -> bool:
-        """Unblock IP using pfctl on macOS"""
-        cmd = ['sudo', 'pfctl', '-t', 'blocked_ips', '-T', 'delete', ip]
-        result = subprocess.run(cmd, capture_output=True, text=True)
-        return result.returncode == 0
-    
+        cmd = ["sudo", "pfctl", "-t", "blocked_ips", "-T", "delete", ip]
+        res = subprocess.run(cmd, capture_output=True, text=True)
+        return res.returncode == 0 or "not found" in res.stderr.lower()
+
+
     def _block_ip_windows(self, ip: str) -> bool:
-        """Block IP using Windows Firewall (netsh)"""
-        rule_name = f"SimpleFirewall_Block_{ip.replace('.', '_')}"
+        rule_name = f"SimpleFirewall_Block_{ip.replace(':', '_').replace('.', '_')}"
         cmd = [
-            'netsh', 'advfirewall', 'firewall', 'add', 'rule',
-            f'name={rule_name}',
-            'dir=in',
-            'action=block',
-            f'remoteip={ip}'
+            "netsh", "advfirewall", "firewall", "add", "rule",
+            f"name={rule_name}",
+            "dir=in", "action=block", f"remoteip={ip}"
         ]
-        try:
-            result = subprocess.run(cmd, capture_output=True, text=True)
-            if result.returncode == 0:
-                self.logger.debug(f"netsh add rule stdout: {result.stdout.strip()}")
-                return True
-            else:
-                self.logger.error(f"netsh add rule failed: rc={result.returncode} stdout={result.stdout.strip()} stderr={result.stderr.strip()}")
-                return False
-        except Exception as e:
-            self.logger.error(f"Exception when running netsh add rule: {e}")
-            return False
+        res = subprocess.run(cmd, capture_output=True, text=True)
+        return res.returncode == 0
 
     def _unblock_ip_windows(self, ip: str) -> bool:
-        """Unblock IP using Windows Firewall (netsh)"""
-        rule_name = f"SimpleFirewall_Block_{ip.replace('.', '_')}"
-        cmd = [
-            'netsh', 'advfirewall', 'firewall', 'delete', 'rule',
-            f'name={rule_name}'
-        ]
-        try:
-            result = subprocess.run(cmd, capture_output=True, text=True)
-            if result.returncode == 0:
-                self.logger.debug(f"netsh delete rule stdout: {result.stdout.strip()}")
-                return True
-            else:
-                self.logger.error(f"netsh delete rule failed: rc={result.returncode} stdout={result.stdout.strip()} stderr={result.stderr.strip()}")
-                return False
-        except Exception as e:
-            self.logger.error(f"Exception when running netsh delete rule: {e}")
-            return False
+        rule_name = f"SimpleFirewall_Block_{ip.replace(':', '_').replace('.', '_')}"
+        cmd = ["netsh", "advfirewall", "firewall", "delete", "rule",
+               f"name={rule_name}"]
+        res = subprocess.run(cmd, capture_output=True, text=True)
+        return res.returncode == 0
+
 
     def get_blocked_ips(self) -> Dict[str, str]:
-        """Get currently blocked IPs with their block times"""
         with self.lock:
-            return {
-                ip: block_time.isoformat()
-                for ip, block_time in self.blocked_ips.items()
-            }
-    
+            return {ip: t.isoformat() for ip, t in self.blocked_ips.items()}
+
     def cleanup_all_blocks(self) -> List[str]:
-        """Remove all blocks (useful for shutdown)"""
-        cleaned_ips = []
-        
-        with self.lock:
-            for ip in list(self.blocked_ips.keys()):
-                if self.unblock_ip(ip):
-                    cleaned_ips.append(ip)
-        
-        return cleaned_ips
-    
+        cleaned = []
+        for ip in list(self.blocked_ips.keys()):
+            if self.unblock_ip(ip):
+                cleaned.append(ip)
+
+        # macOS table flush
+        if self.platform == "darwin" and cleaned:
+            subprocess.run(["sudo", "pfctl", "-t", "blocked_ips", "-T", "flush"])
+
+        return cleaned
+
     def get_stats(self) -> Dict[str, int]:
-        """Get blocking statistics"""
         with self.lock:
             return {
-                'currently_blocked': len(self.blocked_ips),
-                'whitelist_size': len(self.whitelist)
-            }
+                "currently_blocked": len(self.blocked_ips),
+                "whitelist_size": len(self.whitelist)
+            }