sigma-rules/hunting/markdown.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.

from pathlib import Path
import click
from .definitions import ATLAS_URL, ATTACK_URL, STATIC_INTEGRATION_LINK_MAP, Hunt
from .utils import load_index_file, load_toml, save_index_file, validate_link


class MarkdownGenerator:
    """Class to generate or update Markdown documentation from TOML or YAML files."""
    def __init__(self, base_path: Path):
        """Initialize with the base path and load the hunting index."""
        self.base_path = base_path
        self.hunting_index = load_index_file()

    def process_file(self, file_path: Path) -> None:
        """Process a single TOML file and generate its Markdown representation."""
        if not file_path.is_file() or file_path.suffix != '.toml':
            raise ValueError(f"The provided path is not a valid TOML file: {file_path}")

        click.echo(f"Processing specific TOML file: {file_path}")
        hunt_config = load_toml(file_path)
        markdown_content = self.convert_toml_to_markdown(hunt_config, file_path)

        docs_folder = self.create_docs_folder(file_path)
        markdown_path = docs_folder / f"{file_path.stem}.md"
        self.save_markdown(markdown_path, markdown_content)

        self.update_or_add_entry(hunt_config, file_path)

    def process_folder(self, folder: str) -> None:
        """Process all TOML files in a specified folder and generate their Markdown representations."""
        folder_path = self.base_path / folder / "queries"
        docs_folder = self.base_path / folder / "docs"

        if not folder_path.is_dir() or not docs_folder.is_dir():
            raise ValueError(f"Queries folder {folder_path} or docs folder {docs_folder} does not exist.")

        click.echo(f"Processing all TOML files in folder: {folder_path}")
        toml_files = folder_path.rglob("*.toml")

        for toml_file in toml_files:
            self.process_file(toml_file)

    def process_all_files(self) -> None:
        """Process all TOML files in the base directory and subfolders."""
        click.echo("Processing all TOML files in the base directory and subfolders.")
        toml_files = self.base_path.rglob("queries/*.toml")

        for toml_file in toml_files:
            self.process_file(toml_file)

    def convert_toml_to_markdown(self, hunt_config: Hunt, file_path: Path) -> str:
        """Convert a Hunt configuration to Markdown format."""
        markdown = f"# {hunt_config.name}\n\n---\n\n"
        markdown += "## Metadata\n\n"
        markdown += f"- **Author:** {hunt_config.author}\n"
        markdown += f"- **Description:** {hunt_config.description}\n"
        markdown += f"- **UUID:** `{hunt_config.uuid}`\n"
        markdown += f"- **Integration:** {', '.join(self.generate_integration_links(hunt_config.integration))}\n"
        markdown += f"- **Language:** `{hunt_config.language}`\n".replace("'", "").replace('"', "")
        markdown += f"- **Source File:** [{hunt_config.name}]({(Path('../queries') / file_path.name).as_posix()})\n"
        markdown += "\n## Query\n\n"
        for query in hunt_config.query:
            markdown += f"```sql\n{query}```\n\n"

        if hunt_config.notes:
            markdown += "## Notes\n\n" + "\n".join(f"- {note}" for note in hunt_config.notes)
        if hunt_config.mitre:
            markdown += "\n\n## MITRE ATT&CK Techniques\n\n" + "\n".join(
                f"- [{tech}]({ATLAS_URL if tech.startswith('AML') else ATTACK_URL}"
                f"{tech.replace('.', '/') if tech.startswith('T') else tech})"
                for tech in hunt_config.mitre
            )
        if hunt_config.references:
            markdown += "\n\n## References\n\n" + "\n".join(f"- {ref}" for ref in hunt_config.references)

        markdown += f"\n\n## License\n\n- `{hunt_config.license}`\n"
        return markdown

    def save_markdown(self, markdown_path: Path, content: str) -> None:
        """Save the Markdown content to a file."""
        markdown_path.write_text(content, encoding="utf-8")
        click.echo(f"Markdown generated: {markdown_path}")

    def update_or_add_entry(self, hunt_config: Hunt, toml_path: Path) -> None:
        """Update or add the entry for a TOML file in the hunting index."""
        folder_name = toml_path.parent.parent.name
        uuid = hunt_config.uuid

        entry = {
            'name': hunt_config.name,
            'path': f"./{toml_path.relative_to(self.base_path).as_posix()}",
            'mitre': hunt_config.mitre
        }

        if folder_name not in self.hunting_index:
            self.hunting_index[folder_name] = {uuid: entry}
        else:
            self.hunting_index[folder_name][uuid] = entry

        save_index_file(self.base_path, self.hunting_index)

    def create_docs_folder(self, file_path: Path) -> Path:
        """Create the docs folder if it doesn't exist and return the path."""
        docs_folder = file_path.parent.parent / "docs"
        docs_folder.mkdir(parents=True, exist_ok=True)
        return docs_folder

    def generate_integration_links(self, integrations: list[str]) -> list[str]:
        """Generate integration links for the documentation."""
        base_url = 'https://docs.elastic.co/integrations'
        generated = []
        for integration in integrations:
            if integration in STATIC_INTEGRATION_LINK_MAP:
                link_str = STATIC_INTEGRATION_LINK_MAP[integration]
            else:
                link_str = integration.replace('.', '/')
            link = f'{base_url}/{link_str}'
            validate_link(link)
            generated.append(f'[{integration}]({link})')
        return generated

    def update_index_md(self) -> None:
        """Update the index.md file based on the entries in index.yml."""
        index_file = self.base_path / "index.yml"
        index_content = "# List of Available Queries\n\nHere are the queries currently available:\n"

        if not index_file.exists():
            click.echo(f"No index.yml found at {index_file}. Skipping index.md update.")
            return

        for folder, files in sorted(self.hunting_index.items()):
            index_content += f"\n\n## {folder}\n"
            for file_info in sorted(files.values(), key=lambda x: x['name']):
                md_path = file_info['path'].replace('queries', 'docs').replace('.toml', '.md')
                index_content += f"- [{file_info['name']}]({md_path}) (ES|QL)\n"

        index_md_path = self.base_path / "index.md"
        index_md_path.write_text(index_content, encoding="utf-8")
        click.echo(f"Index Markdown updated at: {index_md_path}")