pentool 0.2.2__tar.gz → 0.2.3__tar.gz

{pentool-0.2.2 → pentool-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pentool
-Version: 0.2.2
+Version: 0.2.3
 Summary: Modern Web Pentesting TUI - Professional security testing tool
 Author-email: Pentool Team <dev@pentool.pro>
 License: AGPL-3.0

{pentool-0.2.2 → pentool-0.2.3}/pentool/__init__.py

@@ -1,6 +1,6 @@
 """Pentool — professional web security testing toolkit with Textual TUI."""
 
-__version__ = "0.2.2"
+__version__ = "0.2.3"
 __author__ = "pentool"
 
 

{pentool-0.2.2 → pentool-0.2.3}/pentool/core/database.py

@@ -95,6 +95,14 @@ CREATE TABLE IF NOT EXISTS site_map (
 );
 
 CREATE UNIQUE INDEX IF NOT EXISTS idx_site_map_host_path ON site_map (host, path);
+
+-- Generic per-project key/value settings (e.g. "proxy.enforce_scope").
+-- Lives in the project .db file itself so it travels with the project,
+-- unlike pentool's global ~/.config/pentool/config.yaml.
+CREATE TABLE IF NOT EXISTS project_settings (
+    key     TEXT PRIMARY KEY,
+    value   TEXT NOT NULL DEFAULT ''
+);
 """
 
 
@@ -170,3 +178,31 @@ def init_db_sync(db_path: str) -> None:
         db_path: Path to the SQLite file.
     """
     asyncio.run(init_db(db_path))
+
+
+async def get_project_setting(db_path: str, key: str, default: str | None = None) -> str | None:
+    """Read a single per-project key/value setting from `project_settings`.
+
+    Returns `default` if the key is missing or the table/column doesn't
+    exist yet (e.g. a brand-new .db before init_db() has run).
+    """
+    try:
+        async with aiosqlite.connect(db_path) as db:
+            async with db.execute(
+                "SELECT value FROM project_settings WHERE key = ?", (key,)
+            ) as cur:
+                row = await cur.fetchone()
+                return row[0] if row else default
+    except Exception:
+        return default
+
+
+async def set_project_setting(db_path: str, key: str, value: str) -> None:
+    """Upsert a single per-project key/value setting into `project_settings`."""
+    async with aiosqlite.connect(db_path) as db:
+        await db.execute(
+            "INSERT INTO project_settings (key, value) VALUES (?, ?) "
+            "ON CONFLICT(key) DO UPDATE SET value = excluded.value",
+            (key, value),
+        )
+        await db.commit()

{pentool-0.2.2 → pentool-0.2.3}/pentool/modules/proxy.py

@@ -150,6 +150,14 @@ class ProxyServer:
 
         self.intercept_enabled: bool = False
         self.scope: list[str] = []  # Empty = intercept everything
+        # When False (default), scope is informational only — ALL traffic is
+        # captured/shown regardless of self.scope. Scope filtering (dropping
+        # out-of-scope requests from capture) only takes effect when this is
+        # explicitly turned on via the "Skip out-of-scope" toggle button.
+        # Without this flag, a leftover/stale `scope` list (e.g. saved from a
+        # previous project) would silently blackhole all non-matching traffic
+        # the moment the proxy starts, with no visible indication why.
+        self.enforce_scope: bool = False
 
         # Match/Replace via dedicated engine
         self._match_replace_engine = MatchReplaceEngine()
@@ -205,7 +213,10 @@ class ProxyServer:
         self._running = False
         if self._server:
             self._server.close()
-            await self._server.wait_closed()
+            try:
+                await asyncio.wait_for(self._server.wait_closed(), timeout=2.0)
+            except asyncio.TimeoutError:
+                logger.debug("ProxyServer.stop: wait_closed timeout, continuing")
             self._server = None
         # Close singleton HTTP client
         if self._http_client:
@@ -214,6 +225,25 @@ class ProxyServer:
             except Exception as e:
                 logger.warning("ProxyServer.stop: http_client.close() error: %s", e)
             self._http_client = None
+        # Cancel all remaining active tasks (open TLS tunnels, keep-alive
+        # connections, intercept waits) so the event loop can exit promptly.
+        # Without this, _handle_connect loops on _READ_TIMEOUT=30s and
+        # _INTERCEPT_TIMEOUT=300s — the proxy thread would not exit for up to
+        # 5 minutes after stop(), blocking project switches and restarts.
+        loop = asyncio.get_running_loop()
+        current = asyncio.current_task()
+        pending = [t for t in asyncio.all_tasks(loop) if t is not current and not t.done()]
+        if pending:
+            logger.debug("ProxyServer.stop: cancelling %d active task(s)", len(pending))
+            for task in pending:
+                task.cancel()
+            try:
+                await asyncio.wait_for(
+                    asyncio.gather(*pending, return_exceptions=True),
+                    timeout=2.0,
+                )
+            except asyncio.TimeoutError:
+                logger.debug("ProxyServer.stop: task cancellation timed out")
         logger.info("Proxy stopped")
 
     async def serve_forever(self) -> None:
@@ -296,6 +326,19 @@ class ProxyServer:
         else:
             self.intercept_enabled = enabled
 
+    def set_enforce_scope(self, enabled: bool) -> None:
+        """Thread-safe setting of enforce_scope flag from any thread.
+
+        Same rationale as set_intercept — the proxy loop reads this flag
+        in _handle_http/_handle_connect, so writes from the TUI thread must
+        go through call_soon_threadsafe to avoid a torn read.
+        """
+        loop = self._loop
+        if loop is not None and loop.is_running():
+            loop.call_soon_threadsafe(setattr, self, "enforce_scope", enabled)
+        else:
+            self.enforce_scope = enabled
+
     def _find_request(self, req_id: str) -> InterceptedRequest | None:
         with self._requests_lock:
             for r in reversed(self.requests):
@@ -367,7 +410,8 @@ class ProxyServer:
         await writer.drain()
 
         # If host is out of scope — just tunnel without interception
-        if not self.is_in_scope(domain):
+        # (only enforced when the user turned on "Skip out-of-scope")
+        if self.enforce_scope and not self.is_in_scope(domain):
             await self._tunnel_raw(domain, host_port, reader, writer)
             return
 
@@ -491,7 +535,8 @@ class ProxyServer:
                     req.method, req.url, is_https, self.intercept_enabled, self.is_in_scope(host))
 
         # Check scope (for HTTP; HTTPS already checked in _handle_connect)
