diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index b120bc07..3c483d1c 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -2,6 +2,9 @@ name: Flake8 on: [push] +permissions: + contents: read + jobs: build: runs-on: ubuntu-latest diff --git a/README.md b/README.md index b9dae779..3305805b 100755 --- a/README.md +++ b/README.md @@ -42,6 +42,7 @@ ### Recent improvements - The CLI now allows `--version` and `--update` to run without requiring a URL. - The crawler now resolves relative links against the current page and skips unsupported or malformed hrefs more gracefully. +- The optional TorBotApp desktop UI can be launched from the CLI with `torbot app` when installed as a sibling checkout. ...(will be updated) @@ -78,6 +79,9 @@ python main.py -u https://example.com --depth 2 --save json # Disable SOCKS5 if you are not using Tor locally python main.py -u https://example.com --disable-socks5 --visualize tree + +# Launch the optional desktop app when TorBotApp is installed +torbot app ``` ### Options @@ -93,6 +97,8 @@ optional arguments: -v Displays DEBUG level logging, default is INFO --version Show the current version of TorBot. --update Update TorBot to the latest stable version + --app Run optional TorBot desktop app + --app-dir APP_DIR Path to a TorBotApp checkout -q, --quiet Prevents display of header and IP address --save FORMAT Save results in a file. (tree, JSON) --visualize FORMAT Visualizes tree of data gathered. (tree, JSON, table) @@ -106,6 +112,36 @@ If you are using Tor locally, keep the SOCKS5 proxy enabled. If you are not usin Read more about torrc here: [Torrc](https://github.com/DedSecInside/TorBoT/blob/master/Tor.md) +## Optional desktop app + +TorBot remains a CLI-first Python tool. The desktop UI lives in the separate +[TorBotApp](https://github.com/KingAkeem/TorBotApp) Electron project so Node and +Electron are not required for CLI users. + +To launch it from this CLI, install TorBotApp as a sibling checkout: + +```text +code/ ++-- TorBot/ ++-- TorBotApp/ +``` + +Then run: + +```sh +torbot app +``` + +If TorBotApp is somewhere else, use either: + +```sh +TORBOT_APP_DIR=/path/to/TorBotApp torbot app +torbot --app --app-dir /path/to/TorBotApp +``` + +See [Desktop App Architecture](docs/DESKTOP_APP.md) for how the CLI, optional UI, +and backend services fit together. + ## Curated Features - [x] Visualization Module Revamp - [x] Implement BFS Search for webcrawler diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 68fc45a4..353936d7 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -2,6 +2,37 @@ -------------------- All notable changes to this project will be documented in this file. +## 4.3.0 - 2026-07-28 + +### Changed +- Made NLP classification train lazily from the packaged CSV and cache the classifier for the process, avoiding model retraining for every crawled page. +- Changed the NLP score returned by `classify()` from a misleading random-split accuracy value to prediction confidence for the selected category. + +### Fixed +- Removed runtime `training_data/` generation from page classification so installed packages do not try to write generated files into `site-packages`. +- Removed process-wide directory changes from the NLP helpers. + +### Tests +- Added direct NLP coverage for HTML text extraction, CSV loading, classifier scoring, and path-safe training-data export. + +## 4.2.0 - 2026-07-28 + +Install from PyPI after publishing with: + +```sh +pip install torbot==4.2.0 +``` + +### Added +- Added optional `torbot app` and `torbot --app` CLI commands for launching the TorBotApp desktop UI when installed separately. +- Added TorBotApp discovery through `--app-dir`, `TORBOT_APP_DIR`, or a sibling `TorBotApp` checkout. +- Documented the optional desktop app architecture while keeping Node and Electron out of the Python package dependencies. +- Added a safer publish helper that checks current-version artifacts and prompts for a PyPI token only when uploading. + +### Fixed +- Moved the CLI implementation into the installed package so the `torbot` console command works from a wheel install. +- Removed the deprecated `sklearn` package dependency in favor of the existing `scikit-learn` dependency so fresh installs do not fail. + ## 4.1.1 ### Fixed diff --git a/docs/DESKTOP_APP.md b/docs/DESKTOP_APP.md new file mode 100644 index 00000000..0efd222e --- /dev/null +++ b/docs/DESKTOP_APP.md @@ -0,0 +1,41 @@ +# Desktop App Architecture + +TorBot is CLI-first. The optional desktop experience is provided by +[TorBotApp](https://github.com/KingAkeem/TorBotApp), a separate Electron and +React project. + +Keeping the app separate preserves a lightweight Python install for CLI users: + +- TorBot owns the Python package, command-line crawler, and library modules. +- TorBotApp owns the desktop UI, Electron packaging, and Node dependency tree. +- Backend crawl services expose the API consumed by TorBotApp. + +## Launching from TorBot + +The CLI can launch TorBotApp when it is installed locally: + +```sh +torbot app +``` + +Discovery order: + +1. `--app-dir /path/to/TorBotApp` +2. `TORBOT_APP_DIR=/path/to/TorBotApp` +3. A sibling `TorBotApp` directory next to this repository + +If the app is not found, TorBot prints setup guidance and exits without changing +normal CLI behavior. + +## Expected development layout + +```text +code/ ++-- TorBot/ ++-- TorBotApp/ ++-- gotor/ +``` + +TorBotApp currently talks to the GoTor job-control API. If this Python package +later grows a compatible `torbot serve` API, TorBotApp can support either +backend without making Electron a required dependency of the CLI. diff --git a/main.py b/main.py index 7b8029a9..b576844e 100755 --- a/main.py +++ b/main.py @@ -1,192 +1,12 @@ #!/usr/bin/env python3 -import os -import argparse -import sys -import logging -import toml -import httpx +from torbot.cli import get_version, main, run, set_arguments -from torbot.modules.api import get_ip -from torbot.modules.color import color -from torbot.modules.updater import check_version -from torbot.modules.info import execute_all, fetch_html -from torbot.modules.linktree import LinkTree - - -def print_tor_ip_address(client: httpx.Client) -> None: - """ - https://check.torproject.org/ tells you if you are using tor and it - displays your IP address which we scape and display - """ - resp = get_ip(client) - print(resp["header"]) - print(color(resp["body"], "yellow")) - - -def print_header(version: str) -> None: - """ - Prints the TorBot banner including version and license. - """ - license_msg = color("LICENSE: GNU Public License v3", "red") - banner = r""" - __ ____ ____ __ ______ - / /_/ __ \/ __ \/ /_ ____/_ __/ - / __/ / / / /_/ / __ \/ __ \/ / - / /_/ /_/ / _, _/ /_/ / /_/ / / - \__/\____/_/ |_/_____/\____/_/ v{VERSION} - """.format(VERSION=version) - banner = color(banner, "red") - - title = r""" - {banner} - ####################################################### - # TorBot - Dark Web OSINT Tool # - # GitHub : https://github.com/DedsecInside/TorBot # - # Help : use -h for help text # - ####################################################### - {license_msg} - """ - - title = title.format(license_msg=license_msg, banner=banner) - print(title) - - -def run(arg_parser: argparse.ArgumentParser, version: str) -> None: - args = arg_parser.parse_args() - - # setup logging - date_fmt = "%d-%b-%y %H:%M:%S" - logging_fmt = "%(asctime)s - %(levelname)s - %(message)s" - logging_lvl = logging.DEBUG if args.v else logging.INFO - logging.basicConfig(level=logging_lvl, format=logging_fmt, datefmt=date_fmt) - - # Print verison then exit - if args.version: - print(f"TorBot Version: {version}") - sys.exit() - - # check version and update if necessary - if args.update: - check_version() - sys.exit() - - # URL is required for crawl-related actions - if not args.url: - arg_parser.print_help() - sys.exit() - - socks5_host = args.host - socks5_port = str(args.port) - socks5_proxy = f"socks5://{socks5_host}:{socks5_port}" - with httpx.Client( - timeout=60, proxies=socks5_proxy if not args.disable_socks5 else None - ) as client: - # print header and IP address if not set to quiet - if not args.quiet: - print_header(version) - print_tor_ip_address(client) - - if args.info: - execute_all(client, args.url) - - tree = LinkTree(url=args.url, depth=args.depth, client=client) - tree.load() - - # save data if desired - if args.save == "tree": - tree.save() - elif args.save == "json": - tree.saveJSON() - - if args.html == "display": - fetch_html(client, args.url, tree) - elif args.html == "save": - fetch_html(client, args.url, tree, save_html=True) - - # always print something, table is the default - if args.visualize == "table" or not args.visualize: - tree.showTable() - elif args.visualize == "tree": - print(tree) - elif args.visualize == "json": - tree.showJSON() - - print("\n\n") - - -def set_arguments() -> argparse.ArgumentParser: - """ - Parses user flags passed to TorBot - """ - parser = argparse.ArgumentParser( - prog="TorBot", usage="Gather and analayze data from Tor sites." - ) - parser.add_argument( - "-u", - "--url", - type=str, - default=None, - help="Specifiy a website link to crawl", - ) - parser.add_argument( - "--depth", type=int, help="Specifiy max depth of crawler (default 1)", default=1 - ) - parser.add_argument( - "--host", type=str, help="IP address for SOCKS5 proxy", default="127.0.0.1" - ) - parser.add_argument("--port", type=int, help="Port for SOCKS5 proxy", default=9050) - parser.add_argument( - "--save", type=str, choices=["tree", "json"], help="Save results in a file" - ) - parser.add_argument( - "--visualize", - type=str, - choices=["table", "tree", "json"], - help="Visualizes data collection.", - ) - parser.add_argument("-q", "--quiet", action="store_true") - parser.add_argument( - "--version", action="store_true", help="Show current version of TorBot." - ) - parser.add_argument( - "--update", - action="store_true", - help="Update TorBot to the latest stable version", - ) - parser.add_argument( - "--info", - action="store_true", - help="Info displays basic info of the scanned site. Only supports a single URL at a time.", - ) - parser.add_argument("-v", action="store_true", help="verbose logging") - parser.add_argument( - "--disable-socks5", - action="store_true", - help="Executes HTTP requests without using SOCKS5 proxy", - ) - parser.add_argument( - "--html", - choices=["save", "display"], - help="Saves / Displays the html of the onion link", - ) - - return parser +__all__ = ["get_version", "main", "run", "set_arguments"] if __name__ == "__main__": try: - arg_parser = set_arguments() - config_file_path = os.path.join( - os.path.dirname(os.path.realpath(__file__)), "pyproject.toml" - ) - try: - with open(config_file_path, "r") as f: - data = toml.load(f) - version = data["project"]["version"] - except Exception as e: - raise Exception("unable to find version from pyproject.toml.\n", e) - - run(arg_parser, version) + main() except KeyboardInterrupt: print("Interrupt received! Exiting cleanly...") diff --git a/pyproject.toml b/pyproject.toml index b9974823..6f8e52c2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "torbot" -version = "4.1.1" +version = "4.3.0" authors = [ { name="Akeem King", email="akeemtlking@gmail.com" }, { name="PS Narayanan", email="thepsnarayanan@gmail.com" }, @@ -38,7 +38,6 @@ dependencies = [ "scikit-learn>=1.5.0", "scipy>=1.13.0", "six>=1.16.0", - "sklearn>=0.0", "soupsieve>=2.8.4", "termcolor>=2.4.0", "texttable>=1.7.0", diff --git a/requirements.txt b/requirements.txt index 3b62b378..85eeaa33 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,8 +5,8 @@ certifi==2024.7.4 charset-normalizer==3.3.2 colorama==0.4.6 ; sys_platform == "win32" decorator==5.1.1 -h11==0.14.0 -httpcore==1.0.5 +h11==0.16.0 +httpcore==1.0.9 httpx[socks]==0.27.0 idna==3.15 igraph==0.11.4 @@ -27,9 +27,8 @@ python-dotenv==1.2.2 pywin32-ctypes==0.2.2 ; sys_platform == "win32" scikit-learn==1.5.0 scipy==1.13.0 -setuptools==78.1.1 +setuptools==83.0.0 six==1.16.0 -sklearn==0.0 sniffio==1.3.1 socksio==1.0.0 soupsieve==2.8.4 diff --git a/scripts/publish.sh b/scripts/publish.sh index c1564f1d..69adee94 100755 --- a/scripts/publish.sh +++ b/scripts/publish.sh @@ -1,7 +1,90 @@ #!/usr/bin/env bash set -euo pipefail +usage() { + cat <<'EOF' +Usage: + scripts/publish.sh [--check|--dry-run] + +By default this builds, validates, and uploads the current pyproject.toml +version to PyPI. Use --check or --dry-run to build and validate without +uploading. +EOF +} + +version="$(python3 - <<'PY' +try: + import tomllib +except ModuleNotFoundError: + import toml as tomllib + +with open("pyproject.toml", "r", encoding="utf-8") as handle: + print(tomllib.loads(handle.read())["project"]["version"]) +PY +)" + +upload=true +case "${1:-}" in + "") + ;; + "--upload") + ;; + "--check"|"--dry-run") + upload=false + ;; + "-h"|"--help") + usage + exit 0 + ;; + *) + usage >&2 + exit 2 + ;; +esac + +has_pypirc_credentials() { + local pypirc="${PYPIRC_PATH:-$HOME/.pypirc}" + + [[ -f "$pypirc" ]] || return 1 + grep -Eq '^[[:space:]]*username[[:space:]]*=' "$pypirc" || return 1 + grep -Eq '^[[:space:]]*password[[:space:]]*=' "$pypirc" || return 1 +} + +ensure_pypi_credentials() { + if [[ -n "${TWINE_PASSWORD:-}" ]]; then + export TWINE_USERNAME="${TWINE_USERNAME:-__token__}" + return + fi + + if has_pypirc_credentials; then + return + fi + + if [[ ! -t 0 ]]; then + echo "PyPI credentials were not found." >&2 + echo "Set TWINE_USERNAME=__token__ and TWINE_PASSWORD, or add ~/.pypirc." >&2 + return 1 + fi + + echo "PyPI API token was not found in TWINE_PASSWORD or ~/.pypirc." + read -r -s -p "Enter PyPI API token: " pypi_token + echo + + if [[ -z "$pypi_token" ]]; then + echo "No token entered; aborting upload." >&2 + return 1 + fi + + export TWINE_USERNAME="__token__" + export TWINE_PASSWORD="$pypi_token" +} + python3 -m build -python3 -m twine check dist/* -echo "Publishing is ready. Run the following when you have your PyPI token:" -echo "python3 -m twine upload dist/*" +python3 -m twine check "dist/torbot-${version}"* +if [[ "$upload" == true ]]; then + ensure_pypi_credentials + python3 -m twine upload "dist/torbot-${version}"* +else + echo "Package artifacts for torbot ${version} are ready." + echo "Run scripts/publish.sh to upload them to PyPI." +fi diff --git a/src/torbot/cli.py b/src/torbot/cli.py index 1eb287a1..62860585 100644 --- a/src/torbot/cli.py +++ b/src/torbot/cli.py @@ -1,27 +1,217 @@ -import os +import argparse +import logging import sys +from importlib import metadata +from pathlib import Path +import httpx import toml -REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) -if REPO_ROOT not in sys.path: - sys.path.insert(0, REPO_ROOT) +from torbot.modules.api import get_ip +from torbot.modules.app_launcher import launch_torbot_app +from torbot.modules.color import color +from torbot.modules.info import execute_all, fetch_html +from torbot.modules.linktree import LinkTree +from torbot.modules.updater import check_version -from main import run as run_cli, set_arguments # noqa: E402 + +def get_version() -> str: + """Return the installed TorBot version with a source checkout fallback.""" + try: + return metadata.version("torbot") + except metadata.PackageNotFoundError: + repo_root = Path(__file__).resolve().parents[2] + config_file_path = repo_root / "pyproject.toml" + try: + with config_file_path.open("r", encoding="utf-8") as handle: + data = toml.load(handle) + return data["project"]["version"] + except Exception as exc: + raise RuntimeError("unable to find version from pyproject.toml.") from exc + + +def print_tor_ip_address(client: httpx.Client) -> None: + """ + https://check.torproject.org/ tells you if you are using tor and it + displays your IP address which we scape and display + """ + resp = get_ip(client) + print(resp["header"]) + print(color(resp["body"], "yellow")) + + +def print_header(version: str) -> None: + """ + Prints the TorBot banner including version and license. + """ + license_msg = color("LICENSE: GNU Public License v3", "red") + banner = r""" + __ ____ ____ __ ______ + / /_/ __ \/ __ \/ /_ ____/_ __/ + / __/ / / / /_/ / __ \/ __ \/ / + / /_/ /_/ / _, _/ /_/ / /_/ / / + \__/\____/_/ |_/_____/\____/_/ v{VERSION} + """.format(VERSION=version) + banner = color(banner, "red") + + title = r""" + {banner} + ####################################################### + # TorBot - Dark Web OSINT Tool # + # GitHub : https://github.com/DedsecInside/TorBot # + # Help : use -h for help text # + ####################################################### + {license_msg} + """ + + title = title.format(license_msg=license_msg, banner=banner) + print(title) + + +def run(arg_parser: argparse.ArgumentParser, version: str) -> None: + args = arg_parser.parse_args() + + # setup logging + date_fmt = "%d-%b-%y %H:%M:%S" + logging_fmt = "%(asctime)s - %(levelname)s - %(message)s" + logging_lvl = logging.DEBUG if args.v else logging.INFO + logging.basicConfig(level=logging_lvl, format=logging_fmt, datefmt=date_fmt) + + # Print verison then exit + if args.version: + print(f"TorBot Version: {version}") + sys.exit() + + # check version and update if necessary + if args.update: + check_version() + sys.exit() + + if args.command == "app" or args.app: + sys.exit(launch_torbot_app(Path.cwd(), args.app_dir)) + + # URL is required for crawl-related actions + if not args.url: + arg_parser.print_help() + sys.exit() + + socks5_host = args.host + socks5_port = str(args.port) + socks5_proxy = f"socks5://{socks5_host}:{socks5_port}" + with httpx.Client( + timeout=60, proxy=socks5_proxy if not args.disable_socks5 else None + ) as client: + # print header and IP address if not set to quiet + if not args.quiet: + print_header(version) + print_tor_ip_address(client) + + if args.info: + execute_all(client, args.url) + + tree = LinkTree(url=args.url, depth=args.depth, client=client) + tree.load() + + # save data if desired + if args.save == "tree": + tree.save() + elif args.save == "json": + tree.saveJSON() + + if args.html == "display": + fetch_html(client, args.url, tree) + elif args.html == "save": + fetch_html(client, args.url, tree, save_html=True) + + # always print something, table is the default + if args.visualize == "table" or not args.visualize: + tree.showTable() + elif args.visualize == "tree": + print(tree) + elif args.visualize == "json": + tree.showJSON() + + print("\n\n") + + +def set_arguments() -> argparse.ArgumentParser: + """ + Parses user flags passed to TorBot + """ + parser = argparse.ArgumentParser( + prog="TorBot", usage="Gather and analayze data from Tor sites." + ) + parser.add_argument( + "command", + nargs="?", + choices=["app"], + help="Run optional TorBot desktop app", + ) + parser.add_argument( + "-u", + "--url", + type=str, + default=None, + help="Specifiy a website link to crawl", + ) + parser.add_argument( + "--depth", type=int, help="Specifiy max depth of crawler (default 1)", default=1 + ) + parser.add_argument( + "--host", type=str, help="IP address for SOCKS5 proxy", default="127.0.0.1" + ) + parser.add_argument("--port", type=int, help="Port for SOCKS5 proxy", default=9050) + parser.add_argument( + "--save", type=str, choices=["tree", "json"], help="Save results in a file" + ) + parser.add_argument( + "--visualize", + type=str, + choices=["table", "tree", "json"], + help="Visualizes data collection.", + ) + parser.add_argument("-q", "--quiet", action="store_true") + parser.add_argument( + "--version", action="store_true", help="Show current version of TorBot." + ) + parser.add_argument( + "--update", + action="store_true", + help="Update TorBot to the latest stable version", + ) + parser.add_argument( + "--app", + action="store_true", + help="Run optional TorBot desktop app", + ) + parser.add_argument( + "--app-dir", + type=str, + help="Path to a TorBotApp checkout", + ) + parser.add_argument( + "--info", + action="store_true", + help="Info displays basic info of the scanned site. Only supports a single URL at a time.", + ) + parser.add_argument("-v", action="store_true", help="verbose logging") + parser.add_argument( + "--disable-socks5", + action="store_true", + help="Executes HTTP requests without using SOCKS5 proxy", + ) + parser.add_argument( + "--html", + choices=["save", "display"], + help="Saves / Displays the html of the onion link", + ) + + return parser def main() -> None: """Run the TorBot CLI from an installed entry point.""" - arg_parser = set_arguments() - config_file_path = os.path.join(REPO_ROOT, "pyproject.toml") - try: - with open(config_file_path, "r", encoding="utf-8") as handle: - data = toml.load(handle) - version = data["project"]["version"] - except Exception as exc: - raise RuntimeError("unable to find version from pyproject.toml.") from exc - - run_cli(arg_parser, version) + run(set_arguments(), get_version()) if __name__ == "__main__": diff --git a/src/torbot/modules/app_launcher.py b/src/torbot/modules/app_launcher.py new file mode 100644 index 00000000..f47b0f78 --- /dev/null +++ b/src/torbot/modules/app_launcher.py @@ -0,0 +1,86 @@ +""" +Optional TorBot desktop app launcher. + +The Electron UI lives outside the Python package so CLI users do not need Node +or Electron dependencies. This module only discovers and starts that sibling +project when the user explicitly requests it. +""" +from __future__ import annotations + +import os +import shutil +import subprocess +from pathlib import Path + + +TORBOT_APP_REPOSITORY_URL = "https://github.com/KingAkeem/TorBotApp" + + +def _looks_like_torbot_app(path: Path) -> bool: + return path.is_dir() and (path / "package.json").is_file() + + +def find_torbot_app( + repo_root: str | os.PathLike[str], + app_dir: str | None = None, +) -> Path | None: + """ + Find an optional TorBotApp checkout. + + Discovery order: + 1. Explicit --app-dir value. + 2. TORBOT_APP_DIR environment variable. + 3. Sibling TorBotApp directory next to this repository. + """ + if app_dir: + resolved = Path(app_dir).expanduser().resolve() + return resolved if _looks_like_torbot_app(resolved) else None + + candidates = [] + env_app_dir = os.environ.get("TORBOT_APP_DIR") + if env_app_dir: + candidates.append(Path(env_app_dir).expanduser()) + + root = Path(repo_root).resolve() + cwd = Path.cwd().resolve() + candidates.extend( + [ + root.parent / "TorBotApp", + cwd / "TorBotApp", + cwd.parent / "TorBotApp", + ] + ) + + for candidate in candidates: + resolved = candidate.resolve() + if _looks_like_torbot_app(resolved): + return resolved + + return None + + +def launch_torbot_app( + repo_root: str | os.PathLike[str], + app_dir: str | None = None, +) -> int: + """ + Launch the optional Electron app with npm. + + Returns a process-style exit code so the CLI can exit cleanly. + """ + torbot_app_dir = find_torbot_app(repo_root, app_dir=app_dir) + if torbot_app_dir is None: + print("TorBotApp is not installed or could not be found.") + print(f"Clone it from {TORBOT_APP_REPOSITORY_URL}") + print("Expected layout: TorBot and TorBotApp as sibling directories.") + print("You can also set TORBOT_APP_DIR or pass --app-dir /path/to/TorBotApp.") + return 1 + + if shutil.which("npm") is None: + print("npm is required to launch TorBotApp, but it was not found on PATH.") + print("Install Node.js/npm, then run this command again.") + return 1 + + print(f"Starting TorBotApp from {torbot_app_dir}") + completed = subprocess.run(["npm", "start"], cwd=str(torbot_app_dir), check=False) + return completed.returncode diff --git a/src/torbot/modules/nlp/README.md b/src/torbot/modules/nlp/README.md index 0073143c..5697c220 100644 --- a/src/torbot/modules/nlp/README.md +++ b/src/torbot/modules/nlp/README.md @@ -1,11 +1,30 @@ -# Natural Language Processing Library +# Natural Language Processing Module -This library provides tool for performing natural language processing on websites. -This library is in it's infancy currently and can only be used for testing. +This module classifies crawled pages into broad website categories using the +packaged `website_classification.csv` dataset. -To test gathering data use: -`python3 gater_data.py` -* This will generate the data necessary to train the classification model +Runtime classification uses `classify(html)` from `main.py`. The classifier is +trained lazily from the CSV on first use and cached for the rest of the process, +so crawling multiple pages does not retrain the model for every page. -To predict the classification of a webiste use: -`classify.py` and provide the url +## Public API + +```python +from torbot.modules.nlp.main import classify + +category, confidence = classify("...") +``` + +`confidence` is the model's confidence for the selected category. It is not a +model-wide accuracy score. + +## Optional training-data export + +The crawler no longer needs a generated `training_data/` directory. If you need +one for experiments with `sklearn.datasets.load_files`, run: + +```sh +python3 -m torbot.modules.nlp.gather_data +``` + +This creates an ignored `training_data/` directory beside the NLP module. diff --git a/src/torbot/modules/nlp/gather_data.py b/src/torbot/modules/nlp/gather_data.py index f47750e5..f20cb3f9 100644 --- a/src/torbot/modules/nlp/gather_data.py +++ b/src/torbot/modules/nlp/gather_data.py @@ -1,15 +1,19 @@ import csv -import os - from pathlib import Path +from typing import Union + -os.chdir(Path(__file__).parent) +NLP_DIRECTORY = Path(__file__).resolve().parent +DEFAULT_CSV_PATH = NLP_DIRECTORY / "website_classification.csv" +DEFAULT_OUTPUT_DIRECTORY = NLP_DIRECTORY / "training_data" -def write_data(): +def write_data( + csv_path: Union[str, Path] = DEFAULT_CSV_PATH, + output_directory: Union[str, Path] = DEFAULT_OUTPUT_DIRECTORY, +) -> None: """ - Writes the training data from the csv file to a directory based on the - scikit-learn.datasets `load_files` specification. + Write CSV rows to a scikit-learn load_files-compatible directory tree. dataset source: https://www.kaggle.com/hetulmehta/website-classification @@ -20,17 +24,19 @@ def write_data(): category_2_folder/ file_43.txt file_44.txt ... """ - - with open("website_classification.csv") as csvfile: - website_reader = csv.reader(csvfile, delimiter=",") + output_path = Path(output_directory) + with Path(csv_path).open(newline="", encoding="utf-8", errors="replace") as csvfile: + website_reader = csv.DictReader(csvfile) for row in website_reader: - [id, website, content, category] = row - if category != "category": - category = category.replace("/", "+") - dir_name = f"training_data/{category}" - Path(dir_name).mkdir(parents=True, exist_ok=True) - with open(f"{dir_name}/{id}.txt", mode="w+") as txtfile: - txtfile.write(content) + row_id = (row.get("id") or "").strip() + content = row.get("cleaned_text") or "" + category = (row.get("category") or "").strip().replace("/", "+") + if not row_id or not content or not category: + continue + + category_directory = output_path / category + category_directory.mkdir(parents=True, exist_ok=True) + (category_directory / f"{row_id}.txt").write_text(content, encoding="utf-8") if __name__ == "__main__": diff --git a/src/torbot/modules/nlp/main.py b/src/torbot/modules/nlp/main.py index 9722c601..1c04f088 100644 --- a/src/torbot/modules/nlp/main.py +++ b/src/torbot/modules/nlp/main.py @@ -1,48 +1,121 @@ -import numpy as np -import os +""" +Website text classification helpers. + +The classifier is trained lazily from the packaged CSV and cached for the +process lifetime. That keeps LinkTree classification cheap after the first +call and avoids writing generated training files into the installed package. +""" +from __future__ import annotations + +import csv +from functools import lru_cache from pathlib import Path +from typing import List, Tuple, Union +import numpy as np from bs4 import BeautifulSoup -from sklearn.model_selection import train_test_split -from sklearn.pipeline import Pipeline +from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.linear_model import SGDClassifier -from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer -from sklearn.datasets import load_files +from sklearn.pipeline import Pipeline -def classify(data): - """ - Classify URL specified by user - """ - soup = BeautifulSoup(data, features="html.parser") - html = soup.get_text() +NLP_DIRECTORY = Path(__file__).resolve().parent +DEFAULT_DATASET_PATH = NLP_DIRECTORY / "website_classification.csv" + + +def extract_text(html: str) -> str: + """Extract normalized visible text from an HTML document.""" + soup = BeautifulSoup(html or "", features="html.parser") + return soup.get_text(" ", strip=True) + + +def load_training_rows( + csv_path: Union[str, Path] = DEFAULT_DATASET_PATH, +) -> Tuple[List[str], List[str]]: + """Load training text and labels from the website classification CSV.""" + path = Path(csv_path) + texts: list[str] = [] + labels: list[str] = [] + + with path.open(newline="", encoding="utf-8", errors="replace") as csvfile: + reader = csv.DictReader(csvfile) + for row in reader: + text = (row.get("cleaned_text") or "").strip() + category = (row.get("category") or "").strip() + if text and category: + texts.append(text) + labels.append(category) + + if not texts: + raise ValueError(f"no training rows found in {path}") + + if len(set(labels)) < 2: + raise ValueError(f"at least two categories are required in {path}") - # create classifier - clf = Pipeline( + return texts, labels + + +def build_classifier(csv_path: Union[str, Path] = DEFAULT_DATASET_PATH) -> Pipeline: + """Train a text classifier from the supplied dataset CSV.""" + texts, labels = load_training_rows(csv_path) + classifier = Pipeline( [ - ("vect", CountVectorizer()), - ("tfidf", TfidfTransformer()), - ("clf", SGDClassifier()), + ( + "tfidf", + TfidfVectorizer( + lowercase=True, + max_features=50000, + ngram_range=(1, 2), + sublinear_tf=True, + ), + ), + ( + "clf", + SGDClassifier( + class_weight="balanced", + loss="modified_huber", + max_iter=1000, + random_state=42, + tol=1e-3, + ), + ), ] ) - try: - os.chdir(Path(__file__).parent) - - dataset = load_files("training_data") - except FileNotFoundError: - print("Training data not found. Obtaining training data...") - print("This may take a while...") - from .gather_data import write_data - - write_data() - print("Training data obtained.") - dataset = load_files("training_data") - pass - x_train, x_test, y_train, y_test = train_test_split(dataset.data, dataset.target) - clf.fit(x_train, y_train) - - # returns an array of target_name values - predicted = clf.predict([html]) - accuracy = np.mean(predicted == y_test) - - return [dataset.target_names[predicted[0]], accuracy] + classifier.fit(texts, labels) + return classifier + + +@lru_cache(maxsize=4) +def get_classifier(csv_path: Union[str, Path] = DEFAULT_DATASET_PATH) -> Pipeline: + """Return a cached classifier for the requested dataset.""" + return build_classifier(Path(csv_path).resolve()) + + +def _prediction_score(classifier: Pipeline, text: str) -> float: + """Return a best-effort confidence score for a single prediction.""" + if hasattr(classifier, "predict_proba"): + probabilities = classifier.predict_proba([text])[0] + return float(np.max(probabilities)) + + if hasattr(classifier, "decision_function"): + scores = np.asarray(classifier.decision_function([text])) + return float(np.max(scores)) + + return 0.0 + + +def classify(data: str) -> List[Union[str, float]]: + """ + Classify the supplied HTML and return [category, confidence]. + + The confidence value is model confidence for the selected category, not a + model-wide accuracy score. + """ + text = extract_text(data) + if not text: + return ["unknown", 0.0] + + classifier = get_classifier() + prediction = classifier.predict([text])[0] + confidence = _prediction_score(classifier, text) + return [str(prediction), confidence] diff --git a/tests/test_cli_and_crawler.py b/tests/test_cli_and_crawler.py index b06cf709..3b33d7ff 100644 --- a/tests/test_cli_and_crawler.py +++ b/tests/test_cli_and_crawler.py @@ -1,4 +1,5 @@ from main import set_arguments +from torbot.modules.app_launcher import find_torbot_app, launch_torbot_app from torbot.modules.linktree import parse_links @@ -10,6 +11,40 @@ def test_version_flag_does_not_require_url() -> None: assert args.url is None +def test_app_command_does_not_require_url() -> None: + parser = set_arguments() + args = parser.parse_args(["app"]) + + assert args.command == "app" + assert args.url is None + + +def test_app_flag_does_not_require_url() -> None: + parser = set_arguments() + args = parser.parse_args(["--app"]) + + assert args.app is True + assert args.url is None + + +def test_find_torbot_app_from_explicit_directory(tmp_path) -> None: + app_dir = tmp_path / "TorBotApp" + app_dir.mkdir() + (app_dir / "package.json").write_text("{}", encoding="utf-8") + + assert find_torbot_app(tmp_path / "TorBot", app_dir=str(app_dir)) == app_dir + + +def test_launch_torbot_app_reports_missing_checkout(tmp_path, capsys) -> None: + exit_code = launch_torbot_app( + tmp_path / "TorBot", + app_dir=str(tmp_path / "missing-app"), + ) + + assert exit_code == 1 + assert "TorBotApp is not installed" in capsys.readouterr().out + + def test_parse_links_resolves_relative_urls_against_base_url() -> None: html = """ diff --git a/tests/test_nlp.py b/tests/test_nlp.py new file mode 100644 index 00000000..489e2078 --- /dev/null +++ b/tests/test_nlp.py @@ -0,0 +1,77 @@ +import importlib +import os +import sys + + +def import_real_nlp_module(): + """Import the real NLP module instead of the lightweight test stub.""" + sys.modules.pop("torbot.modules.nlp.main", None) + return importlib.import_module("torbot.modules.nlp.main") + + +def write_fixture_csv(path): + path.write_text( + "\n".join( + [ + "id,website,cleaned_text,category", + "1,https://sports.test,basketball team score playoff coach,Sports", + "2,https://sports2.test,football league match tournament goal,Sports", + "3,https://biz.test,revenue company shareholder market,Business", + "4,https://biz2.test,corporate earnings sales customer,Business", + ] + ), + encoding="utf-8", + ) + + +def test_extract_text_normalizes_html() -> None: + nlp = import_real_nlp_module() + + assert nlp.extract_text("

