diff --git a/.gitignore b/.gitignore index d5d994e..b2c3dac 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ config.json hashcat.pot hate_crack/__init__.pyc +hate_crack/_version.py __pycache__/ hate_crack/__pycache__/ tests/__pycache__/ diff --git a/Dockerfile.test b/Dockerfile.test index b177a7c..a36f5bb 100644 --- a/Dockerfile.test +++ b/Dockerfile.test @@ -19,7 +19,7 @@ RUN python -m pip install -q uv==0.9.28 COPY . /workspace -RUN make install +RUN make submodules vendor-assets && uv tool install . && make clean-vendor ENV PATH="${HOME}/.local/bin:${PATH}" ENV HATE_CRACK_SKIP_INIT=1 diff --git a/Makefile b/Makefile index 44632b5..9278900 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,5 @@ .DEFAULT_GOAL := submodules -.PHONY: install reinstall dev-install dev-reinstall clean hashcat-utils submodules submodules-pre vendor-assets clean-vendor test coverage lint check ruff mypy +.PHONY: install reinstall update dev-install dev-reinstall clean hashcat-utils submodules submodules-pre vendor-assets clean-vendor test coverage lint check ruff mypy hashcat-utils: submodules $(MAKE) -C hashcat-utils @@ -58,7 +58,11 @@ install: submodules vendor-assets $(MAKE) clean-vendor; \ exit 1; \ fi - @uv tool install . + @uv tool install . --force --reinstall + @$(MAKE) clean-vendor + +update: submodules vendor-assets + @uv tool install . --force --reinstall @$(MAKE) clean-vendor reinstall: uninstall install diff --git a/README.md b/README.md index cca2616..0f9282a 100644 --- a/README.md +++ b/README.md @@ -69,8 +69,6 @@ These are required for certain download/extraction flows: Manual install commands: -Manual install commands: - Ubuntu/Kali: ```bash sudo apt-get update @@ -198,6 +196,12 @@ Install OS dependencies + tool (auto-detects macOS vs Debian/Ubuntu): make install ``` +Rebuild submodules and reinstall the tool (quick update after pulling changes): + +```bash +make update +``` + Reinstall the Python tool in-place (keeps OS deps as-is): ```bash @@ -236,7 +240,7 @@ make test Install the project with optional dev dependencies (includes type stubs, linters, and testing tools): ```bash -pip install -e ".[dev]" +make dev-install ``` ### Continuous Integration @@ -309,7 +313,8 @@ Create `.git/hooks/pre-push` to automatically run checks before pushing: #!/bin/bash set -e .venv/bin/ruff check hate_crack -.venv/bin/mypy hate_crack +.venv/bin/mypy --exclude HashcatRosetta --exclude hashcat-utils --ignore-missing-imports hate_crack +HATE_CRACK_SKIP_INIT=1 HATE_CRACK_RUN_E2E=0 HATE_CRACK_RUN_DOCKER_TESTS=0 HATE_CRACK_RUN_LIVE_TESTS=0 .venv/bin/python -m pytest echo "✓ Local checks passed!" ``` @@ -408,6 +413,21 @@ Set Hashview credentials in `config.json`: } ``` +#### Ollama Configuration + +The LLM Attack (option 15) uses Ollama to generate password candidates. Configure the model and context window in `config.json`: + +```json +{ + "ollamaModel": "qwen2.5", + "ollamaNumCtx": 8192 +} +``` + +- **`ollamaModel`** — The Ollama model to use for candidate generation (default: `qwen2.5`). +- **`ollamaNumCtx`** — Context window size for the model (default: `8192`). +- The Ollama URL defaults to `http://localhost:11434`. Ensure Ollama is running before using the LLM Attack. + #### Automatic Found Hash Merging (Download Left Only) When downloading left hashes (uncracked hashes), hate_crack automatically: @@ -443,8 +463,9 @@ $ ./hate_crack.py 1000 \___|_ /(____ /__| \___ >____\______ /|__| (____ /\___ >__|_ \ \/ \/ \/_____/ \/ \/ \/ \/ Version 2.0 - +``` +------------------------------------------------------------------- ## Testing The test suite is mostly offline and uses mocks/fixtures. Live network checks and @@ -509,7 +530,7 @@ All tests use mocked API calls, so they can run without connectivity to a Hashvi ### Continuous Integration -Tests automatically run on GitHub Actions for every push and pull request (Ubuntu, Python 3.13). +Tests automatically run on GitHub Actions for every push and pull request (Ubuntu, Python 3.9 through 3.14). ------------------------------------------------------------------- @@ -527,6 +548,7 @@ Tests automatically run on GitHub Actions for every push and pull request (Ubunt (12) Thorough Combinator Attack (13) Bandrel Methodology (14) Loopback Attack + (15) LLM Attack (90) Download rules from Hashmob.net (91) Analyze Hashcat Rules @@ -655,6 +677,13 @@ Uses hashcat's loopback mode to feed cracked passwords from the current session * Uses an empty wordlist with the --loopback flag to process previously cracked passwords * Automatically downloads Hashmob rules if no rules are available locally +#### LLM Attack +Uses a local Ollama instance to generate password candidates based on target company information. Prompts for company name, industry, and location, then sends these details to the configured LLM model to produce likely password guesses. The generated candidates are fed into a hashcat wordlist+rules attack. + +* Requires a running Ollama instance (default: `http://localhost:11434`) +* Configurable model and context window via `config.json` (see Ollama Configuration below) +* Prompts for target company name, industry, and location + #### Download Rules from Hashmob.net Downloads the latest rule files from Hashmob.net's rule repository. These rules are curated and optimized for password cracking and can be used with the Quick Crack and Loopback Attack modes. @@ -687,6 +716,12 @@ Interactive menu for downloading and managing wordlists from Weakpass.com via Bi ------------------------------------------------------------------- ### Version History +Version 2.0+ + Added LLM Attack (option 15) using Ollama for AI-generated password candidates + Added Ollama configuration keys (ollamaModel, ollamaNumCtx) + Auto-versioning via setuptools-scm from git tags + CI test fixes across Python 3.9–3.14 + Version 2.0 Modularized codebase into CLI/API/attacks modules Unified CLI options with config overrides (hashview, hashcat, wordlists, pipal) diff --git a/config.json.example b/config.json.example index 5215c2d..4c3aa0f 100644 --- a/config.json.example +++ b/config.json.example @@ -22,5 +22,7 @@ "bandrel_common_basedwords": "welcome,password,p@ssword,p@$$word,changeme,letmein,summer,winter,spring,springtime,fall,autumn,monday,tuesday,wednesday,thursday,friday,saturday,sunday,january,february,march,april,may,june,july,august,september,october,november,december,christmas,easter,covid19", "hashview_url": "http://localhost:8443", "hashview_api_key": "", - "hashmob_api_key": "" + "hashmob_api_key": "", + "ollamaModel": "qwen2.5", + "ollamaNumCtx": 8192 } diff --git a/hate_crack.py b/hate_crack.py index ce16395..f018b00 100755 --- a/hate_crack.py +++ b/hate_crack.py @@ -83,6 +83,7 @@ def get_main_menu_options(): "12": _attacks.thorough_combinator, "13": _attacks.bandrel_method, "14": _attacks.loopback_attack, + "15": _attacks.ollama_attack, "90": download_hashmob_rules, "91": weakpass_wordlist_menu, "92": download_hashmob_wordlists, diff --git a/hate_crack/__init__.py b/hate_crack/__init__.py index cb37f13..6b0d048 100644 --- a/hate_crack/__init__.py +++ b/hate_crack/__init__.py @@ -1 +1,13 @@ # hate_crack package +import re as _re + +from hate_crack._version import __version__ as _raw_version +from hate_crack._version import __version_tuple__ + +# Clean setuptools-scm version for display: +# "2.0.post1.dev0+g05b5d6dc7.d20260214" → "2.0+g05b5d6dc7" +# "2.0.post1.dev1+g1234abc" → "2.0+g1234abc" +# "2.0" → "2.0" +__version__ = _re.sub(r"(\.post\d+\.dev\d+|\.d\d{8})", "", _raw_version) + +__all__ = ["__version__", "__version_tuple__"] diff --git a/hate_crack/api.py b/hate_crack/api.py index d0ee8de..fb9909f 100644 --- a/hate_crack/api.py +++ b/hate_crack/api.py @@ -117,6 +117,29 @@ def get_hcat_tuning_args(): return [] +def get_hcat_potfile_path(): + """Return the resolved potfile path from config, or the default.""" + config_path = _resolve_config_path() + if config_path: + try: + with open(config_path) as f: + config = json.load(f) + raw = (config.get("hcatPotfilePath") or "").strip() + if raw: + return os.path.expanduser(raw) + except Exception: + pass + return os.path.expanduser("~/.hashcat/hashcat.potfile") + + +def get_hcat_potfile_args(): + """Return potfile args list for hashcat, e.g. ['--potfile-path=/path'].""" + pot = get_hcat_potfile_path() + if pot: + return [f"--potfile-path={pot}"] + return [] + + def cleanup_torrent_files(directory=None): """Remove stray .torrent files from the wordlists directory on graceful exit.""" if directory is None: @@ -805,7 +828,6 @@ class HashviewAPI: self, customer_id, hashfile_id, output_file=None, hash_type=None ): import sys - import subprocess url = f"{self.base_url}/v1/hashfiles/{hashfile_id}/left" resp = self.session.get(url, headers=self._auth_headers(), stream=True) @@ -882,7 +904,7 @@ class HashviewAPI: for line in f: line = line.strip() if line: - parts = line.rsplit(":", 1) # Split on last colon + parts = line.split(":", 1) # Split on first colon if len(parts) == 2: hash_part, clear_part = parts hf.write(hash_part + "\n") @@ -894,277 +916,41 @@ class HashviewAPI: f"Split found file into {hashes_count} hashes and {clears_count} clears" ) - # Run hashcat to combine them - combined_file = output_abs + ".out" - try: - # Execute hashcat: hashcat -m hash_type found_hashes found_clears --outfile output.out --outfile-format=1,2 - tuning_args = get_hcat_tuning_args() + # Append found hashes to the left file + with open(output_abs, "a", encoding="utf-8") as lf: + with open(found_hashes_file, "r", encoding="utf-8", errors="ignore") as fhf: + for line in fhf: + line = line.strip() + if line: + lf.write(line + "\n") + print(f"✓ Appended {hashes_count} found hashes to {output_abs}") - # Create temporary outfile for hashcat - temp_outfile = output_abs + ".tmp" - - if self.debug: - print( - f"[DEBUG] download_left_hashes: hash_type={hash_type}, type={type(hash_type)}" - ) - - # Build command with hash type if provided - cmd = ["hashcat", *tuning_args] - if hash_type: - cmd.extend(["-m", str(hash_type)]) - cmd.extend( - [ - found_hashes_file, - found_clears_file, - "--outfile", - temp_outfile, - "--outfile-format=1,2", - ] - ) - - if self.debug: - print(f"[DEBUG] Running command: {' '.join(cmd)}") - - print(f"Running: {' '.join(cmd)}") - - result = subprocess.run( - cmd, - stderr=subprocess.PIPE, - stdout=subprocess.PIPE, - text=True, - timeout=300, - ) - - if result.returncode != 0: - print(f"Warning: hashcat exited with code {result.returncode}") - if result.stderr: - print(f" stderr: {result.stderr}") - - # Append the output to the combined file - if os.path.exists(temp_outfile): - with open( - temp_outfile, "r", encoding="utf-8", errors="ignore" - ) as tmp_f: - with open(combined_file, "a", encoding="utf-8") as out_f: - out_f.write(tmp_f.read()) - - # Count lines appended - with open( - combined_file, "r", encoding="utf-8", errors="ignore" - ) as f: - combined_count = len(f.readlines()) - print( - f"✓ Appended cracked hashes to {combined_file} (total lines: {combined_count})" - ) - - # Clean up temp file - try: - os.remove(temp_outfile) - except Exception: - pass - else: - print("Note: No cracked hashes found") - - except FileNotFoundError: - print("✗ Error: hashcat not found in PATH") - except subprocess.TimeoutExpired: - print("✗ Error: hashcat execution timed out") - except Exception as e: - print(f"✗ Error running hashcat: {e}") - - # Clean up temporary files (keep when debug is enabled) - if not self.debug: - files_to_delete = [found_file, found_hashes_file, found_clears_file] - for temp_file in files_to_delete: - try: - if os.path.exists(temp_file): - os.remove(temp_file) - print(f"Deleted {temp_file}") - except Exception as e: - print(f"Warning: Could not delete {temp_file}: {e}") + # Append found hash:clear pairs to the potfile + potfile_path = get_hcat_potfile_path() + if potfile_path: + appended = 0 + with open(potfile_path, "a", encoding="utf-8") as pf: + with open(found_file, "r", encoding="utf-8", errors="ignore") as ff: + for line in ff: + line = line.strip() + if line and ":" in line: + pf.write(line + "\n") + appended += 1 + combined_count = appended + print(f"✓ Appended {appended} found hashes to potfile: {potfile_path}") else: - print("Debug enabled: keeping found and split files") + print("Warning: No potfile path configured, skipping potfile update") + + # Clean up the two found_ files + for f_path in (found_file, found_hashes_file, found_clears_file): + try: + os.remove(f_path) + except OSError: + pass except Exception as e: # If there's any error downloading found file, just skip it print(f"Note: Could not download found hashes: {e}") - # Clean up found file if it was partially written - if not self.debug: - try: - if os.path.exists(found_file): - os.remove(found_file) - except Exception: - pass - - return { - "output_file": output_file, - "size": downloaded, - "combined_count": combined_count, - "combined_file": combined_file, - } - - def download_found_hashes( - self, customer_id, hashfile_id, output_file=None, hash_type=None - ): - import sys - import subprocess - - url = f"{self.base_url}/v1/hashfiles/{hashfile_id}/found" - resp = self.session.get(url, headers=self._auth_headers(), stream=True) - resp.raise_for_status() - if output_file is None: - output_file = f"found_{customer_id}_{hashfile_id}.txt" - - # Convert to absolute path to ensure file is preserved if CWD changes - output_file = os.path.abspath(output_file) - - total = int(resp.headers.get("content-length", 0)) - downloaded = 0 - chunk_size = 8192 - with open(output_file, "wb") as f: - for chunk in resp.iter_content(chunk_size=chunk_size): - if chunk: - f.write(chunk) - downloaded += len(chunk) - if total > 0: - done = int(50 * downloaded / total) - bar = "[" + "=" * done + " " * (50 - done) + "]" - percent = 100 * downloaded / total - sys.stdout.write( - f"\rDownloading: {bar} {percent:5.1f}% ({downloaded}/{total} bytes)" - ) - sys.stdout.flush() - if total > 0: - sys.stdout.write("\n") - # If content-length is not provided, just print size at end - if total == 0: - print(f"Downloaded {downloaded} bytes.") - - # Split found file into hashes and clears - output_dir = os.path.dirname(os.path.abspath(output_file)) or os.getcwd() - found_hashes_file = os.path.join( - output_dir, f"found_hashes_{customer_id}_{hashfile_id}.txt" - ) - found_clears_file = os.path.join( - output_dir, f"found_clears_{customer_id}_{hashfile_id}.txt" - ) - - hashes_count = 0 - clears_count = 0 - combined_count = 0 - combined_file = None - - try: - with ( - open(found_hashes_file, "w", encoding="utf-8") as hf, - open(found_clears_file, "w", encoding="utf-8") as cf, - ): - with open(output_file, "r", encoding="utf-8", errors="ignore") as f: - for line in f: - line = line.strip() - if line: - parts = line.rsplit(":", 1) # Split on last colon - if len(parts) == 2: - hash_part, clear_part = parts - hf.write(hash_part + "\n") - cf.write(clear_part + "\n") - hashes_count += 1 - clears_count += 1 - - print( - f"✓ Split found file into {hashes_count} hashes and {clears_count} clears" - ) - - # Run hashcat to combine them into an output file - combined_file = output_file + ".out" - try: - tuning_args = get_hcat_tuning_args() - - # Create temporary outfile for hashcat - temp_outfile = output_file + ".tmp" - - if self.debug: - print( - f"[DEBUG] download_found_hashes: hash_type={hash_type}, type={type(hash_type)}" - ) - - # Build command with hash type if provided - cmd = ["hashcat", *tuning_args] - if hash_type: - cmd.extend(["-m", str(hash_type)]) - cmd.extend( - [ - found_hashes_file, - found_clears_file, - "--outfile", - temp_outfile, - "--outfile-format=1,2", - ] - ) - - if self.debug: - print(f"[DEBUG] Running command: {' '.join(cmd)}") - - print(f"Running: {' '.join(cmd)}") - - result = subprocess.run( - cmd, - stderr=subprocess.PIPE, - stdout=subprocess.PIPE, - text=True, - timeout=300, - ) - - if result.returncode != 0: - print(f"Warning: hashcat exited with code {result.returncode}") - if result.stderr: - print(f" stderr: {result.stderr}") - - # Write the output file - if os.path.exists(temp_outfile): - with open( - temp_outfile, "r", encoding="utf-8", errors="ignore" - ) as tmp_f: - with open(combined_file, "w", encoding="utf-8") as out_f: - out_f.write(tmp_f.read()) - - # Count lines in output - with open( - combined_file, "r", encoding="utf-8", errors="ignore" - ) as f: - combined_count = len(f.readlines()) - print(f"✓ Created {combined_file} (total lines: {combined_count})") - - # Clean up temp file - try: - os.remove(temp_outfile) - except Exception: - pass - else: - print("Note: No cracked hashes generated") - - except FileNotFoundError: - print("✗ Error: hashcat not found in PATH") - except subprocess.TimeoutExpired: - print("✗ Error: hashcat execution timed out") - except Exception as e: - print(f"✗ Error running hashcat: {e}") - - # Clean up temporary files (keep when debug is enabled) - if not self.debug: - files_to_delete = [found_hashes_file, found_clears_file] - for temp_file in files_to_delete: - try: - if os.path.exists(temp_file): - os.remove(temp_file) - except Exception as e: - if self.debug: - print(f"Warning: Could not delete {temp_file}: {e}") - else: - print("Debug enabled: keeping split files") - - except Exception as e: - print(f"✗ Error splitting found file: {e}") return { "output_file": output_file, diff --git a/hate_crack/attacks.py b/hate_crack/attacks.py index 1d3b5dc..d6ba8fd 100644 --- a/hate_crack/attacks.py +++ b/hate_crack/attacks.py @@ -486,3 +486,18 @@ def middle_combinator(ctx: Any) -> None: def bandrel_method(ctx: Any) -> None: ctx.hcatBandrel(ctx.hcatHashType, ctx.hcatHashFile) + + +def ollama_attack(ctx: Any) -> None: + print("\n\tLLM Attack") + company = input("Company name: ").strip() + industry = input("Industry: ").strip() + location = input("Location: ").strip() + target_info = { + "company": company, + "industry": industry, + "location": location, + } + ctx.hcatOllama( + ctx.hcatHashType, ctx.hcatHashFile, "target", target_info + ) diff --git a/hate_crack/main.py b/hate_crack/main.py index 9d738e5..7effe45 100755 --- a/hate_crack/main.py +++ b/hate_crack/main.py @@ -19,6 +19,8 @@ import readline import subprocess import shlex import argparse +import urllib.request +import urllib.error from types import SimpleNamespace #!/usr/bin/env python3 @@ -451,9 +453,30 @@ except KeyError as e: default_config.get("hcatDebugLogPath", "./hashcat_debug") ) +ollamaUrl = "http://" + os.environ.get("OLLAMA_HOST", "localhost:11434") +try: + ollamaModel = config_parser["ollamaModel"] +except KeyError as e: + print( + "{0} is not defined in config.json using defaults from config.json.example".format( + e + ) + ) + ollamaModel = default_config.get("ollamaModel", "qwen2.5") +try: + ollamaNumCtx = int(config_parser["ollamaNumCtx"]) +except KeyError as e: + print( + "{0} is not defined in config.json using defaults from config.json.example".format( + e + ) + ) + ollamaNumCtx = int(default_config.get("ollamaNumCtx", 8192)) + hcatExpanderBin = "expander.bin" hcatCombinatorBin = "combinator.bin" hcatPrinceBin = "pp64.bin" +hcatHcstat2genBin = "hcstat2gen.bin" def _resolve_wordlist_path(wordlist, base_dir): @@ -588,7 +611,6 @@ hcatGoodMeasureBaseList = _normalize_wordlist_setting( hcatGoodMeasureBaseList, wordlists_dir ) hcatPrinceBaseList = _normalize_wordlist_setting(hcatPrinceBaseList, wordlists_dir) - if not SKIP_INIT: # Verify hashcat binary is available # hcatBin should be in PATH or be an absolute path (resolved from hcatPath + hcatBin if configured) @@ -653,6 +675,20 @@ if not SKIP_INIT: except SystemExit: print("PRINCE attacks will not be available.") + # Verify hcstat2gen binary (optional, for LLM attacks) + # Note: hcstat2gen is part of hashcat-utils, already in hate_crack repo + hcstat2gen_path = ( + hate_path + "/hashcat-utils/bin/" + hcatHcstat2genBin + ) + try: + ensure_binary( + hcstat2gen_path, + build_dir=os.path.join(hate_path, "hashcat-utils"), + name="hcstat2gen", + ) + except SystemExit: + print("LLM attacks will not be available.") + except Exception as e: print(f"Module initialization error: {e}") if not shutil.which("hashcat") and not os.path.exists("/usr/bin/hashcat"): @@ -737,15 +773,17 @@ def usage(): def ascii_art(): + from hate_crack import __version__ + print(r""" - ___ ___ __ _________ __ + ___ ___ __ _________ __ / | \_____ _/ |_ ____ \_ ___ \____________ ____ | | __ / ~ \__ \\ __\/ __ \ / \ \/\_ __ \__ \ _/ ___\| |/ / -\ Y // __ \| | \ ___/ \ \____| | \// __ \\ \___| < +\ Y // __ \| | \ ___/ \ \____| | \// __ \\ \___| < \___|_ /(____ /__| \___ >____\______ /|__| (____ /\___ >__|_ \ \/ \/ \/_____/ \/ \/ \/ \/ - Version 2.0 + Version """ + __version__ + """ """) @@ -882,20 +920,25 @@ def _write_field_sorted_unique(input_path, output_path, field_index, delimiter=" def _run_hashcat_show(hash_type, hash_file, output_path): + result = subprocess.run( + [ + hcatBin, + "--show", + # Use hashcat's built-in potfile unless configured otherwise. + *([f"--potfile-path={hcatPotfilePath}"] if hcatPotfilePath else []), + "-m", + str(hash_type), + hash_file, + ], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + check=False, + ) with open(output_path, "w") as out: - subprocess.run( - [ - hcatBin, - "--show", - # Use hashcat's built-in potfile unless configured otherwise. - *([f"--potfile-path={hcatPotfilePath}"] if hcatPotfilePath else []), - "-m", - str(hash_type), - hash_file, - ], - stdout=out, - check=False, - ) + for line in result.stdout.decode("utf-8", errors="ignore").splitlines(): + # hashcat --show prints parse errors to stdout; skip non-result lines + if ":" in line and not line.startswith(("Hash parsing error", "* ")): + out.write(line + "\n") # Brute Force Attack @@ -1446,6 +1489,245 @@ def hcatBandrel(hcatHashType, hcatHashFile): hcatProcess.kill() +# Pull an Ollama model via the /api/pull streaming endpoint +def _pull_ollama_model(url, model): + """Pull an Ollama model. Returns True on success, False on failure.""" + print(f"Model '{model}' not found locally. Pulling from Ollama...") + pull_url = f"{url}/api/pull" + payload = json.dumps({"name": model, "stream": True}).encode("utf-8") + req = urllib.request.Request( + pull_url, + data=payload, + headers={"Content-Type": "application/json"}, + ) + try: + with urllib.request.urlopen(req) as resp: + for raw_line in resp: + line = raw_line.decode("utf-8", errors="replace").strip() + if not line: + continue + try: + data = json.loads(line) + except json.JSONDecodeError: + continue + status = data.get("status") + if status: + print(f" {status}") + except urllib.error.HTTPError as e: + print(f"Error pulling model: HTTP {e.code}") + return False + except urllib.error.URLError as e: + print(f"Error: Could not connect to Ollama: {e}") + return False + except Exception as e: + print(f"Error pulling model: {e}") + return False + print(f"Successfully pulled model '{model}'.") + return True + + +# LLM Ollama Attack +def hcatOllama(hcatHashType, hcatHashFile, mode, context_data): + global hcatProcess + candidates_path = f"{hcatHashFile}.ollama_candidates" + + # Step A: Build LLM prompt based on mode + if mode == "wordlist": + wordlist_path = context_data + if not os.path.isfile(wordlist_path): + print(f"Error: Wordlist not found: {wordlist_path}") + return + lines = [] + try: + with open(wordlist_path, "r", errors="ignore") as f: + for line in f: + stripped = line.strip() + if not stripped: + continue + # Use only content after the first colon (e.g. hash:password -> password) + if ":" in stripped: + stripped = stripped.split(":", 1)[1] + if stripped: + lines.append(stripped) + except Exception as e: + print(f"Error reading wordlist: {e}") + return + print(f"Loaded {len(lines)} passwords from wordlist.") + wordlist_sample = "\n".join(lines) + prompt = ( + "Generate baseword to be used in a denylist for keeping users from setting their passwords with these basewords." + "Study the patterns, character choices, and structures. Focus on patterns like capitalization, leetspeak, suffixes, and common substitutions. Here are the sample passwords:\n" + f"{wordlist_sample}" + ) + elif mode == "target": + company = context_data.get("company", "") + industry = context_data.get("industry", "") + location = context_data.get("location", "") + prompt = ( + "Generate baseword to be used in a denylist for keeping users from setting their passwords with these basewords. We are protecting the employees of " + f"{company} in the {industry} industry located in {location}.\n\n" + "Include variations with:\n" + "- Company name variations and abbreviations\n" + "- Common password patterns (Season+Year, Name+Numbers)\n" + "- Keyboard walks and common substitutions (@ for a, 3 for e, etc.)\n" + "- Location-based words and local references\n" + "- Industry-specific terminology\n" + "Output ONLY the passwords, one per line, no numbering or explanation." + ) + else: + print(f"Error: Unknown LLM generation mode: {mode}") + return + + # Step B: Call Ollama API to generate candidates + print(f"Generating password candidates via Ollama ({ollamaModel})...") + api_url = f"{ollamaUrl}/api/generate" + payload = json.dumps({ + "model": ollamaModel, + "prompt": prompt, + "stream": False, + "options": {"num_ctx": ollamaNumCtx}, + }).encode("utf-8") + + if debug_mode: + print(f"[DEBUG] Ollama API URL: {api_url}") + print(f"[DEBUG] Ollama request payload: {payload.decode('utf-8')}") + + try: + req = urllib.request.Request( + api_url, + data=payload, + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(req, timeout=600) as resp: + result = json.loads(resp.read().decode("utf-8")) + if debug_mode: + print(f"[DEBUG] Ollama response: {json.dumps(result, indent=2)}") + except urllib.error.HTTPError as e: + if e.code == 404: + if _pull_ollama_model(ollamaUrl, ollamaModel): + try: + req = urllib.request.Request( + api_url, + data=payload, + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(req, timeout=600) as resp: + result = json.loads(resp.read().decode("utf-8")) + if debug_mode: + print(f"[DEBUG] Ollama response (after pull): {json.dumps(result, indent=2)}") + except Exception as retry_err: + print(f"Error calling Ollama API after pull: {retry_err}") + return + else: + print(f"Could not pull model '{ollamaModel}'. Aborting LLM attack.") + return + else: + print(f"Error: Could not connect to Ollama at {ollamaUrl}: {e}") + print("Ensure Ollama is running (ollama serve) and try again.") + return + except urllib.error.URLError as e: + print(f"Error: Could not connect to Ollama at {ollamaUrl}: {e}") + print("Ensure Ollama is running (ollama serve) and try again.") + return + except Exception as e: + print(f"Error calling Ollama API: {e}") + return + + response_text = result.get("response", "") + if "I'm sorry, but I can't help with that" in response_text: + print("Error: Ollama refused the request. Try a different model or adjust your prompt.") + return + raw_lines = response_text.strip().split("\n") + # Filter out blank lines and lines that look like numbering/explanation + candidates = [] + for line in raw_lines: + stripped = line.strip() + if not stripped: + continue + # Strip leading numbering like "1. " or "1) " or "- " + cleaned = re.sub(r"^\d+[.)]\s*", "", stripped) + cleaned = re.sub(r"^[-*]\s*", "", cleaned) + cleaned = cleaned.strip() + if cleaned and len(cleaned) <= 128: + candidates.append(cleaned) + + if not candidates: + print("Error: Ollama returned no usable password candidates.") + return + + try: + with open(candidates_path, "w") as f: + for candidate in candidates: + f.write(candidate + "\n") + except Exception as e: + print(f"Error writing candidates file: {e}") + return + + print(f"Generated {len(candidates)} password candidates -> {candidates_path}") + if debug_mode: + filtered_count = len(raw_lines) - len(candidates) + print(f"[DEBUG] Filtered out {filtered_count} lines from Ollama response ({len(raw_lines)} raw -> {len(candidates)} candidates)") + + # Step C: Run hashcat wordlist attack with LLM-generated candidates (no rules) + print("Running wordlist attack with LLM-generated candidates...") + cmd = [ + hcatBin, + "-m", + hcatHashType, + hcatHashFile, + "--session", + generate_session_id(), + "-o", + f"{hcatHashFile}.out", + candidates_path, + ] + cmd.extend(shlex.split(hcatTuning)) + _append_potfile_arg(cmd) + hcatProcess = subprocess.Popen(cmd) + try: + hcatProcess.wait() + except KeyboardInterrupt: + print("Killing PID {0}...".format(str(hcatProcess.pid))) + hcatProcess.kill() + return + + # Step D: Run hashcat with LLM candidates against every rule in the rules directory + rule_files = sorted( + f for f in os.listdir(rulesDirectory) if f != ".DS_Store" + ) + if not rule_files: + print("No rule files found in rules directory. Skipping rule-based attacks.") + return + + print(f"\nRunning LLM candidates with {len(rule_files)} rule file(s) from {rulesDirectory}...") + for rule in rule_files: + rule_path = os.path.join(rulesDirectory, rule) + print(f"\n\tRunning with rule: {rule}") + cmd = [ + hcatBin, + "-m", + hcatHashType, + hcatHashFile, + "--session", + generate_session_id(), + "-o", + f"{hcatHashFile}.out", + "-r", + rule_path, + candidates_path, + ] + cmd.extend(shlex.split(hcatTuning)) + _append_potfile_arg(cmd) + hcatProcess = subprocess.Popen(cmd) + try: + hcatProcess.wait() + except KeyboardInterrupt: + print("Killing PID {0}...".format(str(hcatProcess.pid))) + hcatProcess.kill() + return + + + # Middle fast Combinator Attack def hcatMiddleCombinator(hcatHashType, hcatHashFile): global hcatProcess @@ -1961,13 +2243,10 @@ def hashview_api(): menu_options.append(("download_wordlist", "Download Wordlist")) menu_options.append( ( - "download_left", - "Download Left Hashes (with automatic merge if found)", + "download_hashes", + "Download Hashes (left + found to potfile)", ) ) - menu_options.append( - ("download_found", "Download Found Hashes (with automatic split)") - ) if hcatHashFile: menu_options.append( ("upload_hashfile_job", "Upload Hashfile and Create Job") @@ -2364,7 +2643,7 @@ def hashview_api(): except Exception as e: print(f"\n✗ Error uploading hashfile: {str(e)}") - elif option_key == "download_left": + elif option_key == "download_hashes": # Download left hashes try: while True: @@ -2533,157 +2812,6 @@ def hashview_api(): except Exception as e: print(f"\n✗ Error downloading hashes: {str(e)}") - elif option_key == "download_found": - # Download found hashes - try: - while True: - # First, list customers to help user select - customers_result = api_harness.list_customers() - customers = ( - customers_result.get("customers", []) - if isinstance(customers_result, dict) - else customers_result - ) - if customers: - api_harness.display_customers_multicolumn(customers) - else: - print("\nNo customers found.") - - # Select or create customer - customer_input = input( - "\nEnter customer ID or N to create new: " - ).strip() - if customer_input.lower() == "n": - customer_name = input("Enter customer name: ").strip() - if customer_name: - try: - result = api_harness.create_customer(customer_name) - print( - f"\n✓ Success: {result.get('msg', 'Customer created')}" - ) - customer_id = result.get( - "customer_id" - ) or result.get("id") - if not customer_id: - print("\n✗ Error: Customer ID not returned.") - continue - print(f" Customer ID: {customer_id}") - except Exception as e: - print(f"\n✗ Error creating customer: {str(e)}") - continue - else: - print("\n✗ Error: Customer name cannot be empty.") - continue - else: - try: - customer_id = int(customer_input) - except ValueError: - print( - "\n✗ Error: Invalid ID entered. Please enter a numeric ID or N." - ) - continue - - # List hashfiles for the customer - try: - customer_hashfiles = api_harness.get_customer_hashfiles( - customer_id - ) - - if not customer_hashfiles: - print( - f"\nNo hashfiles found for customer ID {customer_id}" - ) - continue - - print("\n" + "=" * 120) - print(f"Hashfiles for Customer ID {customer_id}:") - print("=" * 120) - print(f"{'ID':<10} {'Hash Type':<10} {'Name':<96}") - print("-" * 120) - hashfile_map = {} - for hf in customer_hashfiles: - hf_id = hf.get("id") - hf_name = hf.get("name", "N/A") - hf_type = ( - hf.get("hash_type") or hf.get("hashtype") or "N/A" - ) - if hf_id is None: - continue - # Truncate long names to fit within 120 columns - if len(str(hf_name)) > 96: - hf_name = str(hf_name)[:93] + "..." - if debug_mode: - print( - f"[DEBUG] Hashfile {hf_id}: hash_type={hf.get('hash_type')}, hashtype={hf.get('hashtype')}, combined={hf_type}" - ) - print(f"{hf_id:<10} {hf_type:<10} {hf_name:<96}") - hashfile_map[int(hf_id)] = hf_type - print("=" * 120) - print(f"Total: {len(hashfile_map)} hashfile(s)") - except Exception as e: - print(f"\nWarning: Could not list hashfiles: {e}") - continue - - while True: - try: - hashfile_id_input = input( - "\nEnter hashfile ID: " - ).strip() - hashfile_id = int(hashfile_id_input) - except ValueError: - print( - "\n✗ Error: Invalid ID entered. Please enter a numeric ID." - ) - continue - if hashfile_id not in hashfile_map: - print( - "\n✗ Error: Hashfile ID not in the list. Please try again." - ) - continue - break - break - - # Set output filename automatically - output_file = f"found_{customer_id}_{hashfile_id}.txt" - - # Get hash type for hashcat from the hashfile map - selected_hash_type = hashfile_map.get(hashfile_id) - if debug_mode: - print( - f"[DEBUG] selected_hash_type from map: {selected_hash_type}" - ) - if not selected_hash_type or selected_hash_type == "N/A": - try: - details = api_harness.get_hashfile_details(hashfile_id) - selected_hash_type = details.get("hashtype") - if debug_mode: - print( - f"[DEBUG] selected_hash_type from get_hashfile_details: {selected_hash_type}" - ) - except Exception as e: - if debug_mode: - print(f"[DEBUG] Error fetching hashfile details: {e}") - selected_hash_type = None - - # Download the found hashes - if debug_mode: - print( - f"[DEBUG] Calling download_found_hashes with hash_type={selected_hash_type}" - ) - download_result = api_harness.download_found_hashes( - customer_id, hashfile_id, output_file - ) - print(f"\n✓ Success: Downloaded {download_result['size']} bytes") - print(f" File: {download_result['output_file']}") - if selected_hash_type: - print(f" Hash mode: {selected_hash_type}") - print("\nFound hashes downloaded successfully. These are already cracked hashes.") - - except ValueError: - print("\n✗ Error: Invalid ID entered. Please enter a numeric ID.") - except Exception as e: - print(f"\n✗ Error downloading found hashes: {str(e)}") - elif option_key == "back": break @@ -2756,6 +2884,10 @@ def loopback_attack(): return _attacks.loopback_attack(_attack_ctx()) +def ollama_attack(): + return _attacks.ollama_attack(_attack_ctx()) + + # convert hex words for recycling def convert_hex(working_file): processed_words = [] @@ -2983,6 +3115,7 @@ def get_main_menu_options(): "12": thorough_combinator, "13": bandrel_method, "14": loopback_attack, + "15": ollama_attack, "90": download_hashmob_rules, "91": analyze_rules, "92": download_hashmob_wordlists, @@ -3124,8 +3257,8 @@ def main(): ) hv_download_left = hashview_subparsers.add_parser( - "download-left", - help="Download left hashes for a hashfile", + "download-hashes", + help="Download left hashes and append found hashes to potfile", ) hv_download_left.add_argument( "--customer-id", required=True, type=int, help="Customer ID" @@ -3139,17 +3272,6 @@ def main(): help="Hash type for hashcat (e.g., 1000 for NTLM)", ) - hv_download_found = hashview_subparsers.add_parser( - "download-found", - help="Download found hashes for a hashfile", - ) - hv_download_found.add_argument( - "--customer-id", required=True, type=int, help="Customer ID" - ) - hv_download_found.add_argument( - "--hashfile-id", required=True, type=int, help="Hashfile ID" - ) - hv_upload_hashfile_job = hashview_subparsers.add_parser( "upload-hashfile-job", help="Upload a hashfile and create a job", @@ -3190,8 +3312,7 @@ def main(): hashview_subcommands = [ "upload-cracked", "upload-wordlist", - "download-left", - "download-found", + "download-hashes", "upload-hashfile-job", ] has_hashview_flag = "--hashview" in argv @@ -3315,7 +3436,7 @@ def main(): print(f" Wordlist ID: {result['wordlist_id']}") sys.exit(0) - if args.hashview_command == "download-left": + if args.hashview_command == "download-hashes": download_result = api_harness.download_left_hashes( args.customer_id, args.hashfile_id, @@ -3325,15 +3446,6 @@ def main(): print(f" File: {download_result['output_file']}") sys.exit(0) - if args.hashview_command == "download-found": - download_result = api_harness.download_found_hashes( - args.customer_id, - args.hashfile_id, - ) - print(f"\n✓ Success: Downloaded {download_result['size']} bytes") - print(f" File: {download_result['output_file']}") - sys.exit(0) - if args.hashview_command == "upload-hashfile-job": hashfile_path = resolve_path(args.file) if not hashfile_path or not os.path.isfile(hashfile_path): @@ -3563,6 +3675,7 @@ def main(): print("\t(12) Thorough Combinator Attack") print("\t(13) Bandrel Methodology") print("\t(14) Loopback Attack") + print("\t(15) LLM Attack") print("\n\t(90) Download rules from Hashmob.net") print("\n\t(91) Analyze Hashcat Rules") print("\t(92) Download wordlists from Hashmob.net") diff --git a/pyproject.toml b/pyproject.toml index fe707bf..4414bbc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,10 @@ [build-system] -requires = ["setuptools>=69"] +requires = ["setuptools>=69", "setuptools-scm>=8"] build-backend = "setuptools.build_meta" [project] name = "hate_crack" -version = "2.0" +dynamic = ["version"] description = "Menu driven Python wrapper for hashcat" readme = "README.md" requires-python = ">=3.13" @@ -37,6 +37,10 @@ hate_crack = [ "princeprocessor/**", ] +[tool.setuptools_scm] +version_file = "hate_crack/_version.py" +version_scheme = "no-guess-dev" + [tool.ruff] exclude = [ "build", diff --git a/tests/test_hashview.py b/tests/test_hashview.py index 3f5181f..8a1df3c 100644 --- a/tests/test_hashview.py +++ b/tests/test_hashview.py @@ -270,52 +270,6 @@ class TestHashviewAPI: assert "Cookie" in auth_headers or "uuid" in str(auth_headers) assert HASHVIEW_API_KEY in str(auth_headers) - def test_download_found_hashes(self, api, tmp_path): - """Test downloading found hashes: real API if possible, else mock.""" - hashview_url, hashview_api_key = self._get_hashview_config() - customer_id = os.environ.get("HASHVIEW_CUSTOMER_ID") - hashfile_id = os.environ.get("HASHVIEW_HASHFILE_ID") - if all([hashview_url, hashview_api_key, customer_id, hashfile_id]): - # Real API test - real_api = HashviewAPI(hashview_url, hashview_api_key) - output_file = tmp_path / f"found_{customer_id}_{hashfile_id}.txt" - result = real_api.download_found_hashes( - int(customer_id), int(hashfile_id), output_file=str(output_file) - ) - assert os.path.exists(result["output_file"]) - with open(result["output_file"], "rb") as f: - content = f.read() - print(f"[DEBUG] Downloaded {len(content)} bytes to {result['output_file']}") - assert result["size"] == len(content) - else: - # Mock test - mock_response = Mock() - mock_response.content = b"hash1:pass1\nhash2:pass2\n" - mock_response.raise_for_status = Mock() - mock_response.headers = {"content-length": "0"} - - def iter_content(chunk_size=8192): - yield mock_response.content - - mock_response.iter_content = iter_content - api.session.get.return_value = mock_response - - output_file = tmp_path / "found_1_2.txt" - result = api.download_found_hashes(1, 2, output_file=str(output_file)) - assert os.path.exists(result["output_file"]) - with open(result["output_file"], "rb") as f: - content = f.read() - assert content == b"hash1:pass1\nhash2:pass2\n" - assert result["size"] == len(content) - - # Verify auth headers were passed in the found hashes download call - call_args_list = api.session.get.call_args_list - found_call = [c for c in call_args_list if "found" in str(c)][0] - assert found_call.kwargs.get("headers") is not None - auth_headers = found_call.kwargs.get("headers") - assert "Cookie" in auth_headers or "uuid" in str(auth_headers) - assert HASHVIEW_API_KEY in str(auth_headers) - def test_download_wordlist(self, api, tmp_path): """Test downloading a wordlist: real API if possible, else mock.""" hashview_url, hashview_api_key = self._get_hashview_config() @@ -633,6 +587,10 @@ class TestHashviewAPI: # Set up session.get to return different responses api.session.get.side_effect = [mock_left_response, mock_found_response] + # Mock potfile path so cleanup isn't blocked by missing ~/.hashcat dir + potfile = str(tmp_path / "hashcat.potfile") + monkeypatch.setattr("hate_crack.api.get_hcat_potfile_path", lambda: potfile) + # Download left hashes (should auto-download and split found for hashcat) left_file = tmp_path / "left_1_2.txt" result = api.download_left_hashes(1, 2, output_file=str(left_file)) @@ -640,21 +598,17 @@ class TestHashviewAPI: # Verify left file was created assert os.path.exists(result["output_file"]) - # Verify found file was downloaded and deleted + # Verify found files are cleaned up after merge found_file = tmp_path / "found_1_2.txt" - assert not os.path.exists(found_file), ( - "Found file should be deleted after split" - ) - assert not (other_cwd / "found_1_2.txt").exists() + assert not os.path.exists(found_file), "Found file should be deleted after merge" - # Verify split files were created and deleted found_hashes_file = tmp_path / "found_hashes_1_2.txt" found_clears_file = tmp_path / "found_clears_1_2.txt" assert not os.path.exists(str(found_hashes_file)), ( - "Split hashes file should be deleted" + "Split hashes file should be deleted after merge" ) assert not os.path.exists(str(found_clears_file)), ( - "Split clears file should be deleted" + "Split clears file should be deleted after merge" ) def test_download_left_id_matching(self, api, tmp_path): diff --git a/tests/test_hashview_cli_subcommands.py b/tests/test_hashview_cli_subcommands.py index 7831855..4c59ff2 100644 --- a/tests/test_hashview_cli_subcommands.py +++ b/tests/test_hashview_cli_subcommands.py @@ -28,15 +28,6 @@ class DummyHashviewAPI: "size": 10, } - def download_found_hashes(self, customer_id, hashfile_id, output_file=None): - self.calls.append( - ("download_found_hashes", customer_id, hashfile_id, output_file) - ) - return { - "output_file": output_file or f"found_{customer_id}_{hashfile_id}.txt", - "size": 12, - } - def upload_hashfile( self, file_path, customer_id, hash_type, file_format=5, hashfile_name=None ): diff --git a/tests/test_hashview_cli_subcommands_subprocess.py b/tests/test_hashview_cli_subcommands_subprocess.py index 1e31c88..529e57d 100644 --- a/tests/test_hashview_cli_subcommands_subprocess.py +++ b/tests/test_hashview_cli_subcommands_subprocess.py @@ -77,8 +77,7 @@ def _ensure_customer_one(): "--name", "TestWordlist", ], - ["hashview", "download-left", "--customer-id", "1", "--hashfile-id", "2"], - ["hashview", "download-found", "--customer-id", "1", "--hashfile-id", "2"], + ["hashview", "download-hashes", "--customer-id", "1", "--hashfile-id", "2"], [ "hashview", "upload-hashfile-job", @@ -142,45 +141,25 @@ def test_hashview_subcommands_live_downloads(): base_cmd = [sys.executable, HATE_CRACK_SCRIPT, "hashview"] customer_id = _ensure_customer_one() - left_cmd = base_cmd + [ - "download-left", + dl_cmd = base_cmd + [ + "download-hashes", "--customer-id", str(customer_id), "--hashfile-id", os.environ["HASHVIEW_HASHFILE_ID"], ] - left = subprocess.run( - left_cmd, + dl = subprocess.run( + dl_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, cwd=REPO_ROOT, env=env, ) - left_out = left.stdout + left.stderr - assert left.returncode == 0, left_out - assert "Downloaded" in left_out - assert "left_" in left_out - - found_cmd = base_cmd + [ - "download-found", - "--customer-id", - str(customer_id), - "--hashfile-id", - os.environ["HASHVIEW_HASHFILE_ID"], - ] - found = subprocess.run( - found_cmd, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - cwd=REPO_ROOT, - env=env, - ) - found_out = found.stdout + found.stderr - assert found.returncode == 0, found_out - assert "Downloaded" in found_out - assert "found_" in found_out + dl_out = dl.stdout + dl.stderr + assert dl.returncode == 0, dl_out + assert "Downloaded" in dl_out + assert "left_" in dl_out @pytest.mark.skipif( diff --git a/tests/test_pull_ollama_model.py b/tests/test_pull_ollama_model.py new file mode 100644 index 0000000..54abe21 --- /dev/null +++ b/tests/test_pull_ollama_model.py @@ -0,0 +1,684 @@ +"""Unit tests for _pull_ollama_model helper, hcatOllama steps A-C, and ollama_attack handler.""" + +import io +import json +import os +import urllib.error +import urllib.request +from contextlib import contextmanager +from types import SimpleNamespace +from unittest import mock + +import pytest + +os.environ["HATE_CRACK_SKIP_INIT"] = "1" +from hate_crack import main as hc_main # noqa: E402 + + +# --------------------------------------------------------------------------- +# Shared test infrastructure +# --------------------------------------------------------------------------- + +OLLAMA_URL = "http://localhost:11434" +MODEL = "llama3.2" + + +@pytest.fixture +def ollama_env(tmp_path): + """Create the filesystem layout hcatOllama expects.""" + hash_file = tmp_path / "hashes.txt" + hash_file.touch() + + wordlist = tmp_path / "sample.txt" + wordlist.write_text("password\n123456\nletmein\n") + + return SimpleNamespace( + tmp_path=tmp_path, + hash_file=str(hash_file), + wordlist=str(wordlist), + ) + + +@contextmanager +def ollama_globals(tmp_path, tuning="", potfile=""): + """Patch the hc_main globals that hcatOllama reads.""" + rules_dir = str(tmp_path / "rules") + os.makedirs(rules_dir, exist_ok=True) + with mock.patch.object(hc_main, "ollamaUrl", OLLAMA_URL), \ + mock.patch.object(hc_main, "ollamaModel", MODEL), \ + mock.patch.object(hc_main, "hcatBin", "/usr/bin/hashcat"), \ + mock.patch.object(hc_main, "hcatTuning", tuning), \ + mock.patch.object(hc_main, "hcatPotfilePath", potfile), \ + mock.patch.object(hc_main, "hate_path", str(tmp_path)), \ + mock.patch.object(hc_main, "rulesDirectory", rules_dir), \ + mock.patch("hate_crack.main.generate_session_id", return_value="test_session"): + yield + + +def _make_proc(wait_return=0): + """Create a mock subprocess that works with both wait() and communicate().""" + proc = mock.MagicMock() + proc.wait.return_value = wait_return + proc.communicate.return_value = (b"", b"") + proc.returncode = wait_return + return proc + + +def _generate_response(passwords): + """Build a fake urlopen context-manager that returns an Ollama /api/generate JSON body.""" + body = json.dumps({"response": "\n".join(passwords)}).encode() + resp = mock.MagicMock() + resp.__enter__ = mock.Mock(return_value=io.BytesIO(body)) + resp.__exit__ = mock.Mock(return_value=False) + return resp + + +def _urlopen_with_response(passwords): + """Return a urlopen mock that always succeeds with the given passwords.""" + return mock.patch( + "hate_crack.main.urllib.request.urlopen", + return_value=_generate_response(passwords), + ) + + +# --------------------------------------------------------------------------- +# _pull_ollama_model tests +# --------------------------------------------------------------------------- + +class TestPullOllamaModel: + """Tests for _pull_ollama_model().""" + + OLLAMA_URL = "http://localhost:11434" + MODEL = "llama3.2" + + def _make_stream_response(self, statuses): + """Build a fake streaming response (newline-delimited JSON).""" + lines = [json.dumps({"status": s}).encode() + b"\n" for s in statuses] + return io.BytesIO(b"".join(lines)) + + @mock.patch("hate_crack.main.urllib.request.urlopen") + def test_successful_pull(self, mock_urlopen, capsys): + mock_urlopen.return_value.__enter__ = mock.Mock( + return_value=self._make_stream_response( + ["pulling manifest", "downloading sha256:abc123", "success"] + ) + ) + mock_urlopen.return_value.__exit__ = mock.Mock(return_value=False) + + result = hc_main._pull_ollama_model(self.OLLAMA_URL, self.MODEL) + + assert result is True + captured = capsys.readouterr() + assert "not found locally" in captured.out + assert "Successfully pulled" in captured.out + assert "pulling manifest" in captured.out + + @mock.patch("hate_crack.main.urllib.request.urlopen") + def test_pull_http_error(self, mock_urlopen, capsys): + mock_urlopen.side_effect = urllib.error.HTTPError( + url="http://localhost:11434/api/pull", + code=500, + msg="Internal Server Error", + hdrs=None, + fp=None, + ) + + result = hc_main._pull_ollama_model(self.OLLAMA_URL, self.MODEL) + + assert result is False + captured = capsys.readouterr() + assert "HTTP 500" in captured.out + + @mock.patch("hate_crack.main.urllib.request.urlopen") + def test_pull_url_error(self, mock_urlopen, capsys): + mock_urlopen.side_effect = urllib.error.URLError("Connection refused") + + result = hc_main._pull_ollama_model(self.OLLAMA_URL, self.MODEL) + + assert result is False + captured = capsys.readouterr() + assert "Could not connect" in captured.out + + @mock.patch("hate_crack.main.urllib.request.urlopen") + def test_pull_generic_exception(self, mock_urlopen, capsys): + mock_urlopen.side_effect = RuntimeError("unexpected") + + result = hc_main._pull_ollama_model(self.OLLAMA_URL, self.MODEL) + + assert result is False + captured = capsys.readouterr() + assert "unexpected" in captured.out + + @mock.patch("hate_crack.main.urllib.request.urlopen") + def test_pull_handles_empty_status_lines(self, mock_urlopen, capsys): + """Blank lines and missing status keys should not crash.""" + lines = b'{"status": "downloading"}\n\n{"other": "data"}\n' + mock_urlopen.return_value.__enter__ = mock.Mock( + return_value=io.BytesIO(lines) + ) + mock_urlopen.return_value.__exit__ = mock.Mock(return_value=False) + + result = hc_main._pull_ollama_model(self.OLLAMA_URL, self.MODEL) + + assert result is True + + @mock.patch("hate_crack.main.urllib.request.urlopen") + def test_pull_request_payload(self, mock_urlopen): + """Verify the pull request sends the correct model name and stream=True.""" + mock_urlopen.return_value.__enter__ = mock.Mock( + return_value=self._make_stream_response(["success"]) + ) + mock_urlopen.return_value.__exit__ = mock.Mock(return_value=False) + + hc_main._pull_ollama_model(self.OLLAMA_URL, self.MODEL) + + call_args = mock_urlopen.call_args + req = call_args[0][0] + body = json.loads(req.data.decode("utf-8")) + assert body["name"] == self.MODEL + assert body["stream"] is True + assert req.full_url == f"{self.OLLAMA_URL}/api/pull" + + +# --------------------------------------------------------------------------- +# hcatOllama 404-retry integration tests +# --------------------------------------------------------------------------- + +class TestHcatOllama404Retry: + """Test that hcatOllama auto-pulls on 404 and retries the generate call.""" + + OLLAMA_URL = "http://localhost:11434" + MODEL = "llama3.2" + + def _generate_response(self, passwords): + body = json.dumps({"response": "\n".join(passwords)}).encode() + return io.BytesIO(body) + + def _setup_env(self, tmp_path): + """Create the files/dirs hcatOllama needs before it hits the API.""" + hash_file = str(tmp_path / "hashes.txt") + open(hash_file, "w").close() + + # Create a small wordlist for "wordlist" mode + wordlist = str(tmp_path / "sample.txt") + with open(wordlist, "w") as f: + f.write("password\n123456\nletmein\n") + + return hash_file, wordlist + + @mock.patch("hate_crack.main._pull_ollama_model") + @mock.patch("hate_crack.main.urllib.request.urlopen") + def test_404_triggers_pull_then_retries(self, mock_urlopen, mock_pull, capsys, tmp_path): + """A 404 on generate should trigger a pull, then retry successfully.""" + hash_file, wordlist = self._setup_env(tmp_path) + + # First call: 404, second call: success + generate_ok = mock.MagicMock() + generate_ok.__enter__ = mock.Mock( + return_value=self._generate_response(["Password1", "Summer2024"]) + ) + generate_ok.__exit__ = mock.Mock(return_value=False) + + mock_urlopen.side_effect = [ + urllib.error.HTTPError( + url=f"{self.OLLAMA_URL}/api/generate", + code=404, + msg="Not Found", + hdrs=None, + fp=None, + ), + generate_ok, + ] + mock_pull.return_value = True + + rules_dir = str(tmp_path / "rules") + os.makedirs(rules_dir, exist_ok=True) + with mock.patch.object(hc_main, "ollamaUrl", self.OLLAMA_URL), \ + mock.patch.object(hc_main, "ollamaModel", self.MODEL), \ + mock.patch.object(hc_main, "hcatBin", "/usr/bin/hashcat"), \ + mock.patch.object(hc_main, "hcatTuning", ""), \ + mock.patch.object(hc_main, "hcatPotfilePath", ""), \ + mock.patch.object(hc_main, "hate_path", str(tmp_path)), \ + mock.patch.object(hc_main, "rulesDirectory", rules_dir), \ + mock.patch("hate_crack.main.generate_session_id", return_value="test_session"), \ + mock.patch("subprocess.Popen") as mock_popen: + + proc = _make_proc() + mock_popen.return_value = proc + + hc_main.hcatOllama("0", hash_file, "wordlist", wordlist) + + mock_pull.assert_called_once_with(self.OLLAMA_URL, self.MODEL) + # urlopen called twice: first 404, then retry + assert mock_urlopen.call_count == 2 + + @mock.patch("hate_crack.main._pull_ollama_model") + @mock.patch("hate_crack.main.urllib.request.urlopen") + def test_404_pull_fails_aborts(self, mock_urlopen, mock_pull, capsys, tmp_path): + """If pull fails after 404, hcatOllama should abort gracefully.""" + hash_file, wordlist = self._setup_env(tmp_path) + + mock_urlopen.side_effect = urllib.error.HTTPError( + url=f"{self.OLLAMA_URL}/api/generate", + code=404, + msg="Not Found", + hdrs=None, + fp=None, + ) + mock_pull.return_value = False + + with mock.patch.object(hc_main, "ollamaUrl", self.OLLAMA_URL), \ + mock.patch.object(hc_main, "ollamaModel", self.MODEL), \ + mock.patch.object(hc_main, "hate_path", str(tmp_path)): + + hc_main.hcatOllama("0", hash_file, "wordlist", wordlist) + + captured = capsys.readouterr() + assert "Could not pull model" in captured.out + mock_pull.assert_called_once() + + @mock.patch("hate_crack.main.urllib.request.urlopen") + def test_non_404_http_error_propagates(self, mock_urlopen, capsys, tmp_path): + """Non-404 HTTP errors should not trigger a pull attempt.""" + hash_file, wordlist = self._setup_env(tmp_path) + + mock_urlopen.side_effect = urllib.error.HTTPError( + url=f"{self.OLLAMA_URL}/api/generate", + code=500, + msg="Internal Server Error", + hdrs=None, + fp=None, + ) + + with mock.patch.object(hc_main, "ollamaUrl", self.OLLAMA_URL), \ + mock.patch.object(hc_main, "ollamaModel", self.MODEL), \ + mock.patch.object(hc_main, "hate_path", str(tmp_path)): + + hc_main.hcatOllama("0", hash_file, "wordlist", wordlist) + + captured = capsys.readouterr() + # HTTPError is a subclass of URLError, so it hits the URLError handler + assert "Could not connect to Ollama" in captured.out or "Error calling Ollama API" in captured.out + assert "500" in captured.out + + +# --------------------------------------------------------------------------- +# Step A: Mode routing, prompt construction, early-return error paths +# --------------------------------------------------------------------------- + +class TestHcatOllamaModeRouting: + """Test mode selection, prompt building, and early-return errors.""" + + def test_unknown_mode_prints_error(self, ollama_env, capsys): + """Bad mode string → error message, no API call.""" + with ollama_globals(ollama_env.tmp_path), \ + mock.patch("hate_crack.main.urllib.request.urlopen") as mock_url: + hc_main.hcatOllama("0", ollama_env.hash_file, "bogus", "") + + captured = capsys.readouterr() + assert "Unknown LLM generation mode" in captured.out + mock_url.assert_not_called() + + def test_missing_wordlist_prints_error(self, ollama_env, capsys): + """Non-existent wordlist path → error, no API call.""" + with ollama_globals(ollama_env.tmp_path), \ + mock.patch("hate_crack.main.urllib.request.urlopen") as mock_url: + hc_main.hcatOllama( + "0", ollama_env.hash_file, "wordlist", + "/no/such/wordlist.txt", + ) + + captured = capsys.readouterr() + assert "Wordlist not found" in captured.out + mock_url.assert_not_called() + + def test_wordlist_mode_reads_all_lines(self, ollama_env, capsys): + """All non-blank lines from the wordlist should appear in the prompt.""" + # Write 600 lines to the wordlist + big_wordlist = ollama_env.tmp_path / "big.txt" + big_wordlist.write_text("\n".join(f"pass{i}" for i in range(600)) + "\n") + + captured_payload = {} + + def fake_urlopen(req, **kwargs): + captured_payload["data"] = json.loads(req.data.decode()) + return _generate_response(["Password1"]) + + with ollama_globals(ollama_env.tmp_path), \ + mock.patch("hate_crack.main.urllib.request.urlopen", side_effect=fake_urlopen), \ + mock.patch("subprocess.Popen") as mock_popen: + proc = _make_proc() + mock_popen.return_value = proc + hc_main.hcatOllama( + "0", ollama_env.hash_file, "wordlist", str(big_wordlist), + ) + + prompt_text = captured_payload["data"]["prompt"] + # All lines should be included in the prompt + assert "pass0" in prompt_text + assert "pass499" in prompt_text + assert "pass599" in prompt_text + + def test_target_mode_includes_context_in_prompt(self, ollama_env, capsys): + """Company, industry, and location should appear in the prompt.""" + captured_payload = {} + + def fake_urlopen(req, **kwargs): + captured_payload["data"] = json.loads(req.data.decode()) + return _generate_response(["AcmeCorp2024"]) + + target_info = { + "company": "AcmeCorp", + "industry": "Finance", + "location": "New York", + } + with ollama_globals(ollama_env.tmp_path), \ + mock.patch("hate_crack.main.urllib.request.urlopen", side_effect=fake_urlopen), \ + mock.patch("subprocess.Popen") as mock_popen: + proc = _make_proc() + mock_popen.return_value = proc + hc_main.hcatOllama( + "0", ollama_env.hash_file, "target", target_info, + ) + + prompt_text = captured_payload["data"]["prompt"] + assert "AcmeCorp" in prompt_text + assert "Finance" in prompt_text + assert "New York" in prompt_text + + +# --------------------------------------------------------------------------- +# Step B: Candidate filtering / post-processing +# --------------------------------------------------------------------------- + +class TestHcatOllamaCandidateFiltering: + """Test regex stripping, blank-line removal, 128-char limit, empty-result handling.""" + + def _run_with_response(self, ollama_env, response_text): + """Run hcatOllama with a canned Ollama response_text, return candidates file content.""" + body = json.dumps({"response": response_text}).encode() + resp = mock.MagicMock() + resp.__enter__ = mock.Mock(return_value=io.BytesIO(body)) + resp.__exit__ = mock.Mock(return_value=False) + + candidates_path = f"{ollama_env.hash_file}.ollama_candidates" + captured = {} + + def capture_popen(cmd, **kwargs): + # Read the candidates file before hashcat "runs" (cleanup happens after) + if os.path.isfile(candidates_path): + with open(candidates_path) as f: + captured["content"] = f.read() + return _make_proc() + + with ollama_globals(ollama_env.tmp_path), \ + mock.patch("hate_crack.main.urllib.request.urlopen", return_value=resp), \ + mock.patch("subprocess.Popen", side_effect=capture_popen): + hc_main.hcatOllama( + "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist, + ) + + return captured.get("content") + + def test_strips_numeric_dot_prefix(self, ollama_env): + content = self._run_with_response(ollama_env, "1. Password1\n2. Summer2024") + assert content is not None + lines = [l for l in content.strip().split("\n") if l] + assert "Password1" in lines + assert "Summer2024" in lines + # No line should start with a digit+dot + for line in lines: + assert not line.startswith("1.") + assert not line.startswith("2.") + + def test_strips_dash_and_asterisk_prefix(self, ollama_env): + content = self._run_with_response(ollama_env, "- foo\n* bar") + assert content is not None + lines = [l for l in content.strip().split("\n") if l] + assert "foo" in lines + assert "bar" in lines + + def test_skips_blank_lines(self, ollama_env): + content = self._run_with_response(ollama_env, "alpha\n\n\nbeta\n\n") + assert content is not None + lines = [l for l in content.strip().split("\n") if l] + assert lines == ["alpha", "beta"] + + def test_rejects_over_128_chars(self, ollama_env): + long_pw = "A" * 129 + content = self._run_with_response(ollama_env, f"short\n{long_pw}\nkeep") + assert content is not None + lines = [l for l in content.strip().split("\n") if l] + assert "short" in lines + assert "keep" in lines + assert long_pw not in lines + + def test_accepts_exactly_128_chars(self, ollama_env): + exact_pw = "B" * 128 + content = self._run_with_response(ollama_env, f"{exact_pw}\nother") + assert content is not None + lines = [l for l in content.strip().split("\n") if l] + assert exact_pw in lines + + def test_empty_response_prints_error(self, ollama_env, capsys): + self._run_with_response(ollama_env, "") + captured = capsys.readouterr() + assert "no usable" in captured.out.lower() + + def test_missing_response_key(self, ollama_env, capsys): + """If the JSON has no 'response' key, treat as empty.""" + body = json.dumps({"other": "stuff"}).encode() + resp = mock.MagicMock() + resp.__enter__ = mock.Mock(return_value=io.BytesIO(body)) + resp.__exit__ = mock.Mock(return_value=False) + + with ollama_globals(ollama_env.tmp_path), \ + mock.patch("hate_crack.main.urllib.request.urlopen", return_value=resp): + hc_main.hcatOllama( + "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist, + ) + + captured = capsys.readouterr() + assert "no usable" in captured.out.lower() + + +# --------------------------------------------------------------------------- +# Step B: API error paths (non-404) +# --------------------------------------------------------------------------- + +class TestHcatOllamaApiErrors: + """Test connection errors and generic exceptions during the generate call.""" + + def test_url_error_prints_connection_error(self, ollama_env, capsys): + with ollama_globals(ollama_env.tmp_path), \ + mock.patch( + "hate_crack.main.urllib.request.urlopen", + side_effect=urllib.error.URLError("Connection refused"), + ): + hc_main.hcatOllama( + "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist, + ) + + captured = capsys.readouterr() + assert "Could not connect" in captured.out + assert "Ensure Ollama is running" in captured.out + + def test_generic_exception_prints_error(self, ollama_env, capsys): + with ollama_globals(ollama_env.tmp_path), \ + mock.patch( + "hate_crack.main.urllib.request.urlopen", + side_effect=RuntimeError("boom"), + ): + hc_main.hcatOllama( + "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist, + ) + + captured = capsys.readouterr() + assert "Error calling Ollama API" in captured.out + + def test_generate_request_payload(self, ollama_env): + """Verify the /api/generate request has correct URL, model, stream=false.""" + captured_req = {} + + def fake_urlopen(req, **kwargs): + captured_req["url"] = req.full_url + captured_req["body"] = json.loads(req.data.decode()) + return _generate_response(["Password1"]) + + with ollama_globals(ollama_env.tmp_path), \ + mock.patch("hate_crack.main.urllib.request.urlopen", side_effect=fake_urlopen), \ + mock.patch("subprocess.Popen") as mock_popen: + proc = _make_proc() + mock_popen.return_value = proc + hc_main.hcatOllama( + "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist, + ) + + assert captured_req["url"] == f"{OLLAMA_URL}/api/generate" + assert captured_req["body"]["model"] == MODEL + assert captured_req["body"]["stream"] is False + assert len(captured_req["body"]["prompt"]) > 0 + + +# --------------------------------------------------------------------------- +# Step C: Hashcat command construction +# --------------------------------------------------------------------------- + +class TestHcatOllamaHashcatCommand: + """Test hashcat command flags and process handling.""" + + def _run_full(self, ollama_env, tuning="", potfile=""): + """Run hcatOllama through all steps, returning the list of Popen calls.""" + popen_calls = [] + + def track_popen(cmd, **kwargs): + popen_calls.append((list(cmd), dict(kwargs))) + return _make_proc() + + with ollama_globals(ollama_env.tmp_path, tuning=tuning, potfile=potfile), \ + _urlopen_with_response(["Password1"]), \ + mock.patch("subprocess.Popen", side_effect=track_popen), \ + mock.patch("hate_crack.main.generate_session_id", return_value="test_session"): + hc_main.hcatOllama( + "1000", ollama_env.hash_file, "wordlist", ollama_env.wordlist, + ) + + return popen_calls + + def test_hashcat_command_structure(self, ollama_env): + """First Popen call should be a straight wordlist attack with candidates file, + followed by rule-based attacks.""" + calls = self._run_full(ollama_env) + assert len(calls) >= 1, "Expected at least one Popen call (hashcat)" + + # First call: base wordlist attack (no rules) + hashcat_cmd = calls[0][0] + assert hashcat_cmd[0] == "/usr/bin/hashcat" + assert "-m" in hashcat_cmd + m_idx = hashcat_cmd.index("-m") + assert hashcat_cmd[m_idx + 1] == "1000" + # Should use the candidates wordlist file + candidates_path = f"{ollama_env.hash_file}.ollama_candidates" + assert candidates_path in hashcat_cmd + # First call should NOT have mask attack flags + assert "-a" not in hashcat_cmd + assert "--markov-hcstat2" not in " ".join(hashcat_cmd) + assert "--increment" not in hashcat_cmd + # First call should NOT have -r (rule) flag + assert "-r" not in hashcat_cmd + + # Subsequent calls: rule-based attacks + if len(calls) > 1: + for rule_cmd, _ in calls[1:]: + assert "-r" in rule_cmd + assert candidates_path in rule_cmd + + def test_hashcat_includes_tuning(self, ollama_env): + """-w 3 tuning flag should be appended to hashcat cmd.""" + calls = self._run_full(ollama_env, tuning="-w 3") + hashcat_cmd = calls[0][0] + assert "-w" in hashcat_cmd + w_idx = hashcat_cmd.index("-w") + assert hashcat_cmd[w_idx + 1] == "3" + + def test_hashcat_includes_potfile(self, ollama_env): + """--potfile-path should be present when hcatPotfilePath is set.""" + calls = self._run_full(ollama_env, potfile="/tmp/test.potfile") + hashcat_cmd = calls[0][0] + potfile_flags = [f for f in hashcat_cmd if "--potfile-path=" in f] + assert len(potfile_flags) == 1 + assert potfile_flags[0] == "--potfile-path=/tmp/test.potfile" + + def test_hashcat_keyboard_interrupt(self, ollama_env, capsys): + """KeyboardInterrupt during hashcat should kill the process.""" + proc = mock.MagicMock() + proc.pid = 99999 + proc.wait.side_effect = KeyboardInterrupt() + + with ollama_globals(ollama_env.tmp_path), \ + _urlopen_with_response(["Password1"]), \ + mock.patch("subprocess.Popen", return_value=proc), \ + mock.patch("hate_crack.main.generate_session_id", return_value="test_session"): + hc_main.hcatOllama( + "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist, + ) + + captured = capsys.readouterr() + assert "Killing PID" in captured.out + proc.kill.assert_called_once() + + +# --------------------------------------------------------------------------- +# Step F: Cleanup +# --------------------------------------------------------------------------- + +class TestHcatOllamaWordlistPersistence: + """Test that the generated wordlist file persists after the run.""" + + def test_candidates_file_persists(self, ollama_env): + """ollama_candidates wordlist should remain after the run.""" + with ollama_globals(ollama_env.tmp_path), \ + _urlopen_with_response(["Password1"]), \ + mock.patch("subprocess.Popen", return_value=_make_proc()), \ + mock.patch("hate_crack.main.generate_session_id", return_value="test_session"): + hc_main.hcatOllama( + "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist, + ) + + candidates_path = f"{ollama_env.hash_file}.ollama_candidates" + assert os.path.isfile(candidates_path), "candidates wordlist should persist" + + +# --------------------------------------------------------------------------- +# attacks.py: ollama_attack() UI handler +# --------------------------------------------------------------------------- + +class TestOllamaAttackHandler: + """Test the ollama_attack(ctx) menu handler in attacks.py.""" + + def _make_ctx(self, **overrides): + """Build a mock ctx with the attributes ollama_attack() reads.""" + ctx = mock.MagicMock() + ctx.hcatHashType = "0" + ctx.hcatHashFile = "/tmp/hashes.txt" + ctx.hcatWordlists = "/tmp/wordlists" + for k, v in overrides.items(): + setattr(ctx, k, v) + return ctx + + def test_target_mode(self): + """ollama_attack prompts for company/industry/location and calls hcatOllama.""" + from hate_crack.attacks import ollama_attack + + ctx = self._make_ctx() + with mock.patch("builtins.input", side_effect=["AcmeCorp", "Finance", "NYC"]): + ollama_attack(ctx) + + ctx.hcatOllama.assert_called_once() + call_args = ctx.hcatOllama.call_args + assert call_args[0][2] == "target" + target_info = call_args[0][3] + assert target_info["company"] == "AcmeCorp" + assert target_info["industry"] == "Finance" + assert target_info["location"] == "NYC" + diff --git a/tests/test_ui_menu_options.py b/tests/test_ui_menu_options.py index 5dbe771..9ecf38f 100644 --- a/tests/test_ui_menu_options.py +++ b/tests/test_ui_menu_options.py @@ -25,6 +25,7 @@ MENU_OPTION_TEST_CASES = [ ("12", CLI_MODULE._attacks, "thorough_combinator", "thorough"), ("13", CLI_MODULE._attacks, "bandrel_method", "bandrel"), ("14", CLI_MODULE._attacks, "loopback_attack", "loopback"), + ("15", CLI_MODULE._attacks, "ollama_attack", "ollama"), ("90", CLI_MODULE, "download_hashmob_rules", "hashmob-rules"), ("91", CLI_MODULE, "weakpass_wordlist_menu", "weakpass-menu"), ("92", CLI_MODULE, "download_hashmob_wordlists", "hashmob-wordlists"), diff --git a/tools/ollama_benchmark.py b/tools/ollama_benchmark.py new file mode 100755 index 0000000..e7c7ce9 --- /dev/null +++ b/tools/ollama_benchmark.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 +"""Benchmark multiple Ollama models for password candidate generation. + +Runs the same prompt against each model and compares response time, +throughput, candidate count, and refusal behavior. + +Usage: + python tools/ollama_benchmark.py # defaults + python tools/ollama_benchmark.py mistral phi3 # specific models + python tools/ollama_benchmark.py --prompt "custom prompt" + python tools/ollama_benchmark.py --output results.json + python tools/ollama_benchmark.py --num-ctx 2048 8192 32768 # compare context sizes +""" + +import argparse +import json +import os +import re +import time +import urllib.error +import urllib.request + +DEFAULT_MODELS = ["llama3.2", "mistral", "phi3", "gemma2", "qwen2.5"] + +DEFAULT_PROMPT = ( + "Generate baseword to be used in a denylist for keeping users from " + "setting their passwords with these basewords. We are protecting the " + "employees of Acme Corp in the technology industry located in Austin, TX.\n\n" + "Include variations with:\n" + "- Company name variations and abbreviations\n" + "- Common password patterns (Season+Year, Name+Numbers)\n" + "- Keyboard walks and common substitutions (@ for a, 3 for e, etc.)\n" + "- Location-based words and local references\n" + "- Industry-specific terminology\n" + "Output ONLY the passwords, one per line, no numbering or explanation." +) + +DEFAULT_NUM_CTX = [2048, 8192, 32768] +TIMEOUT = 600 + + +def pull_model(url, model): + """Pull an Ollama model. Returns True on success, False on failure.""" + print(f" Model '{model}' not found locally. Pulling...") + pull_url = f"{url}/api/pull" + payload = json.dumps({"name": model, "stream": True}).encode("utf-8") + req = urllib.request.Request( + pull_url, + data=payload, + headers={"Content-Type": "application/json"}, + ) + try: + with urllib.request.urlopen(req) as resp: + for raw_line in resp: + line = raw_line.decode("utf-8", errors="replace").strip() + if not line: + continue + try: + data = json.loads(line) + except json.JSONDecodeError: + continue + status = data.get("status") + if status: + print(f" {status}") + except (urllib.error.HTTPError, urllib.error.URLError, Exception) as e: + print(f" Error pulling model: {e}") + return False + print(f" Successfully pulled '{model}'.") + return True + + +def filter_candidates(response_text): + """Filter raw LLM response into usable password candidates.""" + raw_lines = response_text.strip().split("\n") + candidates = [] + for line in raw_lines: + stripped = line.strip() + if not stripped: + continue + cleaned = re.sub(r"^\d+[.)]\s*", "", stripped) + cleaned = re.sub(r"^[-*]\s*", "", cleaned) + cleaned = cleaned.strip() + if cleaned and len(cleaned) <= 128: + candidates.append(cleaned) + return candidates + + +def benchmark_model(url, model, prompt, num_ctx): + """Run the prompt against a single model at a given context size. Returns a results dict.""" + api_url = f"{url}/api/generate" + payload = json.dumps({ + "model": model, + "prompt": prompt, + "stream": False, + "options": {"num_ctx": num_ctx}, + }).encode("utf-8") + + result = { + "model": model, + "num_ctx": num_ctx, + "response_time_s": None, + "tokens_per_sec": None, + "candidate_count": 0, + "unique_candidates": 0, + "refusal": False, + "error": None, + } + + req = urllib.request.Request( + api_url, + data=payload, + headers={"Content-Type": "application/json"}, + ) + + start = time.monotonic() + try: + with urllib.request.urlopen(req, timeout=TIMEOUT) as resp: + body = json.loads(resp.read().decode("utf-8")) + except urllib.error.HTTPError as e: + if e.code == 404: + if not pull_model(url, model): + result["error"] = f"could not pull model" + return result + # Retry after pull + req = urllib.request.Request( + api_url, + data=payload, + headers={"Content-Type": "application/json"}, + ) + start = time.monotonic() + try: + with urllib.request.urlopen(req, timeout=TIMEOUT) as resp: + body = json.loads(resp.read().decode("utf-8")) + except Exception as retry_err: + result["error"] = str(retry_err) + return result + else: + result["error"] = f"HTTP {e.code}" + return result + except (urllib.error.URLError, Exception) as e: + result["error"] = str(e) + return result + + elapsed = time.monotonic() - start + result["response_time_s"] = round(elapsed, 2) + + # Extract tokens/sec from Ollama response metadata + eval_count = body.get("eval_count") + eval_duration = body.get("eval_duration") # nanoseconds + if eval_count and eval_duration and eval_duration > 0: + result["tokens_per_sec"] = round(eval_count / (eval_duration / 1e9), 2) + + response_text = body.get("response", "") + + # Refusal detection + if "I'm sorry" in response_text or "I can't help with that" in response_text: + result["refusal"] = True + + candidates = filter_candidates(response_text) + result["candidate_count"] = len(candidates) + result["unique_candidates"] = len(set(candidates)) + + return result + + +def print_table(results): + """Print a formatted comparison table.""" + headers = ["Model", "num_ctx", "Time (s)", "Tok/s", "Candidates", "Unique", "Refusal", "Error"] + rows = [] + for r in results: + rows.append([ + r["model"], + str(r["num_ctx"]), + str(r["response_time_s"] or "-"), + str(r["tokens_per_sec"] or "-"), + str(r["candidate_count"]), + str(r["unique_candidates"]), + "YES" if r["refusal"] else "no", + r["error"] or "", + ]) + + # Calculate column widths + col_widths = [len(h) for h in headers] + for row in rows: + for i, cell in enumerate(row): + col_widths[i] = max(col_widths[i], len(cell)) + + def fmt_row(cells): + return " ".join(cell.ljust(col_widths[i]) for i, cell in enumerate(cells)) + + print() + print(fmt_row(headers)) + print(" ".join("-" * w for w in col_widths)) + for row in rows: + print(fmt_row(row)) + print() + + +def main(): + parser = argparse.ArgumentParser( + description="Benchmark Ollama models for password candidate generation", + ) + parser.add_argument( + "models", + nargs="*", + default=DEFAULT_MODELS, + help=f"Models to benchmark (default: {' '.join(DEFAULT_MODELS)})", + ) + parser.add_argument( + "--prompt", + default=DEFAULT_PROMPT, + help="Override the generation prompt", + ) + parser.add_argument( + "--num-ctx", + nargs="+", + type=int, + default=DEFAULT_NUM_CTX, + metavar="N", + help=f"Context window sizes to test (default: {' '.join(str(n) for n in DEFAULT_NUM_CTX)})", + ) + parser.add_argument( + "--output", + metavar="FILE", + help="Write raw results to a JSON file", + ) + args = parser.parse_args() + + url = os.environ.get("OLLAMA_HOST", "http://localhost:11434").rstrip("/") + if not url.startswith("http"): + url = f"http://{url}" + + print(f"Ollama endpoint: {url}") + print(f"Models: {', '.join(args.models)}") + print(f"num_ctx values: {', '.join(str(n) for n in args.num_ctx)}") + print() + + results = [] + for model in args.models: + for num_ctx in args.num_ctx: + print(f"Benchmarking {model} (num_ctx={num_ctx})...") + r = benchmark_model(url, model, args.prompt, num_ctx) + if r["error"]: + print(f" Error: {r['error']}") + else: + print(f" {r['response_time_s']}s, {r['tokens_per_sec']} tok/s, " + f"{r['candidate_count']} candidates ({r['unique_candidates']} unique)" + f"{', REFUSED' if r['refusal'] else ''}") + results.append(r) + + print_table(results) + + if args.output: + with open(args.output, "w") as f: + json.dump(results, f, indent=2) + print(f"Results written to {args.output}") + + +if __name__ == "__main__": + main()