Fix unclosed resources, clean-up signal handling, bump to python3.12 …

…to avoid signal handling issue in asyncio

setup.py

@@ -8,7 +8,7 @@
     url="https://github.com/delvelabs/tachyon",
     version=__version__,
     packages=find_packages(),
-    python_requires='>=3.6.0,<3.13.0',
+    python_requires='>=3.12.0,<4',
     package_data={'tachyon': ['data/*.json']},
     entry_points={
         'console_scripts': [

tachyon/__main__.py

@@ -263,8 +263,7 @@ def main(*, target_host, cookie_file, json_output, max_retry_count, plugin_setti
     conf.target_host = parsed_url.netloc
     conf.base_url = "%s://%s" % (parsed_url.scheme, parsed_url.netloc)
     conf.pre_crawled_paths = pre_crawled_path or []
-
-    hammertime = None
+
     accumulator = ResultAccumulator(output_manager=output_manager)
 
     output_manager.output_info('Starting Discovery on ' + conf.base_url)
@@ -274,45 +273,58 @@ def main(*, target_host, cookie_file, json_output, max_retry_count, plugin_setti
         plugin, value = option.split(':', 1)
         conf.plugin_settings[plugin].append(value)
 
+    loop = custom_event_loop()
+
+    async def async_main():
+        hammertime = None
+        try:
+            root_path = conf.path_template.copy()
+            root_path['url'] = '/'
+            database.valid_paths.append(root_path)
+            for crawled_path in conf.pre_crawled_paths:
+                new_path = conf.path_template.copy()
+                new_path['url'] = crawled_path
+                database.valid_paths.append(new_path)
+            load_target_paths()
+            load_target_files()
+            conf.cookies = loaders.load_cookie_file(cookie_file)
+            conf.user_agent = user_agent
+            conf.proxy_url = proxy
+            conf.forge_vhost = vhost
+
+            async with configure_hammertime(cookies=conf.cookies, proxy=conf.proxy_url, retry_count=max_retry_count,
+                                            user_agent=conf.user_agent, vhost=conf.forge_vhost,
+                                            confirmation_factor=confirmation_factor,
+                                            concurrency=concurrency,
+                                            har_output_dir=har_output_dir) as hammertime:
+                try:
+                    t = loop.create_task(stat_on_input(hammertime))
+                    await scan(hammertime, accumulator=accumulator,
+                               cookies=conf.cookies, directories_only=directories_only,
+                               files_only=files_only, plugins_only=plugins_only, depth_limit=depth_limit,
+                               recursive=recursive)
+                finally:
+                    t.cancel()
+                    textutils.output_info(format_stats(hammertime.stats))
+
+            output_manager.output_info('Scan completed')
+
+        except (KeyboardInterrupt, asyncio.CancelledError):
+            output_manager.output_error('Keyboard Interrupt Received')
+        except (OfflineHostException, StopRequest):
+            output_manager.output_error("Target host seems to be offline.")
+        except ImportError as e:
+            output_manager.output_error("Additional module is required for the requested options: %s" % e)
+        finally:
+            output_manager.flush()
+
     try:
-        root_path = conf.path_template.copy()
-        root_path['url'] = '/'
-        database.valid_paths.append(root_path)
-        for crawled_path in conf.pre_crawled_paths:
-            new_path = conf.path_template.copy()
-            new_path['url'] = crawled_path
-            database.valid_paths.append(new_path)
-        load_target_paths()
-        load_target_files()
-        conf.cookies = loaders.load_cookie_file(cookie_file)
-        conf.user_agent = user_agent
-        conf.proxy_url = proxy
-        conf.forge_vhost = vhost
-        loop = custom_event_loop()
-        hammertime = loop.run_until_complete(
-            configure_hammertime(cookies=conf.cookies, proxy=conf.proxy_url, retry_count=max_retry_count,
-                                 user_agent=conf.user_agent, vhost=conf.forge_vhost,
-                                 confirmation_factor=confirmation_factor,
-                                 concurrency=concurrency,
-                                 har_output_dir=har_output_dir))
-        t = loop.create_task(stat_on_input(hammertime))
-        loop.run_until_complete(scan(hammertime, accumulator=accumulator,
-                                     cookies=conf.cookies, directories_only=directories_only,
-                                     files_only=files_only, plugins_only=plugins_only, depth_limit=depth_limit,
-                                     recursive=recursive))
-        t.cancel()
-        output_manager.output_info('Scan completed')
+        main_task = loop.create_task(async_main())
 
+        loop.run_until_complete(main_task)
     except (KeyboardInterrupt, asyncio.CancelledError):
         output_manager.output_error('Keyboard Interrupt Received')
-    except (OfflineHostException, StopRequest):
-        output_manager.output_error("Target host seems to be offline.")
-    except ImportError as e:
-        output_manager.output_error("Additional module is required for the requested options: %s" % e)
     finally:
-        if hammertime is not None:
-            textutils.output_info(format_stats(hammertime.stats))
-
         output_manager.flush()
 
 
@@ -330,7 +342,7 @@ async def stat_on_input(hammertime):
     await loop.connect_read_pipe(lambda: reader_protocol, sys.stdin)
 
     expiry = datetime.now()
-    while True:
+    while not hammertime.is_closed:
         await reader.readline()
 
         # Throttle stats printing

tachyon/__version__.py

@@ -16,4 +16,4 @@
 # this program; if not, write to the Free Software Foundation, Inc., 59 Temple
 # Place, Suite 330, Boston, MA  02111-1307  USA
 
-__version__ = "3.5.1"
+__version__ = "3.5.4"

tachyon/config.py

@@ -19,6 +19,7 @@
 
 from aiohttp import ClientSession, TCPConnector
 from aiohttp.cookiejar import DummyCookieJar
+from contextlib import asynccontextmanager
 from hammertime import HammerTime
 from hammertime.config import custom_event_loop
 from hammertime.engine import AioHttpEngine
@@ -40,6 +41,7 @@
                      'Chrome/41.0.2228.0 Safari/537.36'
 
 
+@asynccontextmanager
 async def configure_hammertime(proxy=None, retry_count=3, cookies=None, concurrency=0, **kwargs):
     loop = custom_event_loop()
     engine = AioHttpEngine(loop=loop, verify_ssl=False, proxy=proxy)
@@ -55,12 +57,15 @@ async def configure_hammertime(proxy=None, retry_count=3, cookies=None, concurre
         scale_policy = StaticPolicy(concurrency)
 
     kb = KnowledgeBase()
-    hammertime = HammerTime(loop=loop, request_engine=engine, retry_count=retry_count, proxy=proxy, kb=kb,
-                            scale_policy=scale_policy)
-    setup_hammertime_heuristics(hammertime, **kwargs)
-    hammertime.collect_successful_requests()
-    hammertime.kb = kb
-    return hammertime
+    try:
+        hammertime = HammerTime(loop=loop, request_engine=engine, retry_count=retry_count, proxy=proxy, kb=kb,
+                                scale_policy=scale_policy)
+        setup_hammertime_heuristics(hammertime, **kwargs)
+        hammertime.collect_successful_requests()
+        hammertime.kb = kb
+        yield hammertime
+    finally:
+        await engine.session.close()
 
 
 def setup_hammertime_heuristics(hammertime, *,

tachyon/har.py

@@ -17,7 +17,7 @@ def __call__(self, har: HAR):
         filename = "%s.har" % uuid.uuid4()
         file_path = join(self.dir, filename)
         with open(file_path, "w") as fp:
-            fp.write(json.dumps(har.dump().data, indent=4))
+            fp.write(json.dumps(har.dump(), indent=4))
         return file_path
 
 

tachyon/output.py

@@ -67,10 +67,13 @@ def _get_current_time(self):
 class JSONOutput(OutputManager):
 
     def __init__(self):
+        self.flushed = False
         self.buffer = []
 
     def flush(self):
-        self.output_raw_message(json.dumps({"result": self.buffer, "version": __version__, "from": conf.name}))
+        if not self.flushed:
+            self.output_raw_message(json.dumps({"result": self.buffer, "version": __version__, "from": conf.name}))
+            self.flushed = True
 
     def _add_output(self, text, level, data=None):
         formatted = self._format_output(self._get_current_time(), logging.getLevelName(level), text, data)

test/config_test.py

@@ -57,7 +57,8 @@ async def test_configure_hammertime_add_user_agent_to_request_header(self):
 
         with patch("tachyon.config.SetHeader") as set_header:
             set_header.return_value = SetHeader("a", "b")
-            await config.configure_hammertime(user_agent=user_agent)
+            async with config.configure_hammertime(user_agent=user_agent):
+                pass
 
             set_header.assert_any_call("User-Agent", user_agent)
 
@@ -67,7 +68,8 @@ async def test_configure_hammertime_add_host_header_to_request_header(self):
 
         with patch("tachyon.config.SetHeader") as set_header:
             set_header.return_value = SetHeader("a", "b")
-            await config.configure_hammertime()
+            async with config.configure_hammertime():
+                pass
 
             set_header.assert_any_call("Host", conf.target_host)
 
@@ -78,7 +80,8 @@ async def test_configure_hammertime_use_user_supplied_vhost_for_host_header(self
 
         with patch("tachyon.config.SetHeader") as set_header:
             set_header.return_value = SetHeader("a", "b")
-            await config.configure_hammertime(vhost=forge_vhost)
+            async with config.configure_hammertime(vhost=forge_vhost):
+                pass
 
             set_header.assert_any_call("Host", forge_vhost)
 
@@ -88,7 +91,8 @@ async def test_configure_hammertime_allow_requests_to_user_supplied_vhost(self):
         forge_vhost = "vhost.example.com"
 
         with patch("tachyon.config.FilterRequestFromURL", MagicMock(return_value=FilterRequestFromURL)) as url_filter:
-            await config.configure_hammertime(vhost=forge_vhost)
+            async with config.configure_hammertime(vhost=forge_vhost):
+                pass
 
             _, kwargs = url_filter.call_args
             self.assertEqual(kwargs["allowed_urls"], ("vhost.example.com", "example.com"))
@@ -99,16 +103,16 @@ async def test_configure_hammertime_create_aiohttp_engine_for_hammertime(self, l
         engine.session.close = make_mocked_coro()
         EngineFactory = MagicMock(return_value=engine)
         with patch("tachyon.config.AioHttpEngine", EngineFactory):
-            hammertime = await config.configure_hammertime(proxy="my-proxy")
-
-            EngineFactory.assert_called_once_with(loop=loop, verify_ssl=False, proxy="my-proxy")
-            self.assertEqual(hammertime.request_engine.request_engine, engine)
+            async with config.configure_hammertime(proxy="my-proxy") as hammertime:
+                EngineFactory.assert_called_once_with(loop=loop, verify_ssl=False, proxy="my-proxy")
+                self.assertEqual(hammertime.request_engine.request_engine, engine)
 
     @async_test()
     async def test_configure_hammertime_create_client_session_with_dummy_cookie_jar_if_user_supply_cookies(self):
         cookies = "not none"
         with patch("tachyon.config.ClientSession") as SessionFactory:
-            await config.configure_hammertime(cookies=cookies)
+            async with config.configure_hammertime(cookies=cookies):
+                pass
 
             _, kwargs = SessionFactory.call_args
             self.assertTrue(isinstance(kwargs["cookie_jar"], DummyCookieJar))
@@ -117,7 +121,8 @@ async def test_configure_hammertime_create_client_session_with_dummy_cookie_jar_
     async def test_configure_hammertime_configure_aiohttp_to_resolve_host_only_once(self, loop):
         with patch("tachyon.config.TCPConnector", MagicMock(return_value=TCPConnector(loop=loop))) as \
                 ConnectorFactory:
-            await config.configure_hammertime()
+            async with config.configure_hammertime():
+                pass
 
             _, kwargs = ConnectorFactory.call_args
             self.assertTrue(kwargs["use_dns_cache"])