main.py

#!/usr/bin/env python3
"""
OhMyMalware Security Data Ingestion Script for Elasticsearch

This script follows the specific instructions from the ohmymalware repository:
1. Creates separate indices for events (logs-endpoint.events.imported) and alerts (logs-endpoint.alerts.imported)
2. Uses mapping.total_fields.limit of 6000 as specified in README
3. Downloads and uses the official mapping.json from the repository
4. Creates ingest pipeline with event.ingested field for Elastic Security rule compatibility
5. Processes all Event Data including Episode directories
6. Handles data stream creation if required by Elasticsearch templates
7. Routes documents to correct indices based on content type

Note: Modern Elasticsearch versions may automatically create data streams for logs-* indices.
The script handles this automatically and creates either a regular index or data stream as needed.

This implementation matches the exact requirements from the ohmymalware README for triggering
Elastic Security detections and using the correct index patterns.

Downloads data from jamesspi/ohmymalware GitHub repository and indexes it properly.
"""

import csv
import json
import logging
from datetime import datetime
from typing import Any, Dict, List

import pandas as pd
import requests
from elasticsearch import Elasticsearch, helpers

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class SecurityDataIngester:
    def __init__(
        self,
        es_host: str = "localhost",
        es_port: int = 9200,
        es_user: str = None,
        es_password: str = None,
    ):
        """Initialize Elasticsearch connection"""
        if es_user and es_password:
            self.es = Elasticsearch(
                [f"http://{es_host}:{es_port}"], basic_auth=(es_user, es_password)
            )
        else:
            self.es = Elasticsearch([f"http://{es_host}:{es_port}"])

        self.base_url = (
            "https://raw.githubusercontent.com/jamesspi/ohmymalware/main/Event%20Data"
        )
        self.pipeline_name = "ohmymalware-security-pipeline"
        self.events_index = "logs-endpoint.events.imported"
        self.alerts_index = "logs-endpoint.alerts-imported"

    def create_ingest_pipeline(self):
        """Create ingest pipeline for Elastic Security rule compatibility"""
        pipeline_config = {
            "description": "Ingest pipeline for ohmymalware data to trigger Elastic Security rules",
            "processors": [
                {
                    "date": {
                        "field": "@timestamp",
                        "formats": ["ISO8601"],
                        "output_format": "yyyy-MM-dd'T'HH:mm:ss.SSSSSSSSSXXX",
                    }
                },
                {"set": {"field": "event.ingested", "value": "{{_ingest.timestamp}}"}},
            ],
        }

        try:
            self.es.ingest.put_pipeline(id=self.pipeline_name, body=pipeline_config)
            logger.info(f"Created ingest pipeline: {self.pipeline_name}")
        except Exception as e:
            logger.error(f"Error creating ingest pipeline: {e}")
            raise

    def create_indices_and_pipeline(self):
        """Create Elasticsearch indices/data streams and ingest pipeline following ohmymalware repository instructions"""

        # First create the ingest pipeline
        self.create_ingest_pipeline()

        try:
            # Download the mapping.json file from the repository
            mapping_url = "https://raw.githubusercontent.com/jamesspi/ohmymalware/main/Event%20Data/mapping.json"
            response = requests.get(mapping_url)
            response.raise_for_status()

            # Parse the mapping from the repository
            repo_mapping = json.loads(response.text)
            logger.info("Successfully downloaded mapping.json from repository")

        except Exception as e:
            logger.warning(
                f"Could not download mapping.json: {e}. Using fallback mapping."
            )
            # Fallback to basic mapping if mapping.json is not available
            repo_mapping = {
                "properties": {
                    "@timestamp": {"type": "date"},
                    "event_type": {"type": "keyword"},
                    "source_file": {"type": "keyword"},
                    "episode": {"type": "keyword"},
                    "raw_data": {"type": "text"},
                    "parsed_data": {"type": "object", "enabled": True},
                    "ingestion_timestamp": {"type": "date"},
                    "event": {"type": "object", "enabled": True},
                }
            }

        # Create both indices (events and alerts)
        for index_name in [self.events_index, self.alerts_index]:
            try:
                # Check if index or data stream already exists
                if self.es.indices.exists(index=index_name):
                    logger.info(f"Index/data stream {index_name} already exists")
                    continue

                # Try creating as regular index first
                try:
                    index_config = {
                        "mappings": repo_mapping,
                        "settings": {
                            "number_of_shards": 1,
                            "number_of_replicas": 0,
                            "mapping.total_fields.limit": "6000",
                            "index.default_pipeline": self.pipeline_name,
                        },
                    }

                    self.es.indices.create(index=index_name, body=index_config)
                    logger.info(f"Created index: {index_name}")

                except Exception as index_error:
                    if "data streams only" in str(index_error):
                        # Create as data stream instead
                        logger.info(
                            f"Index template requires data stream for {index_name}. Creating data stream..."
                        )

                        try:
                            # Create index template for the data stream
                            template_name = f"{index_name}-template"
                            template_config = {
                                "index_patterns": [index_name],
                                "data_stream": {},
                                "template": {
                                    "mappings": repo_mapping,
                                    "settings": {
                                        "mapping.total_fields.limit": "6000",
                                        "index.default_pipeline": self.pipeline_name,
                                    },
                                },
                            }

                            # Create the index template
                            self.es.indices.put_index_template(
                                name=template_name, body=template_config
                            )
                            logger.info(f"Created index template: {template_name}")

                            # Create the data stream
                            self.es.indices.create_data_stream(name=index_name)
                            logger.info(f"Created data stream: {index_name}")

                        except AttributeError:
                            # Fallback for older Elasticsearch versions that don't support data streams
                            logger.warning(
                                "Data streams not supported in this Elasticsearch version"
                            )
                            # Try with a different index name
                            fallback_index = f"security-{index_name.replace('logs-', '').replace('.', '-')}"
                            logger.info(
                                f"Trying with alternative index name: {fallback_index}"
                            )

                            index_config = {
                                "mappings": repo_mapping,
                                "settings": {
                                    "number_of_shards": 1,
                                    "number_of_replicas": 0,
                                    "mapping.total_fields.limit": "6000",
                                    "index.default_pipeline": self.pipeline_name,
                                },
                            }
                            self.es.indices.create(
                                index=fallback_index, body=index_config
                            )
                            logger.info(
                                f"Created index with alternative name: {fallback_index}"
                            )

                            # Update index name for this fallback
                            if index_name == self.events_index:
                                self.events_index = fallback_index
                            else:
                                self.alerts_index = fallback_index

                        except Exception as ds_error:
                            logger.error(
                                f"Could not create data stream for {index_name}: {ds_error}"
                            )
                            # Final fallback with different name
                            fallback_index = f"security-{index_name.replace('logs-', '').replace('.', '-')}"
                            logger.info(
                                f"Final fallback with index name: {fallback_index}"
                            )

                            index_config = {
                                "mappings": repo_mapping,
                                "settings": {
                                    "number_of_shards": 1,
                                    "number_of_replicas": 0,
                                    "mapping.total_fields.limit": "6000",
                                    "index.default_pipeline": self.pipeline_name,
                                },
                            }
                            self.es.indices.create(
                                index=fallback_index, body=index_config
                            )
                            logger.info(
                                f"Created index with fallback name: {fallback_index}"
                            )

                            # Update index name for this fallback
                            if index_name == self.events_index:
                                self.events_index = fallback_index
                            else:
                                self.alerts_index = fallback_index
                    else:
                        raise index_error

            except Exception as e:
                logger.error(f"Error creating index/data stream {index_name}: {e}")
                raise

    def get_github_files(self) -> List[Dict[str, str]]:
        """Get list of files from GitHub repository Event Data directory including episodes"""
        api_url = (
            "https://api.github.com/repos/jamesspi/ohmymalware/contents/Event%20Data"
        )

        try:
            response = requests.get(api_url)
            response.raise_for_status()
            items = response.json()

            file_list = []

            # Files to skip (not actual data)
            skip_files = {"README.md", "mapping.json", ".gitignore"}

            for item in items:
                if item["type"] == "file" and item["name"] not in skip_files:
                    file_list.append(
                        {"name": item["name"], "path": item["name"], "episode": None}
                    )
                elif item["type"] == "dir" and "Episode" in item["name"]:
                    # Handle episode directories
                    episode_name = item["name"]
                    episode_api_url = f"https://api.github.com/repos/jamesspi/ohmymalware/contents/Event%20Data/{episode_name}"

                    try:
                        episode_response = requests.get(episode_api_url)
                        episode_response.raise_for_status()
                        episode_files = episode_response.json()

                        for episode_file in episode_files:
                            if (
                                episode_file["type"] == "file"
                                and episode_file["name"] not in skip_files
                            ):
                                file_list.append(
                                    {
                                        "name": episode_file["name"],
                                        "path": f"{episode_name}/{episode_file['name']}",
                                        "episode": episode_name,
                                    }
                                )
                    except Exception as e:
                        logger.warning(f"Could not access episode {episode_name}: {e}")

            logger.info(
                f"Found {len(file_list)} data files across episodes and main directory"
            )
            return file_list

        except Exception as e:
            logger.error(f"Error fetching file list: {e}")
            return []

    def download_file(self, file_path: str) -> str:
        """Download file content from GitHub"""
        url = f"{self.base_url}/{file_path}"

        try:
            response = requests.get(url)
            response.raise_for_status()

            # Handle compressed files
            if file_path.endswith(".gz"):
                import gzip

                content = gzip.decompress(response.content).decode("utf-8")
                logger.info(f"Downloaded and decompressed: {file_path}")
            else:
                content = response.text
                logger.info(f"Downloaded: {file_path}")

            return content
        except Exception as e:
            logger.error(f"Error downloading {file_path}: {e}")
            return None

    def parse_file_content(
        self, content: str, filename: str, episode: str = None
    ) -> List[Dict[str, Any]]:
        """Parse file content based on file type"""
        documents = []
        file_ext = filename.lower().split(".")[-1]

        # Handle compressed files
        if filename.endswith(".gz"):
            file_ext = filename.lower().split(".")[-2]  # Get extension before .gz

        try:
            if file_ext == "json":
                data = json.loads(content)
                if isinstance(data, list):
                    for item in data:
                        documents.append(
                            self.create_document(item, filename, "json", episode)
                        )
                else:
                    documents.append(
                        self.create_document(data, filename, "json", episode)
                    )

            elif file_ext == "ndjson":
                # Handle newline-delimited JSON
                for line_num, line in enumerate(content.strip().split("\n")):
                    if line.strip():
                        try:
                            data = json.loads(line.strip())
                            documents.append(
                                self.create_document(data, filename, "ndjson", episode)
                            )
                        except json.JSONDecodeError as e:
                            logger.warning(
                                f"Invalid JSON on line {line_num + 1} in {filename}: {e}"
                            )

            elif file_ext == "csv":
                from io import StringIO

                csv_data = csv.DictReader(StringIO(content))
                for row in csv_data:
                    documents.append(
                        self.create_document(dict(row), filename, "csv", episode)
                    )

            elif file_ext in ["txt", "log"]:
                lines = content.strip().split("\n")
                for i, line in enumerate(lines):
                    if line.strip():  # Skip empty lines
                        documents.append(
                            self.create_document(
                                {"line_number": i + 1, "content": line},
                                filename,
                                "text",
                                episode,
                            )
                        )
            else:
                # Handle other formats as raw text
                documents.append(
                    self.create_document({"content": content}, filename, "raw", episode)
                )

        except Exception as e:
            logger.error(f"Error parsing {filename}: {e}")
            # Fallback: treat as raw text
            documents.append(
                self.create_document(
                    {"content": content, "parse_error": str(e)},
                    filename,
                    "raw",
                    episode,
                )
            )

        return documents

    def create_document(
        self, data: Dict[str, Any], filename: str, data_type: str, episode: str = None
    ) -> Dict[str, Any]:
        """Create standardized document for Elasticsearch following ohmymalware format"""
        current_time = datetime.now().isoformat()

        # Create raw_data with size limit to prevent immense term errors
        try:
            if isinstance(data, dict):
                raw_data_str = json.dumps(data)
                # Limit raw_data size to prevent immense term errors (Elasticsearch limit is 32766 bytes)
                if len(raw_data_str.encode("utf-8")) > 32000:
                    raw_data_str = raw_data_str[:16000] + "... [truncated due to size]"
            else:
                raw_data_str = str(data)
                if len(raw_data_str.encode("utf-8")) > 32000:
                    raw_data_str = raw_data_str[:16000] + "... [truncated due to size]"
        except Exception as e:
            raw_data_str = f"Error serializing data: {str(e)}"

        doc = {
            "@timestamp": current_time,
            "event_type": data_type,
            "source_file": filename,
            "parsed_data": data,
            "raw_data": raw_data_str,
            "ingestion_timestamp": current_time,
        }

        # Add episode information if available
        if episode:
            doc["episode"] = episode

        # Try to extract timestamp from data if available
        timestamp_fields = [
            "timestamp",
            "time",
            "datetime",
            "date",
            "@timestamp",
            "event_time",
            "TimeCreated",
        ]
        for field in timestamp_fields:
            if field in data and data[field]:
                try:
                    # Try to parse various timestamp formats
                    if isinstance(data[field], str):
                        parsed_time = pd.to_datetime(data[field])
                        doc["@timestamp"] = parsed_time.isoformat()
                        break
                except Exception as e:
                    logger.debug(f"Could not parse timestamp field {field}: {e}")
                    continue

        # If the data already has proper structure, merge it with our document
        if isinstance(data, dict):
            # Don't overwrite our standardized fields and skip conflicting data stream fields
            reserved_fields = {
                "@timestamp",
                "event_type",
                "source_file",
                "parsed_data",
                "raw_data",
                "ingestion_timestamp",
                "episode",
                "data_stream",
            }

            for key, value in data.items():
                if key not in reserved_fields:
                    doc[key] = value

        return doc

    def determine_target_index(self, filename: str, data: Dict[str, Any] = None) -> str:
        """Determine whether document should go to events or alerts index"""
        # Check filename patterns
        if "alert" in filename.lower():
            return self.alerts_index

        # Check document content for alert indicators
        if data and isinstance(data, dict):
            # Look for alert-specific fields
            alert_indicators = [
                "kibana.alert",
                "signal",
                "alert",
                "detection",
                "rule.name",
                "rule.id",
                "rule.type",
            ]

            for indicator in alert_indicators:
                if indicator in data or any(
                    indicator in str(key) for key in data.keys()
                ):
                    return self.alerts_index

        # Default to events index
        return self.events_index

    def bulk_index_documents(
        self, documents: List[Dict[str, Any]], filename: str
    ) -> int:
        """Bulk index documents to Elasticsearch with detailed error logging"""
        if not documents:
            return 0

        # Group documents by target index
        events_docs = []
        alerts_docs = []

        for doc in documents:
            target_index = self.determine_target_index(filename, doc)
            if target_index == self.alerts_index:
                alerts_docs.append(doc)
            else:
                events_docs.append(doc)

        total_success = 0

        # Index events
        if events_docs:
            success_count = self._bulk_index_to_target(events_docs, self.events_index)
            total_success += success_count
            logger.info(f"Indexed {success_count} documents to events index")

        # Index alerts
        if alerts_docs:
            success_count = self._bulk_index_to_target(alerts_docs, self.alerts_index)
            total_success += success_count
            logger.info(f"Indexed {success_count} documents to alerts index")

        return total_success

    def _bulk_index_to_target(
        self, documents: List[Dict[str, Any]], target_index: str
    ) -> int:
        """Helper method to bulk index documents to a specific target index"""
        actions = []
        for doc in documents:
            action = {
                "_op_type": "create",  # Required for data streams
                "_index": target_index,
                "_source": doc,
            }
            actions.append(action)

        try:
            # Use bulk with detailed error reporting
            response = helpers.bulk(
                self.es,
                actions,
                chunk_size=500,
                request_timeout=60,
                raise_on_error=False,
                raise_on_exception=False,
            )

            success_count = response[0]
            failed_items = response[1]

            if failed_items:
                logger.error(
                    f"Failed to index {len(failed_items)} documents to {target_index}. First 3 errors:"
                )
                for i, failed_item in enumerate(failed_items[:3]):
                    # Extract the actual error from the response
                    if "create" in failed_item:
                        error_info = failed_item["create"].get("error", {})
                        error_type = error_info.get("type", "unknown")
                        error_reason = error_info.get("reason", "no reason provided")
                        logger.error(
                            f"  Error {i+1}: Type: {error_type}, Reason: {error_reason}"
                        )

                        # Show document that failed if it helps
                        if "caused_by" in error_info:
                            caused_by = error_info["caused_by"]
                            logger.error(
                                f"    Caused by: {caused_by.get('type', 'unknown')} - {caused_by.get('reason', 'no details')}"
                            )
                    else:
                        logger.error(f"  Error {i+1}: {failed_item}")

                if len(failed_items) > 3:
                    logger.error(f"  ... and {len(failed_items) - 3} more errors")

            return success_count

        except Exception as e:
            logger.error(f"Error during bulk indexing to {target_index}: {e}")
            # Try individual indexing to get better error messages
            if len(documents) <= 5:  # Only for small batches
                logger.info(
                    "Trying individual document indexing for better error messages..."
                )
                success_count = 0
                for i, doc in enumerate(documents):
                    try:
                        result = self.es.create(
                            index=target_index, body=doc
                        )  # Use create instead of index
                        success_count += 1
                    except Exception as doc_error:
                        logger.error(f"Document {i+1} failed: {doc_error}")
                return success_count
            return 0

    def run_ingestion(self):
        """Main ingestion process"""
        logger.info("Starting ohmymalware security data ingestion...")

        # Create index
        self.create_indices_and_pipeline()

        # Get file list
        files = self.get_github_files()
        if not files:
            logger.error("No files found to process")
            return

        total_documents = 0

        # Process each file
        for file_info in files:
            filename = file_info["name"]
            file_path = file_info["path"]
            episode = file_info["episode"]

            logger.info(
                f"Processing: {file_path}"
                + (f" (Episode: {episode})" if episode else "")
            )

            # Download file
            content = self.download_file(file_path)
            if not content:
                continue

            # Parse content
            documents = self.parse_file_content(content, filename, episode)

            # Index documents
            indexed_count = self.bulk_index_documents(documents, filename)
            total_documents += indexed_count

            logger.info(f"Processed {file_path}: {indexed_count} documents indexed")

        logger.info(f"Ingestion complete! Total documents indexed: {total_documents}")
        logger.info(
            f"Data ingested into: Events: {self.events_index}, Alerts: {self.alerts_index}"
        )
        logger.info(
            "Note: Conflicting 'data_stream' fields were filtered out to avoid mapping conflicts"
        )

        # Refresh indices
        for index in [self.events_index, self.alerts_index]:
            try:
                self.es.indices.refresh(index=index)
                logger.info(f"Index/Data stream refreshed: {index}")
            except Exception as e:
                logger.warning(f"Could not refresh index {index}: {e}")

        logger.info("Indices refreshed and ready for use")


