hunting/__main__.py

# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
# or more contributor license agreements. Licensed under the Elastic License
# 2.0; you may not use this file except in compliance with the Elastic License
# 2.0.

import json
import textwrap
from collections import Counter
from dataclasses import asdict
from pathlib import Path

import click
from tabulate import tabulate

from detection_rules.misc import parse_user_config

from .definitions import HUNTING_DIR
from .markdown import MarkdownGenerator
from .run import QueryRunner
from .search import QueryIndex
from .utils import (filter_elasticsearch_params, get_hunt_path, load_all_toml,
                    load_toml, update_index_yml)


@click.group()
def hunting():
    """Commands for managing hunting queries and converting TOML to Markdown."""
    pass


@hunting.command('generate-markdown')
@click.argument('path', required=False)
def generate_markdown(path: Path = None):
    """Convert TOML hunting queries to Markdown format."""
    markdown_generator = MarkdownGenerator(HUNTING_DIR)

    if path:
        path = Path(path)
        if path.is_file() and path.suffix == '.toml':
            click.echo(f"Generating Markdown for single file: {path}")
            markdown_generator.process_file(path)
        elif (HUNTING_DIR / path).is_dir():
            click.echo(f"Generating Markdown for folder: {path}")
            markdown_generator.process_folder(path)
        else:
            raise ValueError(f"Invalid path provided: {path}")
    else:
        click.echo("Generating Markdown for all files.")
        markdown_generator.process_all_files()

    # After processing, update the index
    markdown_generator.update_index_md()


@hunting.command('refresh-index')
def refresh_index():
    """Refresh the index.yml file from TOML files and then refresh the index.md file."""
    click.echo("Refreshing the index.yml and index.md files.")
    update_index_yml(HUNTING_DIR)
    markdown_generator = MarkdownGenerator(HUNTING_DIR)
    markdown_generator.update_index_md()
    click.echo("Index refresh complete.")


@hunting.command('search')
@click.option('--tactic', type=str, default=None, help="Search by MITRE tactic ID (e.g., TA0001)")
@click.option('--technique', type=str, default=None, help="Search by MITRE technique ID (e.g., T1078)")
@click.option('--sub-technique', type=str, default=None, help="Search by MITRE sub-technique ID (e.g., T1078.001)")
@click.option('--data-source', type=str, default=None, help="Filter by data_source like 'aws', 'macos', or 'linux'")
@click.option('--keyword', type=str, default=None, help="Search by keyword in name, description, and notes")
def search_queries(tactic: str, technique: str, sub_technique: str, data_source: str, keyword: str):
    """Search for queries based on MITRE tactic, technique, sub-technique, or data_source."""

    if not any([tactic, technique, sub_technique, data_source, keyword]):
        raise click.UsageError("""Please provide at least one filter (tactic, technique, sub-technique,
                               data_source or keyword) to search queries.""")

    click.echo("Searching for queries based on provided filters...")

    # Create an instance of the QueryIndex class
    query_index = QueryIndex(HUNTING_DIR)

    # Filter out None values from the MITRE filter tuple
    mitre_filters = tuple(filter(None, (tactic, technique, sub_technique)))

    # Call the search method of QueryIndex with the provided MITRE filters, data_source, and keyword
    results = query_index.search(mitre_filter=mitre_filters, data_source=data_source, keyword=keyword)

    if results:
        click.secho(f"\nFound {len(results)} matching queries:\n", fg="green", bold=True)

        # Prepare the data for tabulate
        table_data = []
        for result in results:
            # Customize output to include technique, data_source, and UUID
            data_source_str = result['data_source']
            mitre_str = ", ".join(result['mitre'])
            uuid = result['uuid']
            table_data.append([result['name'], uuid, result['path'], data_source_str, mitre_str])

        # Output results using tabulate
        table_headers = ["Name", "UUID", "Location", "Data Source", "MITRE"]
        click.echo(tabulate(table_data, headers=table_headers, tablefmt="fancy_grid"))

    else:
        click.secho("No matching queries found.", fg="red", bold=True)


@hunting.command('view-hunt')
@click.option('--uuid', type=str, help="View a specific hunt by UUID.")
@click.option('--path', type=str, help="View a specific hunt by file path.")
@click.option('--format', 'output_format', default='toml', type=click.Choice(['toml', 'json'], case_sensitive=False),
              help="Output format (toml or json).")