-        if not is_https and not self.is_in_scope(host):
+        # (only enforced when the user turned on "Skip out-of-scope")
+        if self.enforce_scope and not is_https and not self.is_in_scope(host):
             logger.debug("PROXY: _handle_http: host %s out of scope, forwarding direct", host)
             return await self._forward_direct(req, writer)
 

{pentool-0.2.2 → pentool-0.2.3}/pentool/tui/app.py

@@ -701,6 +701,14 @@ class PentoolApp(App):
                 pass
         self.query_one(ContentSwitcher).current = f"screen-{module_id}"
         self._active_module = module_id
+        # _update_proxy_screen_labels() only refreshes the ProxyScreen label
+        # when it IS the active module — so if the proxy was stopped/started
+        # while the user was on a different tab (e.g. right after creating a
+        # new project, which force-stops the proxy while Dashboard is shown),
+        # the label update was skipped and stayed stale until the next actual
+        # start/stop toggle. Refresh it explicitly on every switch into Proxy.
+        if module_id == "proxy":
+            self._update_proxy_screen_labels()
 
     def get_proxy(self) -> ProxyServer | None:
         return self._proxy
@@ -758,6 +766,21 @@ class PentoolApp(App):
         if self._proxy.is_running:
             self._stop_proxy()
         else:
+            if not self._project_loaded:
+                # Project DB switch/open (auto-open at startup, New/Open
+                # Project) is still finishing in the background — starting
+                # the proxy now would race HttpStorage.switch_db() (its
+                # connection may be mid-close/reopen) and silently lose or
+                # fail to persist the first captured requests. _project_loaded
+                # is set as soon as the DB switch itself completes (see
+                # ProjectManager._do_switch) — the other screens may still be
+                # reloading, but that's independent of Proxy.
+                self.notify(
+                    "Проект ещё открывается — подождите пару секунд перед запуском Proxy",
+                    severity="warning",
+                    timeout=4,
+                )
+                return
             self._start_proxy()
 
     def action_toggle_intercept(self) -> None:
@@ -877,13 +900,25 @@ class PentoolApp(App):
                 self._proxy.stop(), self._proxy_loop
             )
             try:
-                future.result(timeout=5)
+                # stop() now cancels all active tasks and waits up to ~4s
+                # internally, so 6s here is generous headroom
+                future.result(timeout=6)
             except Exception as e:
                 logger.warning("APP: proxy.stop() error or timeout: %s", e)
