Refactor IP validation and whitelist logic in IPBlocker

Refactor IPBlocker to improve IP validation and whitelist handling.

Author: mukeshdhadhariya

src/firewall/blocking.py

@@ -8,55 +8,50 @@
 from colorama import Fore, Style
 from utils.logger import get_logger
 from utils.system import get_platform_firewall_command
-import ipaddress
+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
-        # Normalize whitelist (support IPv4, IPv6, CIDR)
-        self.whitelist = {str(ipaddress.ip_network(ip, strict=False)) for ip in whitelist}
+        self.whitelist = whitelist
         self.blocked_ips: Dict[str, datetime] = {}
         self.lock = threading.Lock()
         self.logger = get_logger(__name__)
         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 network string (IPv4/IPv6/CIDR supported)"""
+        """
+        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}")
 
-    def _is_whitelisted(self, ip: str) -> bool:
-        """Check if IP is within whitelist networks"""
-        try:
-            ip_net = ipaddress.ip_network(ip, strict=False)
-            return any(ip_net.subnet_of(ipaddress.ip_network(wl)) for wl in self.whitelist)
-        except ValueError:
-            return False
-
     def block_ip(self, ip: str, reason: str) -> bool:
         """Block an IP address using the appropriate firewall system"""
         try:
-            ip = self._validate_ip(ip)
+            ip = self._validate_ip(ip)  # Validate IP first
         except ValueError:
             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}")
@@ -65,26 +60,26 @@ 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)
+            ip = self._validate_ip(ip)  # Validate IP before unblocking
         except ValueError:
             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}")
@@ -93,7 +88,7 @@ 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
@@ -106,19 +101,23 @@ def unblock_expired_ips(self) -> List[str]:
         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
             ]
-
+            
             for ip in expired_ips:
                 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:
@@ -134,7 +133,7 @@ def _execute_block_command(self, ip: str) -> bool:
         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:
@@ -150,58 +149,73 @@ def _execute_unblock_command(self, ip: str) -> bool:
         except Exception as e:
             self.logger.error(f"Platform-specific unblocking failed: {e}")
             return False
-
-    # -------- PLATFORM-SPECIFIC IMPLEMENTATIONS -------- #
-
+    
     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, capture_output=True, text=True)
+        result = subprocess.run(cmd, check=True, 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']
         result = subprocess.run(cmd, capture_output=True, text=True)
         return result.returncode == 0
-
+    
     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)
-        subprocess.run(['sudo', 'pfctl', '-e'], capture_output=True, text=True)
+        cmd2 = ['sudo', 'pfctl', '-e']
+        subprocess.run(cmd2, capture_output=True, text=True)
         return result1.returncode == 0
-
+    
     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
-
+    
     def _block_ip_windows(self, ip: str) -> bool:
         """Block IP using Windows Firewall (netsh)"""
-        rule_name = f"SimpleFirewall_Block_{ip.replace(':', '_').replace('.', '_')}"
+        rule_name = f"SimpleFirewall_Block_{ip.replace('.', '_')}"
         cmd = [
             'netsh', 'advfirewall', 'firewall', 'add', 'rule',
             f'name={rule_name}',
             'dir=in',
             'action=block',
             f'remoteip={ip}'
         ]
-        result = subprocess.run(cmd, capture_output=True, text=True)
-        return result.returncode == 0
+        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
 
     def _unblock_ip_windows(self, ip: str) -> bool:
         """Unblock IP using Windows Firewall (netsh)"""
-        rule_name = f"SimpleFirewall_Block_{ip.replace(':', '_').replace('.', '_')}"
+        rule_name = f"SimpleFirewall_Block_{ip.replace('.', '_')}"
         cmd = [
             'netsh', 'advfirewall', 'firewall', 'delete', 'rule',
             f'name={rule_name}'
         ]
-        result = subprocess.run(cmd, capture_output=True, text=True)
-        return result.returncode == 0
-
-    # --------------------------------------------------- #
+        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
 
     def get_blocked_ips(self) -> Dict[str, str]:
         """Get currently blocked IPs with their block times"""
@@ -210,16 +224,18 @@ def get_blocked_ips(self) -> Dict[str, str]:
                 ip: block_time.isoformat()
                 for ip, block_time 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
-
+    
     def get_stats(self) -> Dict[str, int]:
         """Get blocking statistics"""
         with self.lock: