From c05cdb87a053e3f9a5fc3c1bf754a86cc2a62543 Mon Sep 17 00:00:00 2001 From: Psychevus Date: Sat, 26 Jul 2025 19:19:03 +0330 Subject: [PATCH] feat: add realtime websocket and UI skeleton --- README.md | 20 +++++++++++++++++ main.py | 28 +++++++++++++++++++++++- pyproject.toml | 1 + requirements.txt | 1 + ui/babel.config.json | 3 +++ ui/index.html | 12 ++++++++++ ui/package.json | 25 +++++++++++++++++++++ ui/src/App.js | 36 ++++++++++++++++++++++++++++++ ui/src/index.js | 6 +++++ ui/webpack.config.js | 27 +++++++++++++++++++++++ utils/realtime.py | 52 ++++++++++++++++++++++++++++++++++++++++++++ 11 files changed, 210 insertions(+), 1 deletion(-) create mode 100644 ui/babel.config.json create mode 100644 ui/index.html create mode 100644 ui/package.json create mode 100644 ui/src/App.js create mode 100644 ui/src/index.js create mode 100644 ui/webpack.config.js create mode 100644 utils/realtime.py diff --git a/README.md b/README.md index 007b77e..fbf0481 100644 --- a/README.md +++ b/README.md @@ -141,6 +141,26 @@ Results are stored in `output/` with timestamps. Notifications can be sent via --- +## Real-time dashboard + +Start the CLI with a WebSocket port to stream live results: + +```bash +python main.py example.com --pipeline --ws-port 8765 +``` + +Then install and launch the React interface: + +```bash +cd ui && npm install +npm start +``` + +Open to view the dashboard with a live terminal feed +and 3D recon graph. + +--- + ## Plugins Drop Python files into a `plugins/` directory or set the diff --git a/main.py b/main.py index 29727ac..9461e6a 100644 --- a/main.py +++ b/main.py @@ -37,6 +37,7 @@ from utils.deps import DependencyError, check_dependencies from utils.plugins import load_plugins from distributed import DistributedOrchestrator +from utils import realtime load_plugins() @@ -110,6 +111,8 @@ def _run_host_tool(tool_cls: type[Module], host_list: List[str], name: str) -> L logger.error("%s error: %s", name, exc) continue local.extend(res) + for f in res: + realtime.publish(f.asdict()) if name == "nmap": xml_path = res[0].data.get("xml") if xml_path: @@ -149,6 +152,8 @@ def _run_host_tool(tool_cls: type[Module], host_list: List[str], name: str) -> L try: res = fut.result(timeout=getattr(cls, "TIMEOUT", 60)) findings.extend(res) + for f in res: + realtime.publish(f.asdict()) except Exception as exc: # noqa: BLE001 logger.error("%s error: %s", name, exc) if show_summary: @@ -166,11 +171,15 @@ def _run_host_tool(tool_cls: type[Module], host_list: List[str], name: str) -> L if use_pipeline and tool_name == "subfinder": res = tool_cls().run(target) findings.extend(res) + for f in res: + realtime.publish(f.asdict()) hosts = [f.data.get("host") for f in res if f.data.get("host")] summary["subdomains"] += len(res) elif use_pipeline and tool_name == "httpx": res = tool_cls().run("\n".join(hosts)) findings.extend(res) + for f in res: + realtime.publish(f.asdict()) urls = [f.data.get("url") for f in res if f.data.get("url")] hosts = [ f.data.get("host") or f.data.get("ip") @@ -185,6 +194,8 @@ def _run_host_tool(tool_cls: type[Module], host_list: List[str], name: str) -> L continue res = tool_cls().run("\n".join(urls)) findings.extend(res) + for f in res: + realtime.publish(f.asdict()) summary["nuclei"] += len(res) elif use_pipeline and tool_name in {"nmap", "testssl"}: if prev_state is not None and not new_urls: @@ -204,12 +215,16 @@ def _run_host_tool(tool_cls: type[Module], host_list: List[str], name: str) -> L for fut, name in futs.items(): res = fut.result() findings.extend(res) + for f in res: + realtime.publish(f.asdict()) continue else: processed.add(tool_name) for h in hosts: res = tool_cls().run(h) findings.extend(res) + for f in res: + realtime.publish(f.asdict()) if tool_name == "nmap": xml_path = res[0].data.get("xml") if xml_path: @@ -219,7 +234,10 @@ def _run_host_tool(tool_cls: type[Module], host_list: List[str], name: str) -> L if json_path: summary["ssl_issues"] += _count_testssl_issues(json_path) else: - findings.extend(tool_cls().run(target)) + res = tool_cls().run(target) + findings.extend(res) + for f in res: + realtime.publish(f.asdict()) except Exception as exc: # noqa: BLE001 logger.error("%s error: %s", tool_name, exc) @@ -359,8 +377,16 @@ def cli() -> None: "--runners", help="Comma-separated list of runner URLs (default: $PENTEST_TOOLKIT_RUNNERS)", ) + parser.add_argument( + "--ws-port", + type=int, + help="Start a WebSocket server on the given port for live results", + ) args = parser.parse_args() + if args.ws_port: + realtime.start(args.ws_port) + runners: List[str] = [] if args.runners: runners = [r.strip() for r in args.runners.split(',') if r.strip()] diff --git a/pyproject.toml b/pyproject.toml index 1d05500..5dc2462 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,7 @@ dependencies = [ "rich>=13", "weasyprint>=61", "requests>=2", + "websockets>=12", ] [project.optional-dependencies] diff --git a/requirements.txt b/requirements.txt index 3d22ae0..fbf0f1c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ pytest>=8 weasyprint>=61 requests>=2 markdown>=3 +websockets>=12 diff --git a/ui/babel.config.json b/ui/babel.config.json new file mode 100644 index 0000000..2b7bafa --- /dev/null +++ b/ui/babel.config.json @@ -0,0 +1,3 @@ +{ + "presets": ["@babel/preset-env", "@babel/preset-react"] +} diff --git a/ui/index.html b/ui/index.html new file mode 100644 index 0000000..f5dc359 --- /dev/null +++ b/ui/index.html @@ -0,0 +1,12 @@ + + + + + Pentest Toolkit Dashboard + + + +
+ + + diff --git a/ui/package.json b/ui/package.json new file mode 100644 index 0000000..d6a7d57 --- /dev/null +++ b/ui/package.json @@ -0,0 +1,25 @@ +{ + "name": "pentest-toolkit-ui", + "version": "0.1.0", + "private": true, + "scripts": { + "start": "webpack serve --mode development", + "build": "webpack --mode production" + }, + "dependencies": { + "react": "^18.3.0", + "react-dom": "^18.3.0", + "react-force-graph": "^1.47.0", + "three": "^0.163.0", + "react-three-fiber": "^9.0.0" + }, + "devDependencies": { + "@babel/core": "^7.24.5", + "@babel/preset-env": "^7.24.5", + "@babel/preset-react": "^7.24.5", + "babel-loader": "^9.1.3", + "webpack": "^5.91.0", + "webpack-cli": "^5.1.4", + "webpack-dev-server": "^5.0.4" + } +} diff --git a/ui/src/App.js b/ui/src/App.js new file mode 100644 index 0000000..a27ac02 --- /dev/null +++ b/ui/src/App.js @@ -0,0 +1,36 @@ +import React, { useEffect, useState } from 'react'; +import ForceGraph3D from 'react-force-graph'; + +export default function App() { + const [graph, setGraph] = useState({ nodes: [], links: [] }); + const [feed, setFeed] = useState([]); + + useEffect(() => { + const ws = new WebSocket('ws://localhost:8765'); + ws.onmessage = e => { + const msg = JSON.parse(e.data); + setFeed(f => [...f, msg]); + if (msg.type === 'subdomain') { + setGraph(g => ({ ...g, nodes: [...g.nodes, { id: msg.host }] })); + } + if (msg.type === 'httpx') { + setGraph(g => ({ + nodes: g.nodes, + links: [...g.links, { source: msg.host, target: msg.url }] + })); + } + }; + return () => ws.close(); + }, []); + + return ( +
+
+ {feed.map((f, i) =>
{JSON.stringify(f)}
)} +
+
+ +
+
+ ); +} diff --git a/ui/src/index.js b/ui/src/index.js new file mode 100644 index 0000000..04df71c --- /dev/null +++ b/ui/src/index.js @@ -0,0 +1,6 @@ +import React from 'react'; +import { createRoot } from 'react-dom/client'; +import App from './App'; + +const root = createRoot(document.getElementById('root')); +root.render(); diff --git a/ui/webpack.config.js b/ui/webpack.config.js new file mode 100644 index 0000000..db3ddbd --- /dev/null +++ b/ui/webpack.config.js @@ -0,0 +1,27 @@ +const path = require('path'); + +module.exports = { + entry: './src/index.js', + output: { + filename: 'bundle.js', + path: path.resolve(__dirname, 'dist'), + }, + module: { + rules: [ + { + test: /\.jsx?$/, + exclude: /node_modules/, + use: { + loader: 'babel-loader' + } + } + ] + }, + resolve: { + extensions: ['.js', '.jsx'] + }, + devServer: { + static: './', + port: 3000 + } +}; diff --git a/utils/realtime.py b/utils/realtime.py new file mode 100644 index 0000000..5cfcfc9 --- /dev/null +++ b/utils/realtime.py @@ -0,0 +1,52 @@ +import asyncio +import json +import threading +from typing import Any, Dict, Set + +import websockets +from websockets.server import WebSocketServerProtocol + +clients: Set[WebSocketServerProtocol] = set() +loop: asyncio.AbstractEventLoop | None = None +started = False + +async def _handler(ws: WebSocketServerProtocol): + clients.add(ws) + try: + async for _ in ws: + pass + finally: + clients.discard(ws) + +async def _broadcast(msg: Dict[str, Any]) -> None: + dead = [] + for ws in clients: + try: + await ws.send(json.dumps(msg)) + except Exception: + dead.append(ws) + for ws in dead: + clients.discard(ws) + + +def start(port: int = 8765) -> None: + """Launch the WebSocket server in a background thread.""" + global loop, started + if started: + return + loop = asyncio.new_event_loop() + started = True + + def _run() -> None: + asyncio.set_event_loop(loop) + server = websockets.serve(_handler, "0.0.0.0", port) + loop.run_until_complete(server) + loop.run_forever() + + threading.Thread(target=_run, daemon=True).start() + + +def publish(msg: Dict[str, Any]) -> None: + """Queue a message to all connected clients.""" + if started and loop: + loop.call_soon_threadsafe(asyncio.create_task, _broadcast(msg))