agentic-threat-hunting-framework 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
athf/commands/env.py ADDED
@@ -0,0 +1,373 @@
1
+ """Environment management commands."""
2
+
3
+ import subprocess # nosec B404
4
+ import sys
5
+ from pathlib import Path
6
+
7
+ import click
8
+ from rich.console import Console
9
+ from rich.panel import Panel
10
+
11
+ console = Console()
12
+
13
+ ENV_EPILOG = """
14
+ \b
15
+ Examples:
16
+ # Setup virtual environment with default Python
17
+ athf env setup
18
+
19
+ # Setup with specific Python version
20
+ athf env setup --python python3.13
21
+
22
+ # Include dev dependencies
23
+ athf env setup --dev
24
+
25
+ # Clean up existing venv and recreate
26
+ athf env setup --clean
27
+
28
+ \b
29
+ After setup:
30
+ # Activate venv (bash/zsh)
31
+ source .venv/bin/activate
32
+
33
+ # Deactivate
34
+ deactivate
35
+ """
36
+
37
+
38
+ @click.group(epilog=ENV_EPILOG)
39
+ def env() -> None:
40
+ """Manage Python virtual environment.
41
+
42
+ Commands for setting up, cleaning, and managing the Python
43
+ virtual environment for ATHF development.
44
+ """
45
+ pass
46
+
47
+
48
+ @env.command(name="setup")
49
+ @click.option(
50
+ "--python",
51
+ default="python3",
52
+ help="Python executable to use (default: python3)",
53
+ )
54
+ @click.option("--dev", is_flag=True, help="Install development dependencies")
55
+ @click.option("--clean", is_flag=True, help="Remove existing venv before creating")
56
+ def setup(python: str, dev: bool, clean: bool) -> None:
57
+ """Setup Python virtual environment with dependencies.
58
+
59
+ Creates .venv directory and installs athf package with
60
+ all dependencies from pyproject.toml.
61
+
62
+ \b
63
+ Steps:
64
+ 1. Create .venv directory (or clean existing)
65
+ 2. Install athf package in editable mode
66
+ 3. Install scikit-learn for semantic search
67
+ 4. Show activation instructions
68
+
69
+ \b
70
+ Examples:
71
+ athf env setup
72
+ athf env setup --python python3.13
73
+ athf env setup --dev
74
+ athf env setup --clean
75
+ """
76
+ venv_path = Path(".venv")
77
+
78
+ # Check if we're in the ATHF directory
79
+ if not Path("pyproject.toml").exists():
80
+ console.print("[red]Error: Not in ATHF directory (pyproject.toml not found)[/red]")
81
+ console.print("[dim]Run this command from the ATHF root directory[/dim]")
82
+ raise click.Abort()
83
+
84
+ # Clean existing venv if requested
85
+ if clean and venv_path.exists():
86
+ console.print("[yellow]🧹 Removing existing .venv directory...[/yellow]")
87
+ try:
88
+ import shutil
89
+
90
+ shutil.rmtree(venv_path)
91
+ console.print("[green]✅ Removed existing .venv[/green]\n")
92
+ except Exception as e:
93
+ console.print(f"[red]Error removing .venv: {e}[/red]")
94
+ raise click.Abort()
95
+
96
+ # Check if venv already exists
97
+ if venv_path.exists():
98
+ console.print("[yellow]⚠️ .venv already exists[/yellow]")
99
+ console.print("[dim]Use --clean to remove and recreate[/dim]\n")
100
+
101
+ # Show helpful usage instructions
102
+ if sys.platform == "win32":
103
+ activate_cmd = ".venv\\Scripts\\activate"
104
+ else:
105
+ activate_cmd = "source .venv/bin/activate"
106
+
107
+ usage_panel = Panel(
108
+ f"[bold cyan]To use the existing virtual environment:[/bold cyan]\n\n"
109
+ f"[green]1. Activate the venv:[/green]\n"
110
+ f" {activate_cmd}\n\n"
111
+ f"[green]2. Run athf commands:[/green]\n"
112
+ f" athf --version\n"
113
+ f" athf hunt --help\n\n"
114
+ f"[green]3. Or use without activating:[/green]\n"
115
+ f" .venv/bin/athf [command]\n\n"
116
+ f"[dim]💡 Your prompt will show (.venv) when activated[/dim]",
117
+ title="✨ Virtual Environment Ready",
118
+ border_style="cyan",
119
+ )
120
+ console.print(usage_panel)
121
+ raise click.Abort()
122
+
123
+ # Create virtual environment
124
+ console.print(f"[cyan]📦 Creating virtual environment with {python}...[/cyan]")
125
+ try:
126
+ subprocess.run(
127
+ [python, "-m", "venv", ".venv"],
128
+ check=True,
129
+ capture_output=True,
130
+ text=True,
131
+ )
132
+ console.print("[green]✅ Virtual environment created[/green]\n")
133
+ except subprocess.CalledProcessError as e:
134
+ console.print(f"[red]Error creating venv: {e.stderr}[/red]")
135
+ raise click.Abort()
136
+ except FileNotFoundError:
137
+ console.print(f"[red]Error: {python} not found[/red]")
138
+ console.print("[dim]Try: athf env setup --python python3.13[/dim]")
139
+ raise click.Abort()
140
+
141
+ # Determine pip path
142
+ if sys.platform == "win32":
143
+ pip_path = venv_path / "Scripts" / "pip"
144
+ else:
145
+ pip_path = venv_path / "bin" / "pip"
146
+
147
+ # Upgrade pip
148
+ console.print("[cyan]📦 Upgrading pip...[/cyan]")
149
+ try:
150
+ subprocess.run(
151
+ [str(pip_path), "install", "--upgrade", "pip"],
152
+ check=True,
153
+ capture_output=True,
154
+ text=True,
155
+ )
156
+ console.print("[green]✅ pip upgraded[/green]\n")
157
+ except subprocess.CalledProcessError as e:
158
+ console.print(f"[yellow]Warning: Failed to upgrade pip: {e.stderr}[/yellow]\n")
159
+
160
+ # Install athf package
161
+ console.print("[cyan]📦 Installing ATHF package...[/cyan]")
162
+ install_cmd = [str(pip_path), "install", "-e", "."]
163
+ if dev:
164
+ install_cmd.append("[dev]")
165
+
166
+ try:
167
+ subprocess.run(
168
+ install_cmd,
169
+ check=True,
170
+ capture_output=True,
171
+ text=True,
172
+ )
173
+ console.print("[green]✅ ATHF installed[/green]\n")
174
+ except subprocess.CalledProcessError as e:
175
+ console.print(f"[red]Error installing package: {e.stderr}[/red]")
176
+ raise click.Abort()
177
+
178
+ # Install scikit-learn for athf similar command
179
+ console.print("[cyan]📦 Installing scikit-learn for semantic search...[/cyan]")
180
+ try:
181
+ subprocess.run(
182
+ [str(pip_path), "install", "scikit-learn"],
183
+ check=True,
184
+ capture_output=True,
185
+ text=True,
186
+ )
187
+ console.print("[green]✅ scikit-learn installed[/green]\n")
188
+ except subprocess.CalledProcessError as e:
189
+ console.print(f"[yellow]Warning: Failed to install scikit-learn: {e.stderr}[/yellow]")
190
+ console.print("[dim]athf similar command will not work without scikit-learn[/dim]\n")
191
+
192
+ # Success message
193
+ console.print("[bold green]🎉 Virtual environment setup complete![/bold green]\n")
194
+
195
+ # Show activation instructions
196
+ if sys.platform == "win32":
197
+ activate_cmd = ".venv\\Scripts\\activate"
198
+ else:
199
+ activate_cmd = "source .venv/bin/activate"
200
+
201
+ activation_panel = Panel(
202
+ f"[cyan]{activate_cmd}[/cyan]\n\n"
203
+ f"[dim]Then verify installation:[/dim]\n"
204
+ f"[white]athf --version[/white]\n"
205
+ f"[white]athf hunt --help[/white]",
206
+ title="🚀 Next Steps",
207
+ border_style="green",
208
+ )
209
+ console.print(activation_panel)
210
+
211
+
212
+ @env.command(name="clean")
213
+ def clean() -> None:
214
+ """Remove virtual environment.
215
+
216
+ Deletes the .venv directory to start fresh.
217
+
218
+ \b
219
+ Example:
220
+ athf env clean
221
+ athf env setup
222
+ """
223
+ venv_path = Path(".venv")
224
+
225
+ if not venv_path.exists():
226
+ console.print("[yellow]No .venv directory found[/yellow]")
227
+ return
228
+
229
+ console.print("[yellow]🧹 Removing .venv directory...[/yellow]")
230
+ try:
231
+ import shutil
232
+
233
+ shutil.rmtree(venv_path)
234
+ console.print("[green]✅ Virtual environment removed[/green]")
235
+ console.print("[dim]Run 'athf env setup' to recreate[/dim]")
236
+ except Exception as e:
237
+ console.print(f"[red]Error removing .venv: {e}[/red]")
238
+ raise click.Abort()
239
+
240
+
241
+ @env.command(name="info")
242
+ def info() -> None:
243
+ """Show virtual environment information.
244
+
245
+ Display Python version, installed packages, and venv location.
246
+
247
+ \b
248
+ Example:
249
+ athf env info
250
+ """
251
+ venv_path = Path(".venv")
252
+
253
+ if not venv_path.exists():
254
+ console.print("[yellow]No .venv directory found[/yellow]")
255
+ console.print("[dim]Run 'athf env setup' to create[/dim]")
256
+ return
257
+
258
+ # Determine python path
259
+ if sys.platform == "win32":
260
+ python_path = venv_path / "Scripts" / "python"
261
+ else:
262
+ python_path = venv_path / "bin" / "python"
263
+
264
+ if not python_path.exists():
265
+ console.print("[red]Error: Virtual environment appears corrupted[/red]")
266
+ console.print("[dim]Run 'athf env setup --clean' to recreate[/dim]")
267
+ return
268
+
269
+ # Get Python version
270
+ try:
271
+ result = subprocess.run(
272
+ [str(python_path), "--version"],
273
+ check=True,
274
+ capture_output=True,
275
+ text=True,
276
+ )
277
+ python_version = result.stdout.strip()
278
+ except subprocess.CalledProcessError:
279
+ python_version = "Unknown"
280
+
281
+ # Get installed packages count
282
+ pip_path = python_path.parent / "pip"
283
+ package_count: int | str
284
+ try:
285
+ result = subprocess.run(
286
+ [str(pip_path), "list", "--format", "freeze"],
287
+ check=True,
288
+ capture_output=True,
289
+ text=True,
290
+ )
291
+ package_count = len(result.stdout.strip().split("\n"))
292
+ except subprocess.CalledProcessError:
293
+ package_count = "Unknown"
294
+
295
+ # Check for athf installation
296
+ try:
297
+ result = subprocess.run(
298
+ [str(pip_path), "show", "agentic-threat-hunting-framework"],
299
+ check=True,
300
+ capture_output=True,
301
+ text=True,
302
+ )
303
+ athf_installed = "✅ Installed" if result.returncode == 0 else "❌ Not installed"
304
+ except subprocess.CalledProcessError:
305
+ athf_installed = "❌ Not installed"
306
+
307
+ # Check for scikit-learn
308
+ try:
309
+ result = subprocess.run(
310
+ [str(pip_path), "show", "scikit-learn"],
311
+ check=True,
312
+ capture_output=True,
313
+ text=True,
314
+ )
315
+ sklearn_installed = "✅ Installed" if result.returncode == 0 else "❌ Not installed"
316
+ except subprocess.CalledProcessError:
317
+ sklearn_installed = "❌ Not installed"
318
+
319
+ # Display info
320
+ console.print("\n[bold]Virtual Environment Info:[/bold]\n")
321
+ console.print(f" [cyan]Location:[/cyan] {venv_path.absolute()}")
322
+ console.print(f" [cyan]Python:[/cyan] {python_version}")
323
+ console.print(f" [cyan]Packages:[/cyan] {package_count} installed")
324
+ console.print(f" [cyan]athf:[/cyan] {athf_installed}")
325
+ console.print(f" [cyan]scikit-learn:[/cyan] {sklearn_installed} [dim](required for athf similar)[/dim]")
326
+ console.print()
327
+
328
+
329
+ @env.command(name="activate")
330
+ def activate() -> None:
331
+ """Show command to activate virtual environment.
332
+
333
+ Note: Cannot activate directly (subprocesses can't modify parent shell).
334
+ Copy and run the printed command to activate.
335
+
336
+ \b
337
+ Example:
338
+ athf env activate
339
+ # Then copy and run the printed command
340
+ """
341
+ venv_path = Path(".venv")
342
+
343
+ if not venv_path.exists():
344
+ console.print("[yellow]No .venv directory found[/yellow]")
345
+ console.print("[dim]Run 'athf env setup' to create[/dim]")
346
+ raise click.Abort()
347
+
348
+ # Determine activation command based on platform
349
+ if sys.platform == "win32":
350
+ activate_cmd = ".venv\\Scripts\\activate"
351
+ else:
352
+ activate_cmd = "source .venv/bin/activate"
353
+
354
+ console.print("\n[bold cyan]To activate the virtual environment, run:[/bold cyan]\n")
355
+ console.print(f" [green]{activate_cmd}[/green]\n")
356
+ console.print("[dim]💡 Tip: Copy the command above and run it in your shell[/dim]\n")
357
+
358
+
359
+ @env.command(name="deactivate")
360
+ def deactivate_cmd() -> None:
361
+ """Show command to deactivate virtual environment.
362
+
363
+ Note: Cannot deactivate directly (subprocesses can't modify parent shell).
364
+ Copy and run the printed command to deactivate.
365
+
366
+ \b
367
+ Example:
368
+ athf env deactivate
369
+ # Then copy and run the printed command
370
+ """
371
+ console.print("\n[bold cyan]To deactivate the virtual environment, run:[/bold cyan]\n")
372
+ console.print(" [green]deactivate[/green]\n")
373
+ console.print("[dim]💡 This will return you to your system Python[/dim]\n")
athf/commands/hunt.py CHANGED
@@ -502,21 +502,57 @@ def search(query: str) -> None:
502
502
  console.print()
