Improved IPv6 Handling for Nmap Agent

agent/nmap_agent.py

@@ -45,6 +45,10 @@
 DEFAULT_MASK_IPV6 = 128
 # scan up to 65536 host
 IPV6_CIDR_LIMIT = 112
+# More IPv6-specific constants
+IPV6_MIN_PREFIX = 8  # Minimum safe prefix length for IPv6
+IPV6_DEFAULT_PING_TIMEOUT = 1000  # ms
+IPV6_MAX_BATCH_SIZE = 100  # Maximum number of IPv6 addresses to scan in one batch
 
 BLACKLISTED_SERVICES = ["tcpwrapped"]
 
@@ -74,6 +78,33 @@
         self._scope_domain_regex: Optional[str] = self.args.get("scope_domain_regex")
         self._vpn_config: Optional[str] = self.args.get("vpn_config")
         self._dns_config: Optional[str] = self.args.get("dns_config")
+        self._validate_ipv6_settings()
+
+    def _validate_ipv6_settings(self) -> None:
+        """Validate IPv6-specific settings."""
+        max_mask = int(self.args.get("max_network_mask_ipv6", "128"))
+        if max_mask < IPV6_MIN_PREFIX or max_mask > 128:
+            raise ValueError(f"IPv6 prefix must be between {IPV6_MIN_PREFIX} and 128")
+
+    def _normalize_ipv6_address(self, address: str) -> str:
+        """Normalize IPv6 address format."""
+        try:
+            return str(ipaddress.IPv6Address(address).compressed)
+        except ipaddress.AddressValueError as e:
+            logger.error("Invalid IPv6 address: %s", e)
+            raise
+
+    def _get_ipv6_scan_options(self) -> dict[str, Any]:
+        """Get IPv6-specific scan options."""
+        return {
+            "ping_timeout": self.args.get(
+                "ipv6_ping_timeout", IPV6_DEFAULT_PING_TIMEOUT
+            ),
+            "min_rate": self.args.get("ipv6_min_rate", 100),
+            "max_rate": self.args.get("ipv6_max_rate", 1000),
+            "max_retries": self.args.get("ipv6_max_retries", 2),
+            "fragment_mtu": self.args.get("ipv6_fragment_mtu", 1280),
+        }
 
     def start(self) -> None:
         if self._vpn_config is not None and self._dns_config is not None:
@@ -102,20 +133,37 @@
             else:
                 hosts = [(host, mask)]
         elif "v6" in message.selector:
+            # Handle IPv6 address normalization
+            try:
+                normalized_host = self._normalize_ipv6_address(host)
+            except ipaddress.AddressValueError:
+                logger.error("Invalid IPv6 address: %s", host)
+                return
+
             mask = int(message.data.get("mask", DEFAULT_MASK_IPV6))
             if mask < IPV6_CIDR_LIMIT:
