InvestigationDetailedSummaryParse

Source

from collections import defaultdict
from dataclasses import dataclass
from typing import NamedTuple

from CommonServerPython import *

FoundTechnique = NamedTuple("FoundTechnique", (("mitre_id", str), ("name", str), ("tactics", list[str])))
Technique = NamedTuple("Technique", (("code", str), ("name", str)))
Tactic = NamedTuple("Tactic", (("code", str), ("name", str), ("default_techniques", Optional[tuple[Technique]])))
Table = NamedTuple("Table", (("title", str), ("content", list[dict[str, str]])))

CONTEXT_PREFIX = "InvestigationDetailedSummary"


@dataclass
class TacticFinding:
    tactic: Tactic
    found_techniques: tuple[Technique] | tuple = ()
    missing_techniques: tuple[Technique] | tuple = ()

    def command_result(self):
        result = {technique.name: True for technique in self.found_techniques} | {
            technique.name: False for technique in self.missing_techniques
        }

        return CommandResults(outputs=result, outputs_prefix=f"{CONTEXT_PREFIX}.{self.tactic.name}")


TACTICS = (
    Tactic("TA0002", "Execution", default_techniques=(Technique("T1059", "Command and Scripting Interpreter"),)),
    Tactic("TA0004", "Privilege Escalation", default_techniques=(Technique("T1547", "Boot or Logon Autostart Execution"),)),
    Tactic("TA0008", "Lateral Movement", default_techniques=(Technique("T1021", "Remote Services"),)),
    Tactic("TA0005", "Defense Evasion", default_techniques=(Technique("T1070", "Indicator Removal on Host"),)),
    Tactic("TA0003", "Persistence", default_techniques=(Technique("T1547", "Boot or Logon Autostart Execution"),)),
)


class ContextParser:
    def __init__(self, context: dict):
        self.tactic_to_found_technique = defaultdict(list)
        self.technique_id_to_found_technique = defaultdict(list)

        techniques = ContextParser._parse_techniques(context)
        for technique in techniques:
            self.technique_id_to_found_technique[technique.mitre_id].append(technique)

            for tactic in technique.tactics:
                self.tactic_to_found_technique[tactic].append(technique)

    @staticmethod
    def _parse_techniques(context: dict) -> tuple[FoundTechnique]:  # type:ignore[arg-type]
        attack_pattern_context = context.get("AttackPattern", ())
        if isinstance(attack_pattern_context, dict):  # can be dict or list of dicts.
            attack_pattern_context = (attack_pattern_context,)

        result = []
        for attack_dict in attack_pattern_context:
            if "value" not in (k.lower() for k in attack_dict):  # outdated MITREv2 pack, <= v1.1.0
                raise DemistoException("please make sure the MITRE ATT&CK v2 pack is up-to-date (v1.1.1 or newer)")

            result.append(
                FoundTechnique(
                    mitre_id=attack_dict.get("MITREID", ""),
                    name=attack_dict.get("Value", "") or attack_dict.get("value", ""),
                    tactics=attack_dict.get("KillChainPhases", []),
                )
            )
        return tuple(result)  # type:ignore[return-value]

    def find(self) -> tuple[TacticFinding]:
        result = []
        for tactic in TACTICS:
            found = tuple(self.tactic_to_found_technique.get(tactic.name, ()))
            missing = tuple(filter(lambda t: t.code not in self.technique_id_to_found_technique, tactic.default_techniques or ()))
            result.append(TacticFinding(tactic, found, missing))  # type:ignore[arg-type]

        searched_tactic_names = {tactic.name for tactic in TACTICS}
        for tactic in set(self.tactic_to_found_technique.keys()).difference(searched_tactic_names):
            result.append(
                TacticFinding(
                    tactic=Tactic(name=tactic, code="", default_techniques=None),
                    found_techniques=tuple(self.tactic_to_found_technique[tactic]),
                )  # type:ignore[arg-type]
            )
        return tuple(result)  # type:ignore[return-value]


def parse_command(context: dict) -> list[CommandResults]:
    parser = ContextParser(context)
    return [finding.command_result() for finding in parser.find()]


def main():
    try:
        context = demisto.callingContext.get("context", {}).get("ExecutionContext", {})
        result = parse_command(context)
        return_results(result)

    except Exception as ex:
        demisto.error(traceback.format_exc())  # print the traceback
        return_error(f"Failed to execute InvestigationDetailedSummaryParse. Error: {ex!s}")


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

README