from abc import ABC
from enum import Enum

from CommonServerPython import *

SUPPORTED_HASHES = 'SHA256', 'SHA512', 'SHA1', 'MD5'
CONTEXT_FINDING_PREFIX = 'InvestigationSummary'


def get(value: dict, key: str) -> Optional[set]:
    """
    Fetches a value, entering lists when necessary.
    NOTE: None values are removed!

    >>> get(value={'foo': {'inner':'value'}}, key='foo.bar') # returns None
    >>> get(value={'foo': {'inner':'value'}}, key='foo.inner')
    {'value'}
    >>> sorted(list(get(value={'foo': [{'inner': 'value2'}, {'inner': 'value1'}, {'bar':'baz'}]}, key='foo.inner')))
    ['value1', 'value2']
    >>> sorted(list(get(value={'foo': [{'inner': ['value1', 'value2']}, {'bar':'baz'}]}, key='foo.inner')))
    ['value1', 'value2']
    """
    parts = key.split(".")
    for i, part in enumerate(parts):
        if isinstance(value, dict):
            if part in value:
                value = value[part]
                if i == len(parts) - 1:  # got to the last key part
                    if isinstance(value, list):
                        return set((sub_value for sub_value in value if sub_value is not None))
                    else:
                        return {value}
            else:
                break
        elif isinstance(value, list) and isinstance(value[0], dict) and i != len(parts):
            # NOTE: None values are removed!
            next_values = set()
            for sub_dict in value:
                next_values.update(get(sub_dict, ".".join(parts[i:])) or {})
            return set((value for value in next_values if value is not None))
    return None


class Result(Enum):
    SUSPICIOUS = 1
    NOT_DETECTED = 0


class Source(Enum):
    SANDBOX = 'Sandbox'
    EDR = 'EDR'


class MalwareFinding(ABC):
    name: str
    context: dict
    tactic: Optional[str]
    result: Optional[Result]
    sources: Set[Source]
    additional_attributes: Optional[Dict[str, str]] = {}

    def __init__(self, name: str, context: dict, tactic: Optional[str] = None):
        """
        Represents a query to the context.
        :param name: Name of the query
        :param context: Context to search in.
        :param tactic: The tactic matching this query, to be shown to be user.
        """
        self.name = name
        self.context = context
        self.tactic = tactic
        self.result = Result.NOT_DETECTED  # default
        self.sources = set()
        # make sure you self.parse_context in every inheriting class' __init__!

    def to_context(self) -> dict:
        """
        Creates a context section from the finding.
        Use `False` when creating context for aggregated results (max of findings)
        :return: a dictionary to be used as CommandResult.output.
        """
        context: Dict[str, Any] = {}

        if self.sources:
            # sources are not always defined, e.g. in Verdict.
            context['Sources'] = sorted([source.value for source in self.sources])  # sorting for deterministic results

        if self.tactic:
            context['Tactic'] = self.tactic

        if self.result:
            context['Result'] = self.result.name.title().replace('_', ' ')

        if self.additional_attributes:
            context.update(self.additional_attributes)

        return context

    def _process_value(self, value: str, expected_value: Union[str, bool], result: Result, source: Source):
        """
        Compares the value to an expected value.
        If the two are equal, the result and source are set accordingly.
        """
        if isinstance(expected_value, bool):
            try:
                value = argToBoolean(value)
            except ValueError:
                return  # if a value can't be converted to bool, it is definitely not the expected boolean.

        if value == expected_value:
            self.result = result
            self.sources.add(source)  # not using break here, in case there are multiple sources

    def _parse_key(self, key: str, expected_value: Union[str, bool], result: Result, source: Source,
                   context: dict = None):
        if context is None:
            context = self.context

        for value in get(context, key) or {}:
            self._process_value(value, expected_value, result, source)

    @abstractmethod
    def _parse_context(self) -> None:
        """
        Determines all class attribute values (result, sources, etc.) from the context.
        Naturally, this is customized in each inheriting class.
        """
        pass

    @abstractmethod
    def to_command_results(self) -> CommandResults:
        pass


class KillChain(MalwareFinding, ABC):
    def __init__(self, name: str, tactic: str, context: dict, search_value: str):
        self.search_value = search_value
        super().__init__(name=name, tactic=tactic, context=context)  # calls self._parse_context
        self._parse_context()

    def _parse_context(self) -> None:
        self._parse_key('csfalconx.resource.sandbox.mitre_attacks.tactic', self.search_value,
                        Result.SUSPICIOUS, Source.SANDBOX)
        self._parse_key('incident.mitretacticname', self.search_value, Result.SUSPICIOUS, Source.EDR)
        self._parse_key('MITREATTACK.value', self.search_value, Result.SUSPICIOUS, Source.EDR)
        self._parse_key('AttackPattern.KillChainPhases', self.search_value, Result.SUSPICIOUS, Source.EDR)

    def to_command_results(self) -> CommandResults:
        return CommandResults(outputs=self.to_context(),
                              outputs_prefix=f'{CONTEXT_FINDING_PREFIX}.{self.name}')


class Persistence(KillChain):
    def __init__(self, context: dict):
        super().__init__(
            name='EvidenceOfPersistence',
            tactic='Persistence',
            context=context,
            search_value=ThreatIntel.KillChainPhases.PERSISTENCE
        )


class DefenseEvasion(KillChain):
    def __init__(self, context: dict):
        super().__init__(
            name='EvidenceOfDefenseEvasion',
            tactic='Defense Evasion',
            context=context,
            search_value=ThreatIntel.KillChainPhases.DEFENSE_EVASION
        )


class Execution(KillChain):
    def __init__(self, context: dict):
        super().__init__(
            name='EvidenceOfExecution',
            tactic='Execution',
            context=context,
            search_value=ThreatIntel.KillChainPhases.EXECUTION
        )


class LateralMovement(KillChain):
    def __init__(self, context: dict):
        super().__init__(
            name='EvidenceOfLateralMovement',
            tactic='Lateral Movement',
            context=context,
            search_value=ThreatIntel.KillChainPhases.LATERAL_MOVEMENT
        )


class PrivilegeEscalation(KillChain):
    def __init__(self, context: dict):
        super().__init__(
            name='EvidenceOfPrivilegeEscalation',
            tactic='Privilege Escalation',
            context=context,
            search_value=ThreatIntel.KillChainPhases.PRIVILEGE_ESCALATION
        )


class CommandAndControl(KillChain):
    def __init__(self, context: dict):
        super().__init__(
            name='EvidenceOfCommandAndControl',
            tactic='Command and Control',
            context=context,
            search_value=ThreatIntel.KillChainPhases.COMMAND_AND_CONTROL
        )


def parse_command(context: dict) -> List[CommandResults]:
    return [
        finding.to_command_results() for finding in (
            Persistence(context),
            DefenseEvasion(context),
            Execution(context),
            LateralMovement(context),
            PrivilegeEscalation(context),
            CommandAndControl(context),
        )
    ]


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

    except Exception as e:
        demisto.error(traceback.format_exc())  # print the traceback
        return_error(f'Failed to execute FindingsParse. Error: {str(e)}')


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