Merge pull request #347 from 13ph03nix/fix-ipv6

fix: ipv6 compatibility issue in build_url
knownsec · Dec 12, 2022 · 5d111ed · 5d111ed
2 parents e7bc557 + 614a6f1
commit 5d111ed
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 4 deletions.
diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
@@ -56,3 +56,6 @@ Rook1e <https://github.com/0x2E>
 
 ekszz <https://github.com/ekszz>
 * contributing to customize poc protocol and default port #321
+
+HomerQing <https://github.com/HomerQing>
+* contributing to fix ipv6 compatibility issue in build_url
diff --git a/pocsuite3/lib/core/poc.py b/pocsuite3/lib/core/poc.py
@@ -192,6 +192,7 @@ def build_url(self):
 
         try:
             pr = urlparse(target)
+            is_ipv6 = pr.netloc.startswith('[')
             self.scheme = pr.scheme
             self.rhost = pr.hostname
             self.rport = pr.port or self.current_protocol_port
@@ -209,7 +210,7 @@ def build_url(self):
                 # adjust port
                 if not self.rport:
                     self.rport = protocol_default_port_map[self.current_protocol]
-            self.netloc = f'{self.rhost}:{self.rport}'
+            self.netloc = f'[{self.rhost}]:{self.rport}' if is_ipv6 else f'{self.rhost}:{self.rport}'
             pr = pr._replace(scheme=self.scheme)
             pr = pr._replace(netloc=self.netloc)
             target = pr.geturl()
@@ -366,9 +367,6 @@ def _check(self, dork='', allow_redirects=False, return_obj=False, is_http=True,
         if self.url.split('://')[0] != self.scheme:
             logger.warn(f'auto correct url: {mosaic(origin_url)} -> {mosaic(self.url)}')
             self.scheme = 'https' if self.url.startswith('https') else 'http'
-            port = urlparse(self.url).port
-            self.rport = port if port else 443 if self.scheme.startswith('https') else 80
-            self.netloc = f'{self.rhost}:{self.rport}'
 
         if return_obj:
             return res

diff --git a/tests/test_build_url.py b/tests/test_build_url.py
@@ -133,6 +133,23 @@ def _verify(self):
             self.assertEqual(res[0]["result"]["VerifyInfo"]["rport"], 8443)
             self.assertEqual(res[0]["result"]["VerifyInfo"]["netloc"], "127.0.0.1:8443")
 
+            # [fd12:3456:789a:1::2]:8443
+            f.seek(0)
+            config = {
+                "url": "[fd12:3456:789a:1::2]:8443",
+                "poc": f.name,
+            }
+            init_pocsuite(config)
+            start_pocsuite()
+            res = get_results()
+            self.assertEqual(
+                res[0]["result"]["VerifyInfo"]["url"], "https://[fd12:3456:789a:1::2]:8443"
+            )
+            self.assertEqual(res[0]["result"]["VerifyInfo"]["scheme"], "https")
+            self.assertEqual(res[0]["result"]["VerifyInfo"]["rhost"], "fd12:3456:789a:1::2")
+            self.assertEqual(res[0]["result"]["VerifyInfo"]["rport"], 8443)
+            self.assertEqual(res[0]["result"]["VerifyInfo"]["netloc"], "[fd12:3456:789a:1::2]:8443")
+
     def test_url_protocol_correct(self):
         with CustomNamedTemporaryFile("w+t") as f:
             poc_content = textwrap.dedent(