Merge pull request #1417 from maurosoria/fix-session

Fix session file does not work (in both async and normal modes)

lib/connection/requester.py

@@ -63,8 +63,8 @@
 class BaseRequester:
     def __init__(self) -> None:
         self._url: str = ""
-        self._proxy_cred: str = ""
         self._rate = 0
+        self.proxy_cred = options["proxy_auth"]
         self.headers = CaseInsensitiveDict(options["headers"])
         self.agents: list[str] = []
         self.session = None
@@ -108,17 +108,14 @@ def set_proxy(self, proxy: str) -> None:
         if not proxy.startswith(PROXY_SCHEMES):
             proxy = f"http://{proxy}"
 
-        if self._proxy_cred and "@" not in proxy:
+        if self.proxy_cred and "@" not in proxy:
             # socks5://localhost:9050 => socks5://[credential]@localhost:9050
-            proxy = proxy.replace("://", f"://{self._proxy_cred}@", 1)
+            proxy = proxy.replace("://", f"://{self.proxy_cred}@", 1)
 
         self.session.proxies = {"https": proxy}
         if not proxy.startswith("https://"):
             self.session.proxies["http"] = proxy
 
-    def set_proxy_auth(self, credential: str) -> None:
-        self._proxy_cred = credential
-
     def is_rate_exceeded(self) -> bool:
         return self._rate >= options["max_rate"] > 0
 
@@ -162,6 +159,9 @@ def __init__(self):
                 ),
             )
 
+        if options["auth"]:
+            self.set_auth(options["auth_type"], options["auth"])
+
     def set_auth(self, type: str, credential: str) -> None:
         if type in ("bearer", "jwt"):
             self.session.auth = HTTPBearerAuth(credential)
@@ -312,16 +312,19 @@ def __init__(self) -> None:
         )
         self.replay_session = None
 
+        if options["auth"]:
+            self.set_auth(options["auth_type"], options["auth"])
+
     def parse_proxy(self, proxy: str) -> str:
         if not proxy:
             return None
 
         if not proxy.startswith(PROXY_SCHEMES):
             proxy = f"http://{proxy}"
 
-        if self._proxy_cred and "@" not in proxy:
+        if self.proxy_cred and "@" not in proxy:
             # socks5://localhost:9050 => socks5://[credential]@localhost:9050
-            proxy = proxy.replace("://", f"://{self._proxy_cred}@", 1)
+            proxy = proxy.replace("://", f"://{self.proxy_cred}@", 1)
 
         return proxy
 

lib/controller/controller.py

@@ -26,6 +26,10 @@
 import re
 import time
 import mysql.connector
+try:
+    import cPickle as pickle
+except ModuleNotFoundError:
+    import pickle
 
 from urllib.parse import urlparse
 
@@ -60,27 +64,18 @@
 from lib.report.manager import ReportManager
 from lib.utils.common import lstrip_once
 from lib.utils.file import FileUtils
-from lib.utils.pickle import pickle, unpickle
 from lib.utils.schemedet import detect_scheme
 from lib.view.terminal import interface
 
-if options["async_mode"]:
-    from lib.connection.requester import AsyncRequester as Requester
-    from lib.core.fuzzer import AsyncFuzzer as Fuzzer
-
-    try:
-        import uvloop
-        asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
-    except ImportError:
-        pass
-else:
-    from lib.connection.requester import Requester
-    from lib.core.fuzzer import Fuzzer
-
 
 class Controller:
     def __init__(self) -> None:
         if options["session_file"]:
+            print("WARNING: Running an untrusted session file might lead to unwanted code execution!")
+            interface.in_line("[c]continue / [q]uit: ")
+            if input() != "c":
+                exit(1)
+
             self._import(options["session_file"])
             self.old_session = True
         else:
@@ -92,26 +87,30 @@ def __init__(self) -> None:
     def _import(self, session_file: str) -> None:
         try:
             with open(session_file, "rb") as fd:
