Improvements for Windows compatibility

wapitiCore/attack/attack.py

@@ -27,8 +27,8 @@
 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 @@ class Attack:
     # 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 @@ class Attack:
 
     @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 @@ def __init__(self, scope_request: Request, scope="folder", session_dir=None, con
         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 @@ async def load_and_attack(self, attack_module: Attack, attacked_ids: Set[int]) -
                 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, signum, frame) -> None:
         """
         Attack handler for Ctrl+C interruption.
         """
@@ -495,7 +496,8 @@ def handle_user_interruption(self, task) -> None:
             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 def attack(self):
         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 @@ async def attack(self):
                     )
 
                 # 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(signum, frame):
     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 @@ def generate_report(self, output_path):
                     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"
         )