indicator-enrichment

import demistomock as demisto
from CommonServerPython import *
from AggregatedCommandApiModule import *
from enum import Enum
from typing import NamedTuple, Any, Optional
from collections import defaultdict
import traceback


MAX_INDICATORS = 100

MSG_NO_VALID_INDICATORS = (
    "No valid indicators provided. You must provide at least one valid indicator " "in arguments: indicator_list or text."
)

ERR_LIMIT_TEMPLATE = (
    "Error: Indicator limit exceeded. Found {found} indicators, but the limit is {limit}. "
    "To process all indicators, set the ignore_indicator_limit argument to true. "
    "Note: This should be avoided if possible."
)

MSG_FATAL_EXECUTION_ERROR = "Fatal error while executing enrichment. View the logs or try again.\n{error}"


class ValidationError(Exception):
    """
    Raised when business logic constraints are violated.
    This exception should result in a return_error action.
    """


class AllExecutionsFailedError(Exception):
    """
    Raised when none of the executions of the underlying enrichment scripts have succeeded.
    """


class GracefulExit(Exception):
    """
    Raised when the script should stop execution and return a status message.
    This exception should result in a return_results action (success status), not an error.
    """


class IndicatorType(Enum):
    """
    Configuration registry for supported indicator types.
    Acts as the single source of truth for command mapping, argument keys, and context prefixes.
    """

    IP = ("ip_list", "IP", "ip-enrichment", "IPEnrichment", ["ip", "ipv4"])
    URL = ("url_list", "URL", "url-enrichment", "URLEnrichment", ["url"])
    DOMAIN = ("domain_list", "Domain", "domain-enrichment", "DomainEnrichment", ["domain"])
    CVE = ("cve_list", "CVE", "cve-enrichment", "CVEEnrichment", ["cve"])
    FILE = ("file_hash", "File", "file-enrichment", "FileEnrichment", ["file"])

    def __init__(self, argument_key: str, display_name: str, command_name: str, context_prefix: str, aliases: list[str]):
        """
        Initializes the IndicatorType configuration.

        Args:
            argument_key: The argument name expected by the child command (e.g., 'ip_list').
            display_name: The human-readable name for headers (e.g., 'IP').
            command_name: The XSOAR command to execute (e.g., 'ip-enrichment').
            context_prefix: The prefix used in EntryContext (e.g., 'IPEnrichment').
            aliases: A list of string representations used for type detection.
        """
        self.argument_key = argument_key
        self.display_name = display_name
        self.command_name = command_name
        self.context_prefix = context_prefix
        self.aliases = aliases

    @classmethod
    def resolve_from_string(cls, value: str) -> Optional["IndicatorType"]:
        """
        Maps a raw string type (e.g., 'IPv4', 'sha256') to the corresponding IndicatorType Enum.

        Args:
            value: The raw type string to resolve.

        Returns:
            The matching IndicatorType member, or None if the type is not supported.
        """
        if not value:
            return None
        normalized_value = value.lower()
        for member in cls:
            if normalized_value in member.aliases:
                return member
        return None


class UnsupportedIndicator(NamedTuple):
    """
    Represents an indicator that was detected but is not supported by the configured Enrichment commands.
    """

    type: str
    value: str


class EnrichmentTask(NamedTuple):
    """
    Represents a unit of work: a specific XSOAR command configured for a specific indicator type.
    """

    indicator_type: IndicatorType
    command: Command


class EnrichmentRequest:
    """
    Immutable Data Transfer Object (DTO) representing the sanitized, deduplicated,
    and validated input data ready for processing.
    """

    def __init__(
        self,
        valid_indicators_by_type: defaultdict[IndicatorType, list[str]],
        unsupported_items: list[UnsupportedIndicator],
        unknown_items: list[str],
        duplicates_removed_count: int,
        sub_command_arguments: dict[str, Any],
        include_raw_context: bool,
    ):
        """
        Initializes the EnrichmentRequest.

        Args:
            valid_indicators_by_type: A mapping of IndicatorType to a list of valid indicator values.
            unsupported_items: A list of items that have a type but are not supported by this script.
            unknown_items: A list of items where type detection failed.
            duplicates_removed_count: The count of duplicate items filtered out during building.
            sub_command_arguments: Arguments to be passed down to child commands (e.g., brands).
            include_raw_context: Whether to include the raw enrichment keys in the final output.
        """
        self.valid_indicators_by_type = valid_indicators_by_type
        self.unsupported_items = unsupported_items
        self.unknown_items = unknown_items
        self.duplicates_removed_count = duplicates_removed_count
        self.sub_command_arguments = sub_command_arguments
        self.include_raw_context = include_raw_context

    @property
    def total_valid_count(self) -> int:
        """Returns the total number of valid indicators across all types."""
        return sum(len(indicators) for indicators in self.valid_indicators_by_type.values())


