OpenHound/src/openhound/cli/saved_search.py at 30886fe53af8fe3d090fe223e2294c110be7f5d9 · SpecterOps/OpenHound · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
from enum import Enum
from pathlib import Path
from typing import Annotated

import typer
from rich.console import Console

from openhound.core.models.saved_search import QueryBundle
from openhound.core.progress import Progress
from openhound.core.saved_searches import SavedSearches, Strategy

saved_searches = typer.Typer(help="Saved searches management")


class Format(str, Enum):
    json = "json"
    yaml = "yaml"


class OutputFormat(str, Enum):
    json = "json"
    zip = "zip"


@saved_searches.command(help="Upload saved searches to BloodHound")
def upload(
        path: Annotated[
            Path,
            typer.Argument(
                exists=True,
                file_okay=False,
                dir_okay=True,
                readable=True,
                resolve_path=True,
                help="Directory where saved searches are located",
            ),
        ],
        file_format: Format = typer.Option(
            default=Format.json,
            help="File format of the saved searches (json or yaml)",
        ),
        strategy: Strategy = typer.Option(
            default=Strategy.skip,
            help="Skip or overwrite saved search if the name already exists",
        ),
        progress: Progress = typer.Option(
            Progress.tqdm, help="Select progress tracker option"
        ),
):
    search_files = (
        path.rglob("**/*.json")
        if file_format == Format.json
        else path.rglob("**/*.yaml")
    )
    pipeline = SavedSearches(progress=progress, strategy=strategy)
    results = pipeline.run(list(search_files), file_format=file_format)
    return results


@saved_searches.command(help="Create a single saved searches json/zip bundle")
def bundle(
        path: Annotated[
            Path,
            typer.Argument(
                exists=True,
                file_okay=False,
                dir_okay=True,
                readable=True,
                resolve_path=True,
                help="Directory where saved searches are located",
            ),
        ],
        output_path: Annotated[typer.FileTextWrite, typer.Argument(
            help="Output file path for the generated saved-search bundle (including filename)")],
        file_format: Format = typer.Option(
            default=Format.json,
            help="File format of the saved searches (json or yaml)",
        ),
        output_format: OutputFormat = typer.Option(
            default=OutputFormat.json,
            help="File format for the saved searches bundle (json or zip)",
        )
):
    search_files = list(
        path.rglob("**/*.json")
        if file_format == Format.json
        else path.rglob("**/*.yaml")
    )

    bundle_object = QueryBundle.from_paths(search_files, file_format=file_format)
    bundle_object.save(output_path, output_format=output_format)

    console = Console()
    console.print("[bold green]Saved-search bundle created[/bold green]")
    console.print(f"[bold magenta]Saved searches:[/bold magenta] {len(bundle_object.queries)}")
    console.print(
        f"[bold magenta]Output path:[/bold magenta] [italic]{Path(output_path.name).resolve()}[/italic]"
    )