-                indict, last_output, opt = unpickle(fd)
+                dict_, last_output, opt = pickle.load(fd)
                 options.update(opt)
         except UnpicklingError:
             interface.error(
                 f"{session_file} is not a valid session file or it's in an old format"
             )
             exit(1)
 
-        self.__dict__ = {**indict, **vars(self)}
+        self.__dict__ = {**dict_, **vars(self)}
         print(last_output)
 
     def _export(self, session_file: str) -> None:
         # Save written output
         last_output = interface.buffer.rstrip()
 
-        # Can't pickle Fuzzer class due to _thread.lock objects
-        del self.fuzzer
+        dict_ = vars(self).copy()
+        # Can't pickle some classes due to _thread.lock objects
+        dict_.pop("fuzzer", None)
+        dict_.pop("pause_future", None)
+        dict_.pop("loop", None)
+        dict_.pop("requester", None)
 
         with open(session_file, "wb") as fd:
-            pickle((vars(self), last_output, options), fd)
+            pickle.dump((dict_, last_output, options), fd)
 
     def setup(self) -> None:
         blacklists.update(get_blacklists())
@@ -136,7 +135,6 @@ def setup(self) -> None:
             if options["cookie"]:
                 options["headers"]["cookie"] = options["cookie"]
 
-        self.requester = Requester()
         self.dictionary = Dictionary(files=options["wordlists"])
         self.start_time = time.time()
         self.passed_urls: set[str] = set()
@@ -145,16 +143,6 @@ def setup(self) -> None:
         self.errors = 0
         self.consecutive_errors = 0
 
-        if options["async_mode"]:
-            self.loop = asyncio.new_event_loop()
-            self.loop.add_signal_handler(signal.SIGINT, self.handle_pause)
-
-        if options["auth"]:
-            self.requester.set_auth(options["auth_type"], options["auth"])
-
-        if options["proxy_auth"]:
-            self.requester.set_proxy_auth(options["proxy_auth"])
-
         if options["log_file"]:
             try:
                 FileUtils.create_dir(FileUtils.parent(options["log_file"]))
@@ -187,6 +175,19 @@ def setup(self) -> None:
             interface.log_file(options["log_file"])
 
     def run(self) -> None:
+        if options["async_mode"]:
+            from lib.connection.requester import AsyncRequester as Requester
+            from lib.core.fuzzer import AsyncFuzzer as Fuzzer
+
+            try:
+                import uvloop
+                asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+            except ImportError:
+                pass
+        else:
+            from lib.connection.requester import Requester
+            from lib.core.fuzzer import Fuzzer
+
         # match_callbacks and not_found_callbacks callback values:
         #  - *args[0]: lib.connection.Response() object
         #
@@ -200,6 +201,11 @@ def run(self) -> None:
         )
         error_callbacks = (self.raise_error, self.append_error_log)
 
+        self.requester = Requester()
+        if options["async_mode"]:
+            self.loop = asyncio.new_event_loop()
+            self.loop.add_signal_handler(signal.SIGINT, self.handle_pause)
+
         while options["urls"]:
             url = options["urls"][0]
             self.fuzzer = Fuzzer(
@@ -446,14 +452,6 @@ def handle_pause(self) -> None:
             option = input()
 
             if option.lower() == "q":
-                if options["async_mode"]:
-                    quitexc = QuitInterrupt("Canceled by the user")
-                    if options["async_mode"]:
-                        self.pause_future.set_exception(quitexc)
-                        break
-                    else:
-                        raise quitexc
-
                 interface.in_line("[s]ave / [q]uit without saving: ")
 
                 option = input()

lib/core/options.py

@@ -39,9 +39,6 @@ def parse_options() -> dict[str, Any]:
     opt = merge_config(parse_arguments())
 
     if opt.session_file:
-        if opt.async_mode:
-            print("Cannot resume a session in asynchronous mode")
-            exit(1)
         return vars(opt)
 
     opt.http_method = opt.http_method.upper()

lib/parse/cmdline.py

@@ -168,6 +168,7 @@ def parse_arguments() -> Values:
         help="Number of threads",
     )
     general.add_option(
+        "-a",
         "--async",
         action="store_true",
         dest="async_mode",