Add subdomain scope option and unit test

bretfourbe · Oct 16, 2023 · 26753fe · 26753fe
1 parent ea6ff7a
commit 26753fe
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 14 deletions.
diff --git a/tests/web/test_scope.py b/tests/web/test_scope.py
@@ -38,7 +38,11 @@ async def test_scopes():
         "http://perdu.com/subdir/page.html?k=v",
         "http://perdu.com/subdir/page.html",
         "http://lost.com/lost.html",
-        "http://external.tld/external.html"
+        "http://external.tld/external.html",
+        "https://subdomain.perdu.com/",
+        "http://subdomain.perdu.com/",
+        "http://subdomain.perdu.com/page.html",
+        "http://subdomain.perdu.com/subdir/subdirpage.html",
     }
 
     scope = Scope(Request("http://perdu.com/subdir/"), "folder")
@@ -69,7 +73,19 @@ async def test_scopes():
         "http://sub.perdu.com/page.html",
         "https://perdu.com/secure.html",
         "http://perdu.com/subdir/page.html?k=v",
-        "http://perdu.com/subdir/page.html"
+        "http://perdu.com/subdir/page.html",
+        "http://subdomain.perdu.com/",
+        "http://subdomain.perdu.com/page.html",
+        "http://subdomain.perdu.com/subdir/subdirpage.html",
+        "https://subdomain.perdu.com/"
+    }
+
+    scope = Scope(Request("http://subdomain.perdu.com/subdir/page.html?k=v"), "subdomain")
+    assert scope.filter(links) == {
+        "http://subdomain.perdu.com/",
+        "http://subdomain.perdu.com/page.html",
+        "http://subdomain.perdu.com/subdir/subdirpage.html",
+        "https://subdomain.perdu.com/"
     }
 
     scope = Scope(Request("http://perdu.com/subdir/page.html?k=v"), "punk")

diff --git a/wapitiCore/net/scope.py b/wapitiCore/net/scope.py
@@ -46,29 +46,37 @@ def name(self) -> str:
         return self._scope
 
     def check(self, resource: Union[Request, str]) -> bool:
+
+        checked = None
+
         if not resource:
             return False
 
-        if self._scope == "punk":
-            # Life is short
-            return True
-
         if isinstance(resource, Request):
             url = resource.url
         else:
             url = resource
 
-        if self._scope == "domain":
-            return is_same_domain(url, self._base_request)
+        if self._scope == "punk":
+            # Life is short
+            checked = True
+
+        elif self._scope == "domain":
+            checked = is_same_domain(url, self._base_request)
+
+        elif self._scope == "subdomain":
+            checked = urlparse(url).hostname == self._base_request.hostname
 
-        if self._scope == "folder":
-            return url.startswith(self._base_request.path)
+        elif self._scope == "folder":
+            checked = url.startswith(self._base_request.path)
 
-        if self._scope == "page":
-            return url.split("?")[0] == self._base_request.path
+        elif self._scope == "page":
+            checked = url.split("?")[0] == self._base_request.path
 
         # URL
-        return url == self._base_request.url
+        if checked is None:
+            checked = url == self._base_request.url
+        return checked
 
     def filter(self, resources: Iterable[Union[Request, str]]) -> Set[Union[Request, str]]:
         return {resource for resource in resources if self.check(resource)}

diff --git a/wapitiCore/parsers/commandline.py b/wapitiCore/parsers/commandline.py
@@ -45,7 +45,7 @@ def parse_args():
         "--scope",
         help="Set scan scope",
         default="folder",
-        choices=["page", "folder", "domain", "url", "punk"]
+        choices=["page", "folder", "subdomain", "domain", "url", "punk"]
     )
 
     parser.add_argument(