fix IPs persisted as domain names

agent/helpers.py

@@ -89,12 +89,14 @@ def build_vuln_location(
     asset: ipv4_asset.IPv4 | ipv6_asset.IPv6 | domain_asset.DomainName
     ip = None
     port = None
-    if is_ipv4(matched_at) is True:
-        ip, port = split_ipv4(matched_at)
+    potential_ip = matched_at
+    if target.scheme != "":
+        potential_ip = potential_ip.replace(f"{target.scheme}://", "")
+    if is_ipv4(potential_ip) is True:
+        ip, port = split_ipv4(potential_ip)
         asset = ipv4_asset.IPv4(host=str(ip), version=4, mask="32")
-    elif is_ipv6(matched_at) is not False:
-        ip = matched_at
-        asset = ipv6_asset.IPv6(host=str(ip), version=4, mask="128")
+    elif is_ipv6(potential_ip) is True:
+        asset = ipv6_asset.IPv6(host=str(potential_ip), version=6, mask="128")
     else:
         asset = domain_asset.DomainName(name=prepare_domain_asset(matched_at))
 

tests/agent_test.py

@@ -343,7 +343,7 @@ def testAgentNuclei_whenIpv6Scanned_emitsExactIpWhereVulnWasFound(
         "ipv6": {
             "host": "FE80:CD00:0000:0CDE:1257:0000:211E:729C",
             "mask": "128",
-            "version": 4,
+            "version": 6,
         }
     }
 

tests/helpers_test.py

@@ -0,0 +1,63 @@
+"""Unit tests for the helpers module."""
+
+from ostorlab.assets import domain_name
+from ostorlab.assets import ipv4
+from ostorlab.assets import ipv6
+
+from agent import helpers
+
+
+def testBuildVulnLocation_whenMatchedAtIsIpv4_returnsVulnLocation() -> None:
+    """Ensure that when matched_at is an IPv4, BuildVulnLocation returns a valid VulnLocation."""
+    matched_at = "70.70.70.70:443"
+
+    vuln_location = helpers.build_vuln_location(matched_at)
+
+    assert vuln_location is not None
+    ipv4_asset = vuln_location.asset
+    assert isinstance(ipv4_asset, ipv4.IPv4)
+    assert ipv4_asset.host == "70.70.70.70"
+    assert ipv4_asset.version == 4
+    assert ipv4_asset.mask == "32"
+
+
+def testBuildVulnLocation_whenMatchedAtIsIpv6_returnsVulnLocation() -> None:
+    """Ensure that when matched_at is an IPv6, BuildVulnLocation returns a valid VulnLocation."""
+    matched_at = "2001:0db8:85a3:0000:0000:8a2e:0370:7334"
+
+    vuln_location = helpers.build_vuln_location(matched_at)
+
+    assert vuln_location is not None
+    ipv6_asset = vuln_location.asset
+    assert isinstance(ipv6_asset, ipv6.IPv6)
+    assert ipv6_asset.host == "2001:0db8:85a3:0000:0000:8a2e:0370:7334"
+    assert ipv6_asset.version == 6
+    assert ipv6_asset.mask == "128"
+
+
+def testBuildVulnLocation_whenMatchedAtIsDomain_returnsVulnLocation() -> None:
+    """Ensure that when matched_at is a domain, BuildVulnLocation returns a valid VulnLocation."""
+    matched_at = "https://www.google.com"
+
+    vuln_location = helpers.build_vuln_location(matched_at)
+
+    assert vuln_location is not None
+    domain_asset = vuln_location.asset
+    assert isinstance(domain_asset, domain_name.DomainName)
+    assert domain_asset.name == "www.google.com"
+
+
+def testBuildVulnLocation_whenMatchedAtIsIpv4WithScheme_returnsValidVulnLocation() -> (
+    None
+):
+    """Ensure that when a scheme is present, BuildVulnLocation returns a valid VulnLocation."""
+    matched_at = "https://70.70.70.70"
+
+    vuln_location = helpers.build_vuln_location(matched_at)
+
+    assert vuln_location is not None
+    ipv4_asset = vuln_location.asset
+    assert isinstance(ipv4_asset, ipv4.IPv4)
+    assert ipv4_asset.host == "70.70.70.70"
+    assert ipv4_asset.version == 4
+    assert ipv4_asset.mask == "32"