-                raise ValueError(
-                    f"Subnet mask below {IPV6_CIDR_LIMIT} is not supported"
+                logger.error(
+                    "IPv6 subnet mask below %s is not supported", IPV6_CIDR_LIMIT
                 )
+                return
 
+            # Validate the mask
             max_mask = int(self.args.get("max_network_mask_ipv6", "128"))
             if mask < max_mask:
-                for subnet in ipaddress.ip_network(
-                    f"{host}/{mask}", strict=False
-                ).subnets(new_prefix=max_mask):
-                    hosts.append((str(subnet.network_address), max_mask))
+                try:
+                    # Only process valid subnets
+                    for subnet in ipaddress.ip_network(
+                        f"{normalized_host}/{mask}", strict=False
+                    ).subnets(new_prefix=max_mask):
+                        if isinstance(subnet, ipaddress.IPv6Network):
+                            hosts.append((str(subnet.network_address), max_mask))
+                        else:
+                            logger.error("Unexpected subnet type for IPv6: %s", subnet)
+                except ValueError as e:
+                    logger.error("Failed to process IPv6 network: %s", e)
+                    return
             else:
-                hosts = [(host, mask)]
+                hosts = [(normalized_host, mask)]
 
         domain_name = self._prepare_domain_name(
             message.data.get("name"), message.data.get("url")
@@ -163,6 +211,8 @@
             logger.error("Neither host or domain are set.")
 
     def _scan_host(self, host: str, mask: int) -> Tuple[Dict[str, Any], str]:
+        is_ipv6 = ":" in host  # Simple check for IPv6 address
+
         options = nmap_options.NmapOptions(
             dns_resolution=False,
             ports=self.args.get("ports"),
@@ -175,11 +225,25 @@
             script_default=self.args.get("script_default", False),
             version_detection=self.args.get("version_info", False),
         )
-        client = nmap_wrapper.NmapWrapper(options)
 
+        if is_ipv6 is True:
+            # Add IPv6-specific options
+            ipv6_options = self._get_ipv6_scan_options()
+            options.ping_timeout = ipv6_options["ping_timeout"]
+            options.min_rate = ipv6_options["min_rate"]
+            options.max_rate = ipv6_options["max_rate"]
+            options.max_retries = ipv6_options["max_retries"]
+            options.fragment_mtu = ipv6_options["fragment_mtu"]
+
+        client = nmap_wrapper.NmapWrapper(options)
         logger.info("scanning target %s/%s with options %s", host, mask, options)
-        scan_results, normal_results = client.scan_hosts(hosts=host, mask=mask)
-        return scan_results, normal_results
+
+        try:
+            scan_results, normal_results = client.scan_hosts(hosts=host, mask=mask)
+            return scan_results, normal_results
+        except subprocess.CalledProcessError as e:
+            logger.error("Nmap scan failed for IPv6 host %s: %s", host, e)
+            raise
 
     def _scan_domain(self, domain_name: str) -> Tuple[Dict[str, Any], str]:
         options = nmap_options.NmapOptions(

agent/nmap_options.py

@@ -63,6 +63,12 @@ class NmapOptions:
     no_ping: bool = True
     privileged: Optional[bool] = None
 
+    ping_timeout: int | None = None  # Example: Timeout for pings (in milliseconds)
+    min_rate: int | None = None  # Minimum number of packets per second
+    max_rate: int | None = None  # Maximum number of packets per second
+    max_retries: int | None = None  # Maximum number of retries for probes
+    fragment_mtu: int | None = None  # MTU for fragmented packets
+
     def _set_os_detection_option(self) -> List[str]:
         """Appends the os detection option to the list of nmap options."""
         command_options = []
@@ -151,6 +157,21 @@ def _run_scripts_command(self, scripts: List[str]) -> List[str]:
 
         return command
 
+    def _set_additional_options(self) -> List[str]:
+        """Appends additional options for fine-tuning Nmap behavior."""
+        command_options = []
+        if self.ping_timeout is not None:
+            command_options.append(f"--host-timeout {self.ping_timeout}ms")
+        if self.min_rate is not None:
+            command_options.append(f"--min-rate {self.min_rate}")
+        if self.max_rate is not None:
+            command_options.append(f"--max-rate {self.max_rate}")
+        if self.max_retries is not None:
+            command_options.append(f"--max-retries {self.max_retries}")
+        if self.fragment_mtu is not None:
+            command_options.append(f"--mtu {self.fragment_mtu}")
+        return command_options
+
     @property
     def command_options(self) -> List[str]:
         """Computes the list of nmap options."""
@@ -165,4 +186,5 @@ def command_options(self) -> List[str]:
         command_options.extend(self._set_privileged())
         command_options.extend(self._set_scripts())
         command_options.extend(self._set_script_default())
+        command_options.extend(self._set_additional_options())
         return command_options

tests/conftest.py

@@ -290,15 +290,41 @@
     )
 
 
+@pytest.fixture
+def invalid_ipv6_msg() -> message.Message:
+    """Creates an invalid IPv6 message for testing error handling."""
+    return message.Message.from_data(
+        selector="v3.asset.ip.v6",
+        data={
+            "version": 6,
+            "host": "invalid_ipv6",
+            "mask": "112",
+        },
+    )
+
+
+@pytest.fixture
+def large_subnet_ipv6_msg() -> message.Message:
+    """Creates a message with a large IPv6 subnet."""
+    return message.Message.from_data(
+        selector="v3.asset.ip.v6",
+        data={
+            "version": 6,
+            "host": "2600:3c01:224a:6e00::",
+            "mask": "112",
+        },
+    )
+
+
 @pytest.fixture
 def ipv6_msg_above_limit() -> message.Message:
-    """Creates a dummy message of type v3.asset.ip.v6 for testing purposes."""
+    """Creates a message with an IPv6 subnet above the allowed limit (below mask 112)."""
     return message.Message.from_data(
         selector="v3.asset.ip.v6",
         data={
             "version": 6,
-            "host": "2600:3c01:224a:6e00:f03c:91ff:fe18:bb2f",
-            "mask": "64",
+            "host": "2600:3c01:224a:6e00::",
+            "mask": "96",  # Below IPV6_CIDR_LIMIT of 112
         },
     )