from collections import defaultdict
from dataclasses import dataclass

from typing import NamedTuple, Tuple

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 = tuple()
    missing_techniques: Tuple[Technique] | Tuple = 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 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'),
                                         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: {str(ex)}')


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