class EnrichmentResult:
    """
    Accumulator for the results of the enrichment execution phase.
    """

    def __init__(self):
        self.enriched_data: list[ContextResult] = []
        self.raw_context: ContextResult = {}
        self.markdown_sections: list[str] = []
        self.execution_errors: List[str] = []


class EnrichmentRequestBuilder:
    """
    Responsible for parsing raw XSOAR arguments, extracting indicators from text and lists,
    deduplicating inputs, and enforcing business validation rules.
    """

    def __init__(self, raw_arguments: dict[str, Any]):
        """
        Args:
            raw_arguments: The dictionary of arguments provided to the command (demisto.args()).
        """
        self.raw_arguments = raw_arguments
        # this is for the deduplication tracking
        self._seen_indicators: set[str] = set()

        self.valid_indicators_by_type: defaultdict[IndicatorType, list[str]] = defaultdict(list)
        self.unsupported_items: list[UnsupportedIndicator] = []
        self.unknown_items: list[str] = []
        self.duplicates_count: int = 0

    def build_and_validate(self) -> EnrichmentRequest:
        """
        Orchestrates the creation of an EnrichmentRequest.

        Returns:
            A fully populated and validated EnrichmentRequest.

        Raises:
            ValidationError: If inputs are invalid or limits are exceeded.
            GracefulExit: If valid inputs are provided but result in no actionable work.
        """
        self._extract_indicators_from_text_argument()
        self._extract_indicators_from_list_argument()
        self._enforce_enrichment_arg_rules()

        return EnrichmentRequest(
            valid_indicators_by_type=self.valid_indicators_by_type,
            unsupported_items=self.unsupported_items,
            unknown_items=self.unknown_items,
            duplicates_removed_count=self.duplicates_count,
            sub_command_arguments=self._filter_passthrough_arguments(),
            include_raw_context=argToBoolean(self.raw_arguments.get("raw_context", False)),
        )

    def _classify_and_store(self, raw_type: str, value: str) -> None:
        """
        Normalizes, deduplicates, and classifies a single indicator value.
        Populates the internal state buckets (valid, unknown, unsupported).

        Args:
            raw_type: The string type representation (e.g., 'IP', 'Unknown').
            value: The indicator value string.
        """
        if not value:
            return

        clean_value = str(value).strip()
        if not clean_value:
            return

        if clean_value in self._seen_indicators:
            self.duplicates_count += 1
            return

        self._seen_indicators.add(clean_value)

        indicator_type = IndicatorType.resolve_from_string(raw_type)

        if indicator_type:
            self.valid_indicators_by_type[indicator_type].append(clean_value)
        elif not raw_type or raw_type.lower() == "unknown":
            self.unknown_items.append(clean_value)
        else:
            self.unsupported_items.append(UnsupportedIndicator(raw_type, clean_value))

    def _extract_indicators_from_text_argument(self) -> None:
        """
        Extracts indicators from the 'text' argument using the 'extractIndicators' command.

        Example of Response of the extractIndicators command:
        {
        "IP": ["1.1.1.1"],
        "Domain": ["google.com"]
        }
        """
        text_argument = self.raw_arguments.get("text", "")
        if not text_argument:
            return

        res = execute_command("extractIndicators", {"text": text_argument}, extract_contents=False)
        if not res or not isinstance(res, list):
            return

        entry_context = res[0].get("EntryContext", {}) or {}
        extracted_indicators = entry_context.get("ExtractedIndicators", {}) or {}

        for type_str, values in extracted_indicators.items():
            items = values if isinstance(values, list) else [values]
            for val in items:
                self._classify_and_store(type_str, val)

    def _extract_indicators_from_list_argument(self) -> None:
        """
        Parses the 'indicator_list' argument and attempts to auto-detect the type of each item.
        """
        raw_list = argToList(self.raw_arguments.get("indicator_list"))
        for item in raw_list:
            detected_type = auto_detect_indicator_type(item)  # type: ignore
            self._classify_and_store(detected_type or "Unknown", item)

    def _enforce_enrichment_arg_rules(self) -> None:
        """
        Validates the state of collected indicators against business rules.

        Raises:
            ValidationError: If:
             - indicator_list and text are not given
             - indicator_list given with no valid indicators
             - Over 100 indicators given (after clearing invalid and duplicates)
            GracefulExit: If text input yields no valid indicators (non-error state).
        """
        total_valid = sum(len(v) for v in self.valid_indicators_by_type.values())
        has_invalid = bool(self.unsupported_items or self.unknown_items)

        text_provided = bool(self.raw_arguments.get("text", "").strip())
        list_key_present = "indicator_list" in self.raw_arguments
        list_has_values = bool(argToList(self.raw_arguments.get("indicator_list")))

        if not text_provided and not list_has_values and not list_key_present:
            raise ValidationError(MSG_NO_VALID_INDICATORS)

        if list_key_present and (not list_has_values or total_valid == 0):
            raise ValidationError(MSG_NO_VALID_INDICATORS)

        if text_provided and not list_key_present and total_valid == 0 and not has_invalid:
            raise GracefulExit(MSG_NO_VALID_INDICATORS)

        ignore_limit = argToBoolean(self.raw_arguments.get("ignore_indicator_limit", False))
        if not ignore_limit and total_valid > MAX_INDICATORS:
            raise ValidationError(ERR_LIMIT_TEMPLATE.format(found=total_valid, limit=MAX_INDICATORS))

    def _filter_passthrough_arguments(self) -> dict[str, Any]:
        """
        Filters the raw arguments to include only those intended for child commands.

        Returns:
            A dictionary of arguments to pass to the enrichment commands.
        """
        keys = ["external_enrichment", "brands", "additional_fields"]
        return {k: self.raw_arguments.get(k) for k in keys if self.raw_arguments.get(k) is not None}