@click.option('--query-only', is_flag=True, help="Only display the query content.")
def view_hunt(uuid: str, path: str, output_format: str, query_only: bool):
    """View a specific hunt by UUID or file path in the specified format (TOML or JSON)."""

    # Get the hunt path or error message
    hunt_path, error_message = get_hunt_path(uuid, path)

    if error_message:
        raise click.ClickException(error_message)

    # Load the TOML data
    hunt = load_toml(hunt_path)

    # Handle query-only option
    if query_only:
        click.secho("Available queries:", fg="blue", bold=True)
        # Format queries for display using tabulate and textwrap
        table_data = [(i, textwrap.fill(query, width=120)) for i, query in enumerate(hunt.query)]
        table_headers = ["Query"]
        click.echo(tabulate(table_data, headers=table_headers, tablefmt="fancy_grid"))
        return

    # Output the hunt in the requested format
    if output_format == 'toml':
        click.echo(hunt_path.read_text())
    elif output_format == 'json':
        hunt_dict = asdict(hunt)
        click.echo(json.dumps(hunt_dict, indent=4))


@hunting.command('hunt-summary')
@click.option('--breakdown', type=click.Choice(['platform', 'integration', 'language'],
                                               case_sensitive=False), default='platform',
              help="Specify how to break down the summary: 'platform', 'integration', or 'language'.")
def hunt_summary(breakdown: str):
    """
    Generate a summary of hunt queries, broken down by platform, integration, or language.
    """
    click.echo(f"Generating hunt summary broken down by {breakdown}...")

    # Load all hunt queries
    all_hunts = load_all_toml(HUNTING_DIR)

    # Use Counter for more concise counting
    platform_counter = Counter()
    integration_counter = Counter()
    language_counter = Counter()

    for hunt, path in all_hunts:
        # Get the platform based on the folder name
        platform = path.parent.parent.stem
        platform_counter[platform] += 1

        # Count integrations
        integration_counter.update(hunt.integration)

        # Count languages, renaming 'SQL' to 'OSQuery'
        languages = ['OSQuery' if lang == 'SQL' else lang for lang in hunt.language]
        language_counter.update(languages)

    # Prepare and display the table based on the selected breakdown
    if breakdown == 'platform':
        table_data = [[platform, count] for platform, count in platform_counter.items()]
        table_headers = ["Platform (Folder)", "Hunt Count"]
    elif breakdown == 'integration':
        table_data = [[integration, count] for integration, count in integration_counter.items()]
        table_headers = ["Integration", "Hunt Count"]
    elif breakdown == 'language':
        table_data = [[language, count] for language, count in language_counter.items()]
        table_headers = ["Language", "Hunt Count"]

    click.echo(tabulate(table_data, headers=table_headers, tablefmt="fancy_grid"))


@hunting.command('run-query')
@click.option('--uuid', help="The UUID of the hunting query to run.")
@click.option('--file-path', help="The file path of the hunting query to run.")
@click.option('--all', 'run_all', is_flag=True, help="Run all eligible queries in the file.")
@click.option('--wait-time', 'wait_time', default=180, help="Time to wait for query completion.")
def run_query(uuid: str, file_path: str, run_all: bool, wait_time: int):
    """Run a hunting query by UUID or file path. Only ES|QL queries are supported."""

    # Get the hunt path or error message
    hunt_path, error_message = get_hunt_path(uuid, file_path)

    if error_message:
        click.echo(error_message)
        return

    # Load the user configuration
    config = parse_user_config()
    if not config:
        click.secho("No configuration found. Please add a `detection-rules-cfg` file.", fg="red", bold=True)
        return

    es_config = filter_elasticsearch_params(config)

    # Create a QueryRunner instance
    query_runner = QueryRunner(es_config)

    # Load the hunting data
    hunting_data = query_runner.load_hunting_file(hunt_path)

    # Display description
    wrapped_description = textwrap.fill(hunting_data.description, width=120)
    click.secho("\nHunting Description:", fg="blue", bold=True)
    click.secho(f"\n{wrapped_description}\n", bold=True)

    # Extract eligible queries
    eligible_queries = {i: query for i, query in enumerate(hunting_data.query) if "from" in query}
    if not eligible_queries:
        click.secho("No eligible queries found in the file.", fg="red", bold=True)
        return

    if run_all:
        # Run all eligible queries if the --all flag is set
        query_runner.run_all_queries(eligible_queries, wait_time)
        return

    # Display available queries
    click.secho("Available queries:", fg="blue", bold=True)
    for i, query in eligible_queries.items():
        click.secho(f"\nQuery {i + 1}:", fg="green", bold=True)
        click.echo(query_runner._format_query(query))
        click.secho("\n" + "-" * 120, fg="yellow")

    # Handle query selection
    while True:
        try:
            query_number = click.prompt("Enter the query number", type=int)
            if query_number - 1 in eligible_queries:
                selected_query = eligible_queries[query_number - 1]
                break
            else:
                click.secho(f"Invalid query number: {query_number}. Please try again.", fg="yellow")
        except ValueError:
            click.secho("Please enter a valid number.", fg="yellow")

    # Run the selected query
    query_runner.run_individual_query(selected_query, wait_time)


if __name__ == "__main__":
    hunting()