Hello

World

") == "Hello World" + + +def test_load_training_rows_uses_csv_without_training_directory(tmp_path) -> None: + nlp = import_real_nlp_module() + csv_path = tmp_path / "website_classification.csv" + write_fixture_csv(csv_path) + + texts, labels = nlp.load_training_rows(csv_path) + + assert len(texts) == 4 + assert sorted(set(labels)) == ["Business", "Sports"] + assert not (tmp_path / "training_data").exists() + + +def test_build_classifier_returns_category_and_confidence(tmp_path) -> None: + nlp = import_real_nlp_module() + csv_path = tmp_path / "website_classification.csv" + write_fixture_csv(csv_path) + + classifier = nlp.build_classifier(csv_path) + text = "basketball team wins playoff tournament" + prediction = classifier.predict([text])[0] + confidence = nlp._prediction_score(classifier, text) + + assert prediction == "Sports" + assert 0.0 <= confidence <= 1.0 + + +def test_classify_empty_html_returns_unknown() -> None: + nlp = import_real_nlp_module() + + assert nlp.classify("") == ["unknown", 0.0] + + +def test_gather_data_writes_training_files_without_changing_cwd(tmp_path) -> None: + from torbot.modules.nlp.gather_data import write_data + + csv_path = tmp_path / "website_classification.csv" + output_path = tmp_path / "training_data" + write_fixture_csv(csv_path) + before = os.getcwd() + + write_data(csv_path=csv_path, output_directory=output_path) + + assert os.getcwd() == before + assert (output_path / "Sports" / "1.txt").read_text(encoding="utf-8") + assert (output_path / "Business" / "3.txt").read_text(encoding="utf-8")