def main():
    """Main function - configure your Elasticsearch connection here"""

    # Configuration - MODIFY THESE VALUES
    ES_HOST = "localhost"  # Your Elasticsearch host
    ES_PORT = 9200  # Your Elasticsearch port
    ES_USER = "elastic"  # Your Elasticsearch username (if auth required)
    ES_PASSWORD = "changeme"  # Your Elasticsearch password (if auth required)

    try:
        ingester = SecurityDataIngester(
            es_host=ES_HOST, es_port=ES_PORT, es_user=ES_USER, es_password=ES_PASSWORD
        )

        # Test connection
        if not ingester.es.ping():
            logger.error("Cannot connect to Elasticsearch")
            return

        logger.info("Connected to Elasticsearch successfully")

        # Check for conflicting templates
        try:
            templates = ingester.es.indices.get_index_template()
            logs_templates = [
                t
                for t in templates.get("index_templates", [])
                if any(
                    "logs" in pattern
                    for pattern in t.get("index_template", {}).get("index_patterns", [])
                )
            ]
            if logs_templates:
                logger.info(
                    f"Found {len(logs_templates)} index templates that might affect logs-* indices"
                )
        except Exception as e:
            logger.warning(f"Could not check index templates: {e}")

        # Run ingestion
        ingester.run_ingestion()

        # Final summary
        logger.info("=" * 60)
        logger.info("INGESTION COMPLETE - READY FOR ELASTIC SECURITY")
        logger.info("=" * 60)
        logger.info(f"Events Index: {ingester.events_index}")
        logger.info(f"Alerts Index: {ingester.alerts_index}")
        logger.info(f"Pipeline: {ingester.pipeline_name}")
        logger.info("The data is now compatible with Elastic Security rules!")
        logger.info('In Kibana, filter data with: _index:"logs-endpoint.*"')
        logger.info("=" * 60)

    except Exception as e:
        logger.error(f"Fatal error: {e}")


if __name__ == "__main__":
    main()