graceful exit for payloads

src/wfuzz/core.py

@@ -97,6 +97,7 @@ def __init__(self, options):
         self.options = options
         self.seed = FuzzResultFactory.from_options(options)
         self.baseline = FuzzResultFactory.from_baseline(self.seed, options)
+        self._payload_list = []
         self.dictio = self.get_dictio()
 
         self.stats = FuzzStats.from_requestGenerator(self)
@@ -175,6 +176,10 @@ def __next__(self):
 
             return FuzzResultFactory.from_seed(self.seed, n, self.options)
 
+    def close(self):
+        for payload in self._payload_list:
+            payload.close()
+
     def get_dictio(self):
         class wrapper(object):
             def __init__(self, iterator):
@@ -190,6 +195,7 @@ def __next__(self):
                 return str(next(self._it))
 
         selected_dic = []
+        self._payload_list = []
 
         if self.options["dictio"]:
             for d in [wrapper(x) for x in self.options["dictio"]]:
@@ -205,6 +211,7 @@ def __next__(self):
                     raise FuzzExceptBadOptions("You must supply a list of payloads in the form of [(name, {params}), ... ]")
 
                 p = Facade().payloads.get_plugin(name)(params)
+                self._payload_list.append(p)
                 pp = dictionary(p, params["encoder"]) if "encoder" in params else p
                 selected_dic.append(sliceit(pp, slicestr) if slicestr else pp)
 

src/wfuzz/options.py

@@ -194,7 +194,9 @@ def export_json(self):
 
     def payload(self, **kwargs):
         self.data.update(kwargs)
-        return requestGenerator(self).get_dictio()
+        # TODO: maybe compile here? or think how to make work and not deadlock from the api
+        self.data['compiled_genreq'] = requestGenerator(self)
+        return self.data['compiled_genreq'].get_dictio()
 
     def fuzz(self, **kwargs):
         self.data.update(kwargs)
@@ -290,6 +292,12 @@ def compile(self):
         return self
 
     def close(self):
-        self.http_pool.deregister()
+        if self.http_pool:
+            self.http_pool.deregister()
+
         if self.fz:
             self.fz.cancel_job()
+
+        # TODO: deadlock when using api.payload()
+        elif self.data['compiled_genreq'] is not None:
+            self.data['compiled_genreq'].close()

src/wfuzz/plugin_api/base.py

@@ -121,6 +121,9 @@ def count(self):
     def __iter__(self):
         raise FuzzExceptPluginError("Method iter not implemented")
 
+    def close(self):
+        pass
+
     def find_file(self, name):
         if os.path.exists(name):
             return name

src/wfuzz/wfuzz.py

@@ -82,7 +82,8 @@ def usage():
 \t--field <field>     : Show a FuzzResult field instead of current payload
 """)
 
-    from .api import payload
+    # TODO: from .api import payload
+    from .api import FuzzSession
     from .exception import FuzzExceptBadOptions
     import getopt
 
@@ -105,18 +106,22 @@ def usage():
         if o in ("--field"):
             field = value
 
+    session = None
+
     try:
         session_options = CLParser(sys.argv).parse_cl()
+        # TODO: api.payload will block with new shodanp
+        session = FuzzSession(**session_options)
         printer = None
 
-        for res in payload(**session_options):
+        for res in session.payload():
             if len(res) > 1:
                 raise FuzzExceptBadOptions("wfpayload can only be used to generate one word dictionaries")
             else:
                 r = res[0]
 
             # TODO: option to not show headers in fuzzres
-            # TODO: all should be same object
+            # TODO: all should be same object type and no need for isinstance
             if isinstance(r, FuzzResult):
                 if printer is None:
                     printer = View(session_options)
@@ -134,6 +139,9 @@ def usage():
         print(("\nFatal exception: %s" % str(e)))
     except Exception as e:
         print(("\nUnhandled exception: %s" % str(e)))
+    finally:
+        if session:
+            session.close()
 
 
 def main_encoder():