+                # Force-cancel anything still running in the proxy loop so
+                # the thread can exit even if stop() itself timed out
+                if self._proxy_loop and not self._proxy_loop.is_closed():
+                    try:
+                        def _cancel_all():
+                            for t in asyncio.all_tasks(self._proxy_loop):
+                                t.cancel()
+                        self._proxy_loop.call_soon_threadsafe(_cancel_all)
+                    except Exception:
+                        pass
         if self._proxy_thread and self._proxy_thread.is_alive():
-            self._proxy_thread.join(timeout=3)
+            self._proxy_thread.join(timeout=5)
             if self._proxy_thread.is_alive():
-                logger.warning("APP: proxy thread did not stop in 3s")
+                logger.warning("APP: proxy thread did not stop in 5s — port 8080 may still be in use")
         self.call_after_refresh(self._update_status)
         self.call_after_refresh(self._update_proxy_screen_labels)
 
@@ -974,7 +1009,14 @@ class PentoolApp(App):
             target = self.query_one(SCREEN_TARGET, TargetScreen)
             api = target._get_api()
             api.sitemap.set_in_scope(msg.host, msg.in_scope)
-            target._load_sitemap()
+            # Rebuild the tree from the in-memory sitemap only — NOT
+            # target._load_sitemap(), which does a full api.load() from
+            # the DB. The sitemap is persisted in batches of 20 rows
+            # (see target/screen.py _do_save_sitemap), so any host/path
+            # added in-memory since the last batch save would be wiped
+            # out by that reload, making a simple "Add to Scope" click
+            # from Proxy appear to clear the whole Target tree.
+            target._refresh_tree()
         except Exception as e:
             logger.debug("on_sync_scope_to_target: %s", e)
 

{pentool-0.2.2 → pentool-0.2.3}/pentool/tui/dialogs/match_replace_dialog.py

@@ -13,7 +13,6 @@ from textual.screen import ModalScreen
 _CSS = (Path(__file__).parent / "match_replace_dialog.tcss").read_text(encoding="utf-8")
 from textual.widgets import (
     Button,
-    Checkbox,
     DataTable,
     Input,
     Label,
@@ -21,6 +20,7 @@ from textual.widgets import (
 )
 
 from pentool.api.proxy_api import MatchReplaceRule
+from pentool.tui.widgets.nice_checkbox import NiceCheckbox as Checkbox
 from pentool.tui.widgets.toolbar_button import ToolbarButton
 
 _TARGET_OPTIONS = [("request", "Request"), ("response", "Response"), ("both", "Both")]

{pentool-0.2.2 → pentool-0.2.3}/pentool/tui/project_manager.py

@@ -325,9 +325,14 @@ class ProjectManager:
             )
             self._app._stop_proxy()
 
+        # Gate Start Proxy (action_toggle_proxy) until _do_switch flips this
+        # back to True once HttpStorage.switch_db() has actually completed —
+        # otherwise a Start Proxy click during the switch could race the
+        # storage connection being closed/reopened underneath it.
+        self._app._project_loaded = False
+
         self._cfg.db_path = path
         self._app._project_path = path
-        # Note: _project_loaded will be set AFTER _reload_project_screens in _do_switch
 
         if self._proxy:
             self._proxy.db_path = path
@@ -337,6 +342,10 @@ class ProjectManager:
                 self._proxy.clear_requests()
                 self._proxy.scope = []
                 self._proxy.match_replace_rules = []
+                # New project's DB has no project_settings row yet — reset to
+                # the documented default (OFF) rather than carrying over
+                # whatever the previous project had in memory.
+                self._proxy.set_enforce_scope(False)
             except Exception:
                 pass
 
@@ -371,12 +380,44 @@ class ProjectManager:
         """Single entry point for all project switches.
 
         Runs in an exclusive worker — no two switches can overlap.
-        Order is strict:
+        Order:
+          0. Wait for the proxy thread to fully exit (if _stop_proxy was
+             called but the thread didn't die within join(timeout) — the
+             old loop may still hold the SQLite WAL lock or the TCP port).
           1. init_db (CREATE TABLE IF NOT EXISTS, migrations)
           2. switch_db on ProxyService (close old connection, open new one)
           3. clear in-memory proxy state if new project
-          4. reload all screens sequentially (no nested run_worker)
+          4. set _project_loaded — DB is ready, Start Proxy is now safe,
+             even though the other screens (Repeater/Scanner/Target/Dashboard)
+             may still be reloading in the background (step 5). Proxy only
+             depends on the HttpStorage connection from step 2, not on any
+             of those screens, so there is no reason to make the user wait
+             for all of them before allowing Start Proxy.
+          5. reload the remaining screens — run concurrently (they read
+             independent tables/state), not sequentially.
         """
