Improvements for Windows compatibility

wapitiCore/attack/attack.py

@@ -19,16 +19,16 @@
 # Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 import dataclasses
 import os
-from os.path import splitext, join as path_join
+from os.path import splitext
 from urllib.parse import quote, urlparse
 from collections import defaultdict
 from enum import Enum, Flag, auto
 import random
 from binascii import hexlify
 from typing import Optional, Iterator, Tuple, List, Callable, Union, Iterable, Type
 import json
+from importlib.resources import files
 
-from pkg_resources import resource_filename
 from httpx import ReadTimeout, RequestError
 
 from wapitiCore.definitions import FindingBase
@@ -191,7 +191,7 @@
     # Must be defined in the code of the module
     require = []
 
-    DATA_DIR = resource_filename("wapitiCore", os.path.join("data", "attacks"))
+    DATA_DIR = str(files("wapitiCore").joinpath("data", "attacks"))
     HOME_DIR = os.getenv("HOME") or os.getenv("USERPROFILE") or "home"
 
     # Color codes
@@ -217,7 +217,7 @@
 
     @staticmethod
     def get_resource(resource_path: str):
-        return resource_filename("wapitiCore", path_join("data", "attacks", resource_path))
+        return str(files("wapitiCore").joinpath("data", "attacks", resource_path))
 
     def __init__(
             self,

wapitiCore/controller/wapiti.py

@@ -198,6 +198,7 @@
         self._wait_time = 2.
         self._buffer = []
         self._user_choice = UserChoice.CONTINUE
+        self._current_attack_task: Optional[asyncio.Task] = None
 
         if session_dir:
             SqlPersister.CRAWLER_DATA_DIR = session_dir
@@ -481,7 +482,7 @@
                 if original_request.path_id is not None:
                     attacked_ids.add(original_request.path_id)
 
-    def handle_user_interruption(self, task) -> None:
+    def handle_user_interruption(self, _, __) -> None:
         """
         Attack handler for Ctrl+C interruption.
         """
@@ -495,7 +496,8 @@
             try:
                 self._user_choice = UserChoice(input("? ").strip().lower())
                 if self._user_choice != UserChoice.CONTINUE:
-                    task.cancel()
+                    if self._current_attack_task is not None:
+                        self._current_attack_task.cancel()
                 return
             except (UnicodeDecodeError, ValueError):
                 print("Invalid choice. Valid choices are r, n, q, and c.")
@@ -540,8 +542,6 @@
         async with AsyncCrawler.with_configuration(self.crawler_configuration) as crawler:
             attack_modules = await self._load_attack_modules(crawler)
 
-            loop = asyncio.get_running_loop()
-
             for attack_module in attack_modules:
                 if attack_module.do_get is False and attack_module.do_post is False:
                     continue
@@ -565,21 +565,21 @@
                     )
 
                 # Create and run each attack module as an asyncio task
-                current_attack_task = asyncio.create_task(
+                self._current_attack_task = asyncio.create_task(
                     self.run_attack_module(attack_module)
                 )
 
                 # Setup signal handler to prompt the user for task cancellation
-                loop.add_signal_handler(signal.SIGINT, self.handle_user_interruption, current_attack_task)
+                signal.signal(signal.SIGINT, self.handle_user_interruption)
 
                 try:
-                    await current_attack_task  # Await the attack module task
+                    await self._current_attack_task  # Await the attack module task
                 except asyncio.CancelledError:
                     # The user chose to stop the current module
                     pass
                 finally:
                     # Clean up the signal handler for the next loop
-                    loop.remove_signal_handler(signal.SIGINT)
+                    signal.signal(signal.SIGINT, signal.SIG_DFL)
 
                 # As the handler directly continue or cancel the current_attack_task module, we don't have
                 # cases where we have to call `continue`. Just check for the two other options

wapitiCore/main/wapiti.py

@@ -44,7 +44,7 @@
 global_stop_event = asyncio.Event()
 
 
-def inner_ctrl_c_signal_handler():
+def inner_ctrl_c_signal_handler(_, __):
     logging.info("Waiting for running crawler tasks to finish, please wait.")
     global_stop_event.set()
 
@@ -463,8 +463,6 @@ async def wapiti_main():
         for url in excluded_urls:
             wap.add_excluded_url(url)
 
-    loop = asyncio.get_event_loop()
-
     try:
         if not args.skip_crawl:
             if await wap.have_attacks_started() and not args.resume_crawl:
@@ -474,10 +472,12 @@ async def wapiti_main():
                     logging.info("[*] Resuming scan from previous session, please wait")
 
                 await wap.load_scan_state()
-                loop.add_signal_handler(signal.SIGINT, inner_ctrl_c_signal_handler)
-                await wap.browse(global_stop_event, parallelism=args.tasks)
-                loop.remove_signal_handler(signal.SIGINT)
-                await wap.save_scan_state()
+                signal.signal(signal.SIGINT, inner_ctrl_c_signal_handler)
+                try:
+                    await wap.browse(global_stop_event, parallelism=args.tasks)
+                finally:
+                    signal.signal(signal.SIGINT, signal.SIG_DFL)
+                    await wap.save_scan_state()
 
         if args.max_parameters:
             count = await wap.persister.remove_big_requests(args.max_parameters)

wapitiCore/report/htmlreportgenerator.py

@@ -21,10 +21,10 @@
 # Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 
 import os
+from importlib.resources import files
 from shutil import copytree, rmtree, copy
 from urllib.parse import urlparse
 import time
-from pkg_resources import resource_filename
 
 from mako.template import Template
 
@@ -71,16 +71,16 @@
                     pass
 
                 copytree(
-                    resource_filename("wapitiCore", os.path.join(self.REPORT_DIR, subdir)),
+                    str(files("wapitiCore").joinpath(self.REPORT_DIR, subdir)),
                     os.path.join(output_path, subdir)
                 )
 
-            copy(resource_filename("wapitiCore", os.path.join(self.REPORT_DIR, "logo_clear.png")), output_path)
+            copy(str(files("wapitiCore").joinpath(self.REPORT_DIR, "logo_clear.png")), output_path)
         else:
-            copytree(resource_filename("wapitiCore", self.REPORT_DIR), output_path)
+            copytree(str(files("wapitiCore").joinpath(self.REPORT_DIR)), output_path)
 
         mytemplate = Template(
-            filename=resource_filename("wapitiCore", os.path.join(self.REPORT_DIR, "report.html")),
+            filename=str(files("wapitiCore").joinpath(self.REPORT_DIR, "report.html")),
             input_encoding="utf-8",
             output_encoding="utf-8"
         )