Fix max-scan-time and missing timeout in headless explorer

Signed-off-by: bretfourbe <gwendal@cyberwatch.fr>

Author: bretfourbe

wapitiCore/controller/wapiti.py

@@ -177,7 +177,7 @@ def __init__(self, scope_request: Request, scope="folder", session_dir=None, con
         self._max_links_per_page = 0
         self._max_files_per_dir = 0
         self._scan_force = "normal"
-        self._max_scan_time = 0
+        self._max_scan_time = None
         self._max_attack_time = None
         self._bug_report = True
         self._logfile = ""
@@ -376,6 +376,16 @@ async def save_scan_state(self):
         # if stopped and self._start_urls:
         #     print(_("The scan will be resumed next time unless you pass the --skip-crawl option."))
 
+    async def explore_and_save_requests(self, explorer):
+        self._buffer = []
+        # Browse URLs are saved them once we have enough in our buffer
+        async for resource, response in explorer.async_explore(self._start_urls, self._excluded_urls):
+            self._buffer.append((resource, response))
+
+            if len(self._buffer) > 100:
+                await self.persister.save_requests(self._buffer)
+                self._buffer = []
+
     async def browse(self, stop_event: asyncio.Event, parallelism: int = 8):
         """Extract hyperlinks and forms from the webpages found on the website"""
         stop_event.clear()
@@ -406,22 +416,21 @@ async def browse(self, stop_event: asyncio.Event, parallelism: int = 8):
         explorer.qs_limit = SCAN_FORCE_VALUES[self._scan_force]
         explorer.load_saved_state(self.persister.output_file[:-2] + "pkl")
 
-        start = datetime.utcnow()
-        buffer = []
-
-        # Browse URLs are saved them once we have enough in our buffer
-        async for resource, response in explorer.async_explore(self._start_urls, self._excluded_urls):
-            buffer.append((resource, response))
-
-            if len(buffer) > 100:
-                await self.persister.save_requests(buffer)
-                buffer = []
+        self._buffer = []
 
-            if not stop_event.is_set() and (datetime.utcnow() - start).total_seconds() > self._max_scan_time >= 1:
-                logging.info("Max scan time was reached, stopping.")
+        try:
+            await asyncio.wait_for(
+               self.explore_and_save_requests(explorer),
+               self._max_scan_time
+            )
+        except asyncio.TimeoutError:
+            logging.info("Max scan time was reached, stopping.")
+            if not stop_event.is_set():
                 stop_event.set()
+        finally:
+            await explorer.clean()
 
-        await self.persister.save_requests(buffer)
+        await self.persister.save_requests(self._buffer)
 
         # Let's save explorer values (limits)
         explorer.save_state(self.persister.output_file[:-2] + "pkl")

wapitiCore/net/explorer.py

@@ -446,6 +446,7 @@ async def async_explore(
             if not task_to_request and (self._stopped.is_set() or not to_explore):
                 break
 
+    async def clean(self):
         self._cookiejar = self._crawler.cookie_jar
         await self._crawler.close()
 

wapitiCore/net/intercepting_explorer.py

@@ -332,7 +332,7 @@ async def launch_headless_explorer(
                         continue
                 else:
                     try:
-                        response = await crawler.async_send(request)
+                        response = await crawler.async_send(request, timeout=crawler.timeout.connect)
                     except httpx.RequestError as exception:
                         logging.error(f"{request} generated an exception: {exception.__class__.__name__}")
                         continue
@@ -389,13 +389,52 @@ def __init__(
         self._final_cookies = None
         self._cookies = cookies or CookieJar()
         self._wait_time = wait_time
+        self._headless_task = None
+
+    async def process_requests(self, excluded_requests, exclusion_regexes):
+        while True:
+            try:
+                request, response = self._queue.get_nowait()
+            except asyncio.QueueEmpty:
+                await asyncio.sleep(.1)
+            except KeyboardInterrupt:
+                break
+            else:
+                self._queue.task_done()
+
+                # Scope check and deduplication are made here
+                if not self._scope.check(request) or request in self._processed_requests:
+                    continue
+
+                # Check for exclusion here because we don't have full control over the headless browser
+                if request in excluded_requests or any(regex.match(request.url) for regex in exclusion_regexes):
+                    continue
+
+                dir_name = request.dir_name
+                if self._max_files_per_dir and self._file_counts[dir_name] >= self._max_files_per_dir:
+                    continue
+
+                self._file_counts[dir_name] += 1
+
+                if self.has_too_many_parameters(request):
+                    continue
+
+                if self._qs_limit and request.parameters_count:
+                    self._pattern_counts[request.pattern] += 1
+
+                yield request, response
+                self._processed_requests.append(request)
+                log_verbose(f"[+] {request}")
+
+            if self._stopped.is_set():
+                break
 
     async def async_explore(
             self,
             to_explore: Deque[Request],
             excluded_urls: list = None
     ) -> AsyncIterator[Tuple[Request, Response]]:
-        queue = asyncio.Queue()
+        self._queue = asyncio.Queue()
 
         exclusion_regexes = []
         excluded_requests = []
@@ -408,10 +447,10 @@ async def async_explore(
                     excluded_requests.append(bad_request)
 
         # Launch proxy as asyncio task
-        mitm_task = asyncio.create_task(
+        self._mitm_task = asyncio.create_task(
             launch_proxy(
                 self._mitm_port,
-                queue,
+                self._queue,
                 self._crawler.headers,
                 self._cookies,
                 self._scope,
@@ -420,12 +459,12 @@ async def async_explore(
             )
         )
 
-        headless_task = None
+        
         if self._headless == "no":
             # No headless crawler, just intercepting mode so no starting URLs
             to_explore.clear()
         else:
-            headless_task = asyncio.create_task(
+            self._headless_task = asyncio.create_task(
                 launch_headless_explorer(
                     self._stopped,
                     self._crawler,
@@ -440,52 +479,23 @@ async def async_explore(
                 )
             )
 
-        while True:
-            try:
-                request, response = queue.get_nowait()
-            except asyncio.QueueEmpty:
-                await asyncio.sleep(.1)
-            except KeyboardInterrupt:
-                break
-            else:
-                queue.task_done()
-
-                # Scope check and deduplication are made here
-                if not self._scope.check(request) or request in self._processed_requests:
-                    continue
-
-                # Check for exclusion here because we don't have full control over the headless browser
-                if request in excluded_requests or any(regex.match(request.url) for regex in exclusion_regexes):
-                    continue
-
-                dir_name = request.dir_name
-                if self._max_files_per_dir and self._file_counts[dir_name] >= self._max_files_per_dir:
-                    continue
-
-                self._file_counts[dir_name] += 1
-
-                if self.has_too_many_parameters(request):
-                    continue
-
-                if self._qs_limit and request.parameters_count:
-                    self._pattern_counts[request.pattern] += 1
-
-                yield request, response
-                self._processed_requests.append(request)
-                log_verbose(f"[+] {request}")
-
+        async for request, response in self.process_requests(excluded_requests, exclusion_regexes):
+            yield request, response
             if self._stopped.is_set():
                 break
 
-        await queue.join()
+    async def clean(self):
+        if not self._queue.empty():
+            await self._queue.join()
+
         # The headless crawler must stop when the stop event is set, let's just wait for it
-        if headless_task:
-            await headless_task
+        if self._headless_task:
+            await self._headless_task
 
         # We are canceling the mitm proxy, but we could have used a special request to shut down the master to.
         # https://docs.mitmproxy.org/stable/addons-examples/#shutdown
-        mitm_task.cancel()
-        self._final_cookies = await mitm_task
+        self._mitm_task.cancel()
+        self._final_cookies = await self._mitm_task
         await self._crawler.close()
 
     @property

wapitiCore/parsers/commandline.py

@@ -329,7 +329,7 @@ def parse_args():
         "--max-scan-time",
         metavar="SECONDS",
         help="Set how many seconds you want the scan to last (floats accepted)",
-        type=float, default=0
+        type=float, default=None
     )
 
     parser.add_argument(