class EnrichmentService:
    """
    Coordinates the execution of child enrichment commands based on the request.
    """

    def __init__(self, request: EnrichmentRequest):
        self._result = EnrichmentResult()
        self._request = request

    def execute(self) -> EnrichmentResult:
        """
        Executes the enrichment plan.

        Args:
            request: The EnrichmentRequest containing grouped indicators.

        Returns:
            An EnrichmentResult object containing aggregated outputs.
        """

        if not self._request.valid_indicators_by_type:
            return self._result

        tasks = self._create_enrichment_tasks()

        batch_output = BatchExecutor().execute_batch([t.command for t in tasks], brands_to_run=None, verbose=False)

        for task, output in zip(tasks, batch_output):
            self._parse_task_output(task, output)

        return self._result

    def _create_enrichment_tasks(self) -> list[EnrichmentTask]:
        """Creates a list of tasks mapping indicator types to configured Commands."""
        tasks = []
        for type_enum, indicators in self._request.valid_indicators_by_type.items():
            cmd = self._build_command(type_enum, indicators)
            tasks.append(EnrichmentTask(type_enum, cmd))
        return tasks

    def _build_command(self, type_enum: IndicatorType, indicators: list[str]) -> Command:
        """
        Constructs a single XSOAR Command pointing to the corresponding underlying script based on the type given..

        Args:
            type_enum: IndicatorType IP, Domain, etc..
            indicators: List of indicators to pass to this command
            extra_args: any other args to pass to this command (i.e. additional_fields)

        Returns:
            Command: A command object which can later be executed.
        """
        cmd_args = {type_enum.argument_key: indicators}
        cmd_args.update(self._request.sub_command_arguments)

        return Command(
            name=type_enum.command_name,
            args=cmd_args,
            command_type=CommandType.INTERNAL,
            ignore_using_brand=True,
            is_multi_input=True,
            is_aggregated_output=True,
        )

    def _parse_task_output(self, task: EnrichmentTask, task_output: CommandProcessResults) -> None:
        """
        Parses the output of a BatchExecutor task and updates the result object.
        Handles Error extraction, HumanReadable aggregation, and Context extraction.
        """
        human_readable_strings = []

        for entry, hr, err in task_output:
            if err:
                self._result.execution_errors.append(err)
                self._result.markdown_sections.append(f"### Error from {task.command.name}\n{err}")
                continue

            valid_hr = hr
            if not valid_hr and isinstance(entry, dict):
                valid_hr = entry.get("HumanReadable", "")

            if valid_hr:
                human_readable_strings.append(valid_hr)

            entry_context = entry.get("EntryContext") or entry.get("Contents") or {}
            if isinstance(entry_context, dict):
                self._result.raw_context.update(entry_context)
                self._add_new_keys_to_enrichment_data(task.indicator_type, entry_context)

        if human_readable_strings:
            combined_md = "\n".join(human_readable_strings)
            header = f"### {task.command.name}"
            self._result.markdown_sections.append(f"{header}\n{combined_md}")

    def _add_new_keys_to_enrichment_data(self, indicator_type: IndicatorType, context: ContextResult) -> None:
        """
        Extracts specific enrichment data based on the indicator type prefix and add "Type" key to it.

        Args:
            indicator_type: The configuration enum for the current indicator type.
            context: The raw EntryContext dictionary from the sub-command.
        """
        for key, value in context.items():
            if key.startswith(indicator_type.context_prefix):
                items = value if isinstance(value, list) else [value]
                for item in items:
                    if isinstance(item, dict):
                        unified_entry = item.copy()
                        unified_entry["Type"] = indicator_type.display_name
                        self._result.enriched_data.append(unified_entry)