+        # 0. Wait for the proxy thread to die before touching the DB.
+        # _stop_proxy() already called join(timeout=5), but if the thread
+        # is still alive (e.g. a 30-second _READ_TIMEOUT blocked task didn't
+        # get cancelled in time) we must not open a new SQLite connection
+        # while the old proxy event loop might still be writing to it.
+        proxy_thread = self._app._proxy_thread
+        if proxy_thread is not None and proxy_thread.is_alive():
+            logger.info("_do_switch: waiting for proxy thread to exit before DB switch…")
+            loop = asyncio.get_running_loop()
+            # Poll in the async event loop so we don't block the TUI thread
+            for _ in range(50):  # up to 5 seconds in 100ms steps
+                alive = await loop.run_in_executor(
+                    None, lambda: proxy_thread.is_alive()
+                )
+                if not alive:
+                    break
+                await asyncio.sleep(0.1)
+            if proxy_thread.is_alive():
+                logger.warning("_do_switch: proxy thread still alive after extra 5s — proceeding anyway")
+            else:
+                logger.info("_do_switch: proxy thread exited, proceeding with DB switch")
+
         # 1. Ensure schema exists (safe for both new and existing DBs)
         await self._init_new_db(path)
 
@@ -405,75 +446,91 @@ class ProjectManager:
             except Exception as exc:
                 logger.debug("_do_switch: clear for new project: %s", exc)
 
-        # 4. Reload all screens — all awaited, never spawning sub-workers
-        await self._reload_project_screens(path, is_new=is_new)
-
-        # 5. Set flag AFTER all screens are reloaded (fixes race condition)
+        # 4. Set flag as soon as the DB itself is ready — Start Proxy only
+        # needs HttpStorage (step 2) done, not the other screens below.
         self._app._project_loaded = True
-        logger.info("_do_switch: project loaded flag set")
+        logger.info("_do_switch: project loaded flag set (DB ready)")
+
+        # 5. Reload the remaining screens concurrently — independent of
+        # each other and of the flag above.
+        await self._reload_project_screens(path, is_new=is_new)
+        logger.info("_do_switch: all screens reloaded")
 
     async def _reload_project_screens(self, path: str, is_new: bool = False) -> None:
-        """Reload data from DB into all screens. Called from _do_switch only."""
+        """Reload data from DB into all screens. Called from _do_switch only.
 
-        # 0. RepeaterScreen — сброс или загрузка вкладок
-        try:
-            from pentool.tui.constants import SCREEN_REPEATER
-            from pentool.tui.screens.repeater.screen import RepeaterScreen
-            repeater = self._app.query_one(SCREEN_REPEATER, RepeaterScreen)
-            if is_new:
-                await repeater.reset_for_new_project()
-            else:
-                await repeater.reload_from_project(path)
-            logger.info("_reload_project_screens: repeater %s",
-                        "reset" if is_new else "reloaded")
-        except Exception as exc:
-            logger.debug("_reload_project_screens repeater: %s", exc)
+        Each screen's reload is independent (different tables/state), so they
+        run concurrently via asyncio.gather instead of one after another —
+        total time is bounded by the slowest single reload, not the sum.
+        """
 
-        # 1. ProxyScreen — must come after switch_db is fully done
-        try:
-            screen = self._app.query_one(SCREEN_PROXY, ProxyScreen)
-            await screen._reload_from_storage()
-            logger.info("_reload_project_screens: proxy reloaded from %s", path)
-        except Exception as exc:
-            logger.debug("_reload_project_screens proxy: %s", exc)
+        async def _reload_repeater() -> None:
+            try:
+                from pentool.tui.constants import SCREEN_REPEATER
+                from pentool.tui.screens.repeater.screen import RepeaterScreen
+                repeater = self._app.query_one(SCREEN_REPEATER, RepeaterScreen)
+                if is_new:
+                    await repeater.reset_for_new_project()
+                else:
+                    await repeater.reload_from_project(path)
+                logger.info("_reload_project_screens: repeater %s",
+                            "reset" if is_new else "reloaded")
+            except Exception as exc:
+                logger.debug("_reload_project_screens repeater: %s", exc)
 
