feat(llm): add wordlist (denylist) mode to LLM attack menu

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Justin Bollinger
2026-07-24 12:08:20 -04:00
co-authored by Claude Sonnet 4.6
parent d6da721414
commit a4da571ab4
2 changed files with 49 additions and 18 deletions
+23 -11
View File
@@ -513,33 +513,45 @@ def bandrel_method(ctx: Any) -> None:
def ollama_attack(ctx: Any) -> None:
_notify.prompt_notify_for_attack("LLM")
print("\n\tLLM Attack")
print("\t1. Target info (company / industry / location)")
print("\t2. Wordlist (generate basewords from a sample wordlist)")
choice = input("\n\tSelect generation mode: ").strip()
if choice == "1":
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)
ctx.hcatOllama(
ctx.hcatHashType,
ctx.hcatHashFile,
"target",
{"company": company, "industry": industry, "location": location},
)
elif choice == "2":
path = _omen_pick_training_wordlist(ctx, title="LLM Sample Wordlists")
if not path:
return
ctx.hcatOllama(ctx.hcatHashType, ctx.hcatHashFile, "wordlist", path)
else:
print("\t[!] Invalid selection.")
def _omen_pick_training_wordlist(ctx: Any):
"""Show wordlist picker for OMEN training. Returns path or None."""
def _omen_pick_training_wordlist(ctx: Any, title: str = "Training Wordlists"):
"""Show wordlist picker. Returns path or None."""
wordlist_files = ctx.list_wordlist_files(ctx.hcatWordlists)
if wordlist_files:
entries = [f"{i}) {f}" for i, f in enumerate(wordlist_files, start=1)]
max_len = max((len(e) for e in entries), default=24)
print_multicolumn_list(
"Training Wordlists",
title,
entries,
min_col_width=max_len,
max_col_width=max_len,
)
print("\tp. Enter a custom path")
sel = input("\n\tSelect wordlist for training: ").strip()
sel = input("\n\tSelect wordlist: ").strip()
if sel.lower() == "p":
path = input("\n\tPath to training wordlist: ").strip()
path = input("\n\tPath to wordlist: ").strip()
return path if path else None
try:
idx = int(sel)
+23 -4
View File
@@ -329,7 +329,7 @@ class TestOllamaAttack:
def test_calls_hcatOllama_with_context(self) -> None:
ctx = _make_ctx()
with patch("builtins.input", side_effect=["ACME", "tech", "NYC"]):
with patch("builtins.input", side_effect=["1", "ACME", "tech", "NYC"]):
ollama_attack(ctx)
ctx.hcatOllama.assert_called_once_with(
@@ -342,7 +342,7 @@ class TestOllamaAttack:
def test_passes_hash_type_and_file(self) -> None:
ctx = _make_ctx(hash_type="1800", hash_file="/tmp/sha512.txt")
with patch("builtins.input", side_effect=["Corp", "finance", "London"]):
with patch("builtins.input", side_effect=["1", "Corp", "finance", "London"]):
ollama_attack(ctx)
call_args = ctx.hcatOllama.call_args[0]
@@ -352,7 +352,7 @@ class TestOllamaAttack:
def test_strips_whitespace_from_inputs(self) -> None:
ctx = _make_ctx()
with patch("builtins.input", side_effect=[" ACME ", " tech ", " NYC "]):
with patch("builtins.input", side_effect=["1", " ACME ", " tech ", " NYC "]):
ollama_attack(ctx)
target_info = ctx.hcatOllama.call_args[0][3]
@@ -363,7 +363,26 @@ class TestOllamaAttack:
def test_target_string_is_literal_target(self) -> None:
ctx = _make_ctx()
with patch("builtins.input", side_effect=["X", "Y", "Z"]):
with patch("builtins.input", side_effect=["1", "X", "Y", "Z"]):
ollama_attack(ctx)
assert ctx.hcatOllama.call_args[0][2] == "target"
def test_wordlist_mode_calls_hcatOllama_with_path(self) -> None:
ctx = _make_ctx()
ctx.list_wordlist_files.return_value = ["rockyou.txt"]
ctx.hcatWordlists = "/tmp/wl"
# mode "2", then pick wordlist "1"
with patch("builtins.input", side_effect=["2", "1"]):
ollama_attack(ctx)
args = ctx.hcatOllama.call_args[0]
assert args[2] == "wordlist"
assert args[3].endswith("rockyou.txt")
def test_invalid_mode_does_not_call_hcatOllama(self) -> None:
ctx = _make_ctx()
with patch("builtins.input", side_effect=["9"]):
ollama_attack(ctx)
ctx.hcatOllama.assert_not_called()