503
503
 
504
504
 
505
+ def _render_progress_bar(covered: int, total: int, width: int = 20) -> str:
506
+ """Render a visual progress bar with filled and empty blocks.
507
+
508
+ Args:
509
+ covered: Number of covered techniques
510
+ total: Total number of techniques
511
+ width: Width of the progress bar in characters
512
+
513
+ Returns:
514
+ ASCII progress bar string using simple characters
515
+ """
516
+ if total == 0:
517
+ return "·" * width
518
+
519
+ # Cap percentage at 100% for visual display
520
+ percentage = min(covered / total, 1.0)
521
+ filled = int(percentage * width)
522
+ empty = width - filled
523
+
524
+ # Use simple characters that render reliably
525
+ filled_char = "■"
526
+ empty_char = "·"
527
+
528
+ return filled_char * filled + empty_char * empty
529
+
530
+
505
531
  @hunt.command()
506
- def coverage() -> None:
532
+ @click.option("--detailed", is_flag=True, help="Show detailed technique coverage with hunt references")
533
+ def coverage(detailed: bool) -> None:
507
534
  """Show MITRE ATT&CK technique coverage across hunts.
508
535
 
509
536
  \b
510
537
  Analyzes and displays:
511
- Which tactics you've hunted (e.g., Persistence, Credential Access)
512
- Which techniques per tactic
513
- Coverage gaps (tactics with few/no hunts)
514
- Hunt distribution across the ATT&CK matrix
538
+ Hunt count per tactic across all 14 ATT&CK tactics
539
+ Technique count per tactic (with caveats - see note below)
540
+ Overall unique technique coverage across all hunts
541
+ Detailed technique-to-hunt mapping (with --detailed)
515
542
 
516
543
  \b
517
- Example:
544
+ Examples:
545
+ # Show coverage overview
518
546
  athf hunt coverage
519
547
 
548
+ # Show detailed technique mapping
549
+ athf hunt coverage --detailed
550
+
551
+ \b
552
+ Note on technique counts:
553
+ Per-tactic technique counts may include duplicates if hunts cover
554
+ multiple tactics. The overall unique technique count (bottom) is accurate.
555
+
520
556
  \b
521
557
  Use this to:
522
558
  • Identify blind spots in your hunting program
@@ -527,23 +563,64 @@ def coverage() -> None:
527
563
 
528
564
  \b
529
565
  Pro tip:
530
- Combine with threat intel to focus on attacker-relevant TTPs.
531
- Example: "Which persistence techniques are we NOT hunting?"
566
+ Focus on tactics with no coverage that align with your threat model.
567
+ Use --detailed to see which specific techniques each hunt covers.
532
568
  """