-        # 2. ScannerScreen
-        try:
-            from pentool.tui.screens.scanner.screen import ScannerScreen
-            scanner_screen = self._app.query_one(SCREEN_SCANNER, ScannerScreen)
-            if scanner_screen._scanner_api is not None:
-                scanner_screen._scanner_api._db_path = path
-            scanner_screen._scanner_api = None
-            scanner_screen._scanner_api = scanner_screen._get_or_create_api(path)
-            scanner_screen._populate_from_db([])
-            scanner_screen._load_findings_worker()
-            logger.info("_reload_project_screens: scanner reloaded")
-        except Exception as exc:
-            logger.debug("_reload_project_screens scanner: %s", exc)
+        async def _reload_proxy() -> None:
+            try:
+                screen = self._app.query_one(SCREEN_PROXY, ProxyScreen)
+                await screen._reload_from_storage()
+                logger.info("_reload_project_screens: proxy reloaded from %s", path)
+            except Exception as exc:
+                logger.debug("_reload_project_screens proxy: %s", exc)
 
-        # 3. TargetScreen
-        try:
-            from pentool.tui.screens.target.screen import TargetScreen
-            target_screen = self._app.query_one(SCREEN_TARGET, TargetScreen)
-            if target_screen._target_api is not None:
-                try:
-                    await target_screen._target_api.save()
-                except Exception:
-                    pass
-            target_screen._target_api = None
-            target_screen._get_api()
-            target_screen._load_sitemap()
-            logger.info("_reload_project_screens: target reloaded from %s", path)
-        except Exception as exc:
-            logger.debug("_reload_project_screens target: %s", exc)
+        async def _reload_scanner() -> None:
+            try:
+                from pentool.tui.screens.scanner.screen import ScannerScreen
+                scanner_screen = self._app.query_one(SCREEN_SCANNER, ScannerScreen)
+                if scanner_screen._scanner_api is not None:
+                    scanner_screen._scanner_api._db_path = path
+                scanner_screen._scanner_api = None
+                scanner_screen._scanner_api = scanner_screen._get_or_create_api(path)
+                scanner_screen._populate_from_db([])
+                scanner_screen._load_findings_worker()
+                logger.info("_reload_project_screens: scanner reloaded")
+            except Exception as exc:
+                logger.debug("_reload_project_screens scanner: %s", exc)
 
-        # 4. Dashboard
-        try:
-            from pentool.tui.screens.dashboard.screen import DashboardScreen
-            dash = self._app.query_one(SCREEN_DASHBOARD, DashboardScreen)
-            dash.refresh_stats()
-        except Exception as exc:
-            logger.debug("_reload_project_screens dashboard: %s", exc)
+        async def _reload_target() -> None:
+            try:
+                from pentool.tui.screens.target.screen import TargetScreen
+                target_screen = self._app.query_one(SCREEN_TARGET, TargetScreen)
+                if target_screen._target_api is not None:
+                    try:
+                        await target_screen._target_api.save()
+                    except Exception:
+                        pass
+                target_screen._target_api = None
+                target_screen._get_api()
+                target_screen._load_sitemap()
+                logger.info("_reload_project_screens: target reloaded from %s", path)
+            except Exception as exc:
+                logger.debug("_reload_project_screens target: %s", exc)
+
+        async def _reload_dashboard() -> None:
+            try:
+                from pentool.tui.screens.dashboard.screen import DashboardScreen
+                dash = self._app.query_one(SCREEN_DASHBOARD, DashboardScreen)
+                dash.refresh_stats()
+            except Exception as exc:
+                logger.debug("_reload_project_screens dashboard: %s", exc)
+
+        await asyncio.gather(
+            _reload_repeater(),
+            _reload_proxy(),
+            _reload_scanner(),
+            _reload_target(),
+            _reload_dashboard(),
+        )
 
     async def _init_new_db(self, path: str) -> None:
         try:

{pentool-0.2.2 → pentool-0.2.3}/pentool/tui/screens/intruder/screen.py

@@ -20,7 +20,6 @@ _CSS = (Path(__file__).parent / "screen.tcss").read_text(encoding="utf-8")
 from textual.message import Message as _Message
 from textual.widgets import (
     Button,
-    Checkbox,
     DataTable,
     Input,
     Label,
@@ -45,6 +44,7 @@ from pentool.core.logging import get_logger
 from pentool.tui.messages import SendToRepeater
 from pentool.tui.mixins.app_mixin import AppMixin
 from pentool.tui.mixins.request_context_menu import RequestContextMenuMixin
+from pentool.tui.widgets.nice_checkbox import NiceCheckbox as Checkbox
 from pentool.tui.widgets.request_editor import HttpView, _load_into_textarea
 from pentool.tui.widgets.resize_handle import ResizeHandle
 from pentool.tui.widgets.toolbar_button import ToolbarButton