class ResponseFormatter:
    """
    Transforms the internal EnrichmentResult into the final XSOAR CommandResults object.
    """

    def format(self, result: EnrichmentResult, request: EnrichmentRequest) -> CommandResults:
        """
        Formats the result into Markdown and Context.

        Args:
            result: The results from the service execution.
            request: The original request (used for error collection and raw_context flags).

        Returns:
            A populated CommandResults object.
        """
        # Fatal Error Check: If we tried to enrich but got 0 results and >0 errors
        if not result.enriched_data and result.execution_errors:
            combined_errors = "\n".join(result.execution_errors)
            raise AllExecutionsFailedError(MSG_FATAL_EXECUTION_ERROR.format(error=combined_errors))

        markdown = self._format_markdown(result, request)
        context = self._format_context(result, request)
        return CommandResults(readable_output=markdown, outputs=context, raw_response=flatten_context_values(context))

    def _format_markdown(self, result: EnrichmentResult, request: EnrichmentRequest) -> str:
        """Generates the final Markdown string, including notes, results, and error tables."""
        sections = []

        if request.duplicates_removed_count > 0:
            sections.append(
                f"Note: Removed {request.duplicates_removed_count} duplicate indicator occurrences before enrichment."
            )

        sections.extend(result.markdown_sections)

        error_table = self._generate_error_table(request)
        if error_table:
            sections.append(error_table)

        return "\n\n".join(sections) if sections else "No enrichment results found."

    def _remove_enrichment_keys_from_ctx(self, raw_context: ContextResult) -> ContextResult:
        filtered_context = {}
        enrichment_prefixes = tuple(indicator_type.context_prefix for indicator_type in IndicatorType)
        for k, v in raw_context.items():
            if not k.startswith(enrichment_prefixes):
                filtered_context[k] = v
        return filtered_context

    def _format_context(self, result: EnrichmentResult, request: EnrichmentRequest) -> ContextResult:
        """Generates the final Context dictionary, handling raw context filtering and unified list generation."""
        final_context = {}

        if not request.include_raw_context:
            filtered_context = self._remove_enrichment_keys_from_ctx(result.raw_context)
            final_context.update(filtered_context)
        else:
            final_context.update(result.raw_context)

        # Construct the enriched data for IndicatorEnrichment key
        final_context_output = result.enriched_data[:]
        final_context_output.extend(self._collect_error_objects(request))

        if final_context_output:
            final_context["IndicatorEnrichment"] = final_context_output

        return final_context

    def _generate_error_table(self, request: EnrichmentRequest) -> str:
        rows = self._collect_error_rows(request)
        if not rows:
            return ""
        return tableToMarkdown("Invalid or unsupported indicators", rows, headers=["Type", "Value", "Status", "Message"])

    def _collect_error_rows(self, request: EnrichmentRequest) -> list[ContextResult]:
        rows = []
        for item in request.unsupported_items:
            rows.append(
                {"Type": item.type, "Value": item.value, "Status": "Error", "Message": "No script supports this indicator type."}
            )
        for value in request.unknown_items:
            rows.append({"Type": "Unknown", "Value": value, "Status": "Error", "Message": "Not a valid indicator."})

        rows.sort(key=lambda x: (x["Type"].lower(), x["Value"].lower()))
        return rows

    def _collect_error_objects(self, request: EnrichmentRequest) -> list[ContextResult]:
        return self._collect_error_rows(request)  # type: ignore


def main():
    """
    Entry point for the indicator enrichment command.
    Orchestrates building the request, executing the service, and formatting the response.
    """
    try:
        enrichment_request_builder = EnrichmentRequestBuilder(demisto.args())
        validated_request = enrichment_request_builder.build_and_validate()

        indicator_enrichment_executor = EnrichmentService(validated_request)
        result = indicator_enrichment_executor.execute()

        command_results = ResponseFormatter().format(result, validated_request)

        return_results(command_results)

    except GracefulExit as info_message:
        return_results(CommandResults(readable_output=str(info_message), outputs={}))

    except ValidationError as error_message:
        demisto.debug(f"Validation Error: {error_message}")
        return_error(str(error_message))

    except AllExecutionsFailedError as error_message:
        demisto.debug(f"Failed to execute !indicator-enrichment. Error: {error_message}\n{str(traceback.format_exc())}")
        return_error(str(error_message))

    except Exception as system_error:
        demisto.error(f"Failed to execute !indicator-enrichment. Error: {str(traceback.format_exc())}")
        return_error(f"Failed to execute !indicator-enrichment. Error: {str(system_error)}")


if __name__ in ("__main__", "__builtin__", "builtins"):
    main()