569
+ from athf.core.attack_matrix import ATTACK_TACTICS, get_sorted_tactics
570
+
533
571
  manager = HuntManager()
534
572
  coverage = manager.calculate_attack_coverage()
535
573
 
536
- if not coverage:
574
+ if not coverage or not coverage.get("by_tactic"):
537
575
  console.print("[yellow]No hunt coverage data available.[/yellow]")
538
576
  return
539
577
 
540
- console.print("\n[bold cyan]🎯 MITRE ATT&CK Coverage[/bold cyan]\n")
578
+ summary = coverage["summary"]
579
+ by_tactic = coverage["by_tactic"]
541
580
 
542
- for tactic, techniques in sorted(coverage.items()):
543
- console.print(f"[bold]{tactic.title()}[/bold] ({len(techniques)} techniques)")
544
- for technique in techniques:
545
- console.print(f" • {technique}")
546
- console.print()
581
+ # Display title
582
+ console.print("\n[bold]MITRE ATT&CK Coverage[/bold]")
583
+ console.print("─" * 60 + "\n")
584
+
585
+ # Display all tactics in ATT&CK order with hunt counts
586
+ for tactic_key in get_sorted_tactics():
587
+ data = by_tactic.get(tactic_key, {})
588
+ tactic_name = ATTACK_TACTICS[tactic_key]["name"]
589
+
590
+ hunt_count = data.get("hunt_count", 0)
591
+ techniques_covered = data.get("techniques_covered", 0)
592
+
593
+ # Format: "Tactic Name 2 hunts, 7 techniques"
594
+ if hunt_count > 0:
595
+ console.print(f"{tactic_name:<24} {hunt_count} hunts, {techniques_covered} techniques")
596
+ else:
597
+ console.print(f"{tactic_name:<24} [dim]no coverage[/dim]")
598
+
599
+ # Display overall coverage
600
+ console.print(
601
+ f"\n[bold]Overall: {summary['unique_techniques']}/{summary['total_techniques']} techniques ({summary['overall_coverage_pct']:.0f}%)[/bold]\n"
602
+ )
603
+
604
+ # Display detailed technique coverage if requested
605
+ if detailed:
606
+ console.print("\n[bold cyan]🔍 Detailed Technique Coverage[/bold cyan]\n")
607
+
608
+ for tactic_key in get_sorted_tactics():
609
+ data = by_tactic.get(tactic_key, {})
610
+ if data.get("hunt_count", 0) == 0:
611
+ continue # Skip tactics with no hunts in detailed view
612
+
613
+ tactic_name = ATTACK_TACTICS[tactic_key]["name"]
614
+ console.print(
615
+ f"\n[bold]{tactic_name}[/bold] ({data['hunt_count']} hunts, {len(data['techniques'])} unique techniques)"
616
+ )
617
+
618
+ # Show techniques with hunt references
619
+ for technique, hunt_ids in sorted(data["techniques"].items()):
620
+ hunt_refs = ", ".join(sorted(set(hunt_ids))) # Remove duplicates and sort
621
+ console.print(f" • [yellow]{technique}[/yellow] - {hunt_refs}")
622
+
623
+ console.print()
547
624
 
548
625
 
549
626
  @hunt.command(hidden=True)