libQueryTpot.py

import logging
from datetime import datetime
from elasticsearch import Elasticsearch, exceptions
from collections import Counter
from typing import List, Dict
from tqdm import tqdm

# Configure logging
logging.basicConfig(
    filename='./pullTPOTDataForOTX.log',
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

class TPotQuery:
    """
    A class to query TPOT Elasticsearch data, excluding noisy sources,
    and collect recent logs for further processing (e.g., STIX bundle creation).
    """

    def __init__(self, es_host, index_name, api_key_id, api_key, batch_size=1000, scroll_time='2m', config=None):
        """
        Initialize the TPotQuery class.

        Args:
            es_host (str): Elasticsearch server hostname or IP.
            index_name (str): Elasticsearch index to query.
            api_key_id (str): Elasticsearch API key ID.
            api_key (str): Elasticsearch API key value.
            batch_size (int): Number of documents to fetch per batch.
            scroll_time (str): Scroll context retention duration.
            config (dict, optional): Additional configuration if needed.
        """
        self.index_name = index_name
        self.batch_size = batch_size
        self.scroll_time = scroll_time
        self.config = config
        self.ignore_types = config['tpot']['ignore_types']
        self.results = []

        try:
            self.es = Elasticsearch(
                hosts=[es_host],
                api_key=(api_key_id, api_key),
                verify_certs=False
            )
        except exceptions.ConnectionError as e:
            logging.error(f"Elasticsearch connection error: {e}")
            raise

    def pull_recent_logs(self, minutes_back=5):
        """
        Pull logs from Elasticsearch within the past `minutes_back` minutes,
        excluding types in the ignore list.

        Args:
            minutes_back (int): Number of minutes to look back from current time.

        Returns:
            list: All matching documents retrieved from Elasticsearch.
        """
        time_range = {
            "range": {
                "@timestamp": {
                    "gte": f"now-{minutes_back}m",
                    "lte": "now"
                }
            }
        }

        exclude_types = {
            "terms": {
                "type": self.ignore_types
            }
        }

        query_body = {
            "query": {
                "bool": {
                    "must": [time_range],
                    "must_not": [exclude_types]
                }
            }
        }

        logging.info(f"Querying Elasticsearch index '{self.index_name}' for recent data...")
        try:
            response = self.es.search(
                index=self.index_name,
                body=query_body,
                scroll=self.scroll_time,
                size=self.batch_size
            )

            scroll_id = response.get('_scroll_id')
            hits = response['hits']['hits']

            while hits:
                self.results.extend(hits)
                response = self.es.scroll(scroll_id=scroll_id, scroll=self.scroll_time)
                scroll_id = response.get('_scroll_id')
                hits = response['hits']['hits']

            self.es.clear_scroll(scroll_id=scroll_id)
            logging.info(f"Retrieved {len(self.results)} records from '{self.index_name}'.")

        except Exception as e:
            logging.error("Failed to retrieve logs from Elasticsearch", exc_info=True)

        return self.results

    def show_results(self):
        count=1
        for item in self.results:
            print (f'[{count}] {item["_source"]['type']} {item["_source"]['src_ip']} : {item["_source"]}')
            count += 1

    def summarize_type_field(self, entries: List[Dict]) -> Dict[str, int]:
        """
        Summarizes the count of each unique 'type' field found within the '_source' key
        of a list of Elasticsearch-style entries.

        Args:
            entries (List[Dict]): List of log entries with '_source' containing the 'type' field.

        Returns:
            Dict[str, int]: Dictionary mapping each unique type to its count.
        """
        type_counts = Counter()

        for entry in entries:
            source = entry.get('_source', {})
            type_value = source.get('type', 'MISSING')
            type_counts[type_value] += 1

        print(f"Total unique types: {len(type_counts)}")
        for t, count in type_counts.items():
            print(f"Type: {t:<20} Count: {count}")

        return dict(type_counts)

    def summarize_credentials(self, entries):
        """
        Summarizes username/password combinations from a list of log entries.

        Args:
            entries (list): List of Elasticsearch-style dicts with '_source' keys.

        Returns:
            dict: Dictionary with (username, password) tuples as keys and their occurrence counts as values.
        """
        cred_counter = Counter()

        for entry in entries:
            source = entry.get('_source', {})
            username = source.get('username')
            password = source.get('password')

            if username is not None and password is not None:
                cred_counter[(username, password)] += 1

        print(f"Total unique username/password pairs: {len(cred_counter)}")
        for (user, pwd), count in cred_counter.items():
            print(f"Username: {user:<20} Password: {pwd:<20} Count: {count}")

        return dict(cred_counter)

    def summarize_hashes(self, entries):
        """
        Summarizes unique file hashes (from the 'shasum' field) in the list of log entries.

        Args:
            entries (list): List of Elasticsearch-style dicts with '_source' keys.

        Returns:
            dict: Dictionary with each unique shasum and its occurrence count.
        """
        hash_counter = Counter()

        for entry in entries:
            source = entry.get('_source', {})
            shasum = source.get('shasum')

            if shasum:
                hash_counter[shasum] += 1

        print(f"Total unique file hashes (shasum): {len(hash_counter)}")
        for shasum, count in hash_counter.items():
            print(f"SHA Sum: {shasum:<64} Count: {count}")

        return dict(hash_counter)

    from collections import Counter

    def summarize_inputs(self, entries):
        """
        Summarizes the count of each unique value in the 'input' field from the log entries.

        Args:
            entries (list): List of Elasticsearch-style dicts with '_source' keys.

        Returns:
            dict: Dictionary with each unique input value and its count.
        """
        input_counter = Counter()

        for entry in entries:
            source = entry.get('_source', {})
            input_value = source.get('input')

            if input_value:
                input_counter[input_value] += 1

        print(f"Total unique input values: {len(input_counter)}")
        for value, count in input_counter.items():
            print(f"Input: {value:<20} Count: {count}")

        return dict(input_counter)