import asyncio
import logging
import uuid
from base64 import b64decode
from datetime import datetime, timezone
from importlib import import_module

import marshmallow as ma

from app.objects.c_ability import Ability, AbilitySchema
from app.objects.secondclass.c_executor import Executor, ExecutorSchema
from app.objects.secondclass.c_fact import Fact, FactSchema, OriginType
from app.objects.secondclass.c_relationship import RelationshipSchema
from app.objects.secondclass.c_visibility import Visibility, VisibilitySchema
from app.utility.base_object import BaseObject
from app.utility.base_parser import PARSER_SIGNALS_FAILURE
from app.utility.base_service import BaseService


NO_STATUS_SET = object()


class LinkSchema(ma.Schema):

    class Meta:
        unknown = ma.EXCLUDE

    id = ma.fields.String(load_default='')
    paw = ma.fields.String()
    command = ma.fields.String()
    plaintext_command = ma.fields.String()
    status = ma.fields.Integer(load_default=-3)
    score = ma.fields.Integer(load_default=0)
    jitter = ma.fields.Integer(load_default=0)
    decide = ma.fields.DateTime(format=BaseObject.TIME_FORMAT)
    pin = ma.fields.Integer(load_default=0)
    pid = ma.fields.String()
    facts = ma.fields.List(ma.fields.Nested(FactSchema()))
    relationships = ma.fields.List(ma.fields.Nested(RelationshipSchema()))
    used = ma.fields.List(ma.fields.Nested(FactSchema()))
    unique = ma.fields.String()
    collect = ma.fields.DateTime(format=BaseObject.TIME_FORMAT, dump_default='')
    finish = ma.fields.String()
    ability = ma.fields.Nested(AbilitySchema())
    executor = ma.fields.Nested(ExecutorSchema())
    cleanup = ma.fields.Integer(load_default=0)
    visibility = ma.fields.Nested(VisibilitySchema())
    host = ma.fields.String(load_default=None)
    output = ma.fields.String()
    deadman = ma.fields.Boolean()
    agent_reported_time = ma.fields.DateTime(format=BaseObject.TIME_FORMAT, load_default=None)

    @ma.pre_load()
    def fix_ability(self, link, **_):
        if 'ability' in link and isinstance(link['ability'], Ability):
            ability = link.pop('ability')
            link['ability'] = ability.schema.dump(ability)
        return link

    @ma.pre_load()
    def fix_executor(self, link, **_):
        if 'executor' in link and isinstance(link['executor'], Executor):
            executor = link.pop('executor')
            link['executor'] = executor.schema.dump(executor)
        return link

    @ma.pre_load()
    def remove_properties(self, data, **_):
        data.pop('unique', None)
        data.pop('decide', None)
        data.pop('pid', None)
        data.pop('facts', None)
        data.pop('collect', None)
        data.pop('finish', None)
        data.pop('visibility', None)
        data.pop('output', None)
        data.pop('used.unique', None)
        return data

    @ma.post_load()
    def build_link(self, data, **kwargs):
        return None if kwargs.get('partial') is True else Link(**data)

    @ma.post_dump()
    def prepare_dump(self, data, **_):
        if data.get('agent_reported_time', None) is None:
            data.pop('agent_reported_time', None)
        return data


class Link(BaseObject):

    schema = LinkSchema()
    display_schema = LinkSchema(exclude=['jitter'])
    load_schema = LinkSchema(exclude=['decide', 'pid', 'facts', 'unique', 'collect', 'finish', 'visibility',
                                      'output', 'used.unique'])

    RESERVED = dict(origin_link_id='#{origin_link_id}')

    EVENT_EXCHANGE = 'link'
    EVENT_QUEUE_STATUS_CHANGED = 'status_changed'

    @property
    def raw_command(self):
        return self.decode_bytes(self.command) if self.command else ''

    @property
    def unique(self):
        return self.hash('%s' % self.id)

    @property
    def pin(self):
        return self._pin

    @pin.setter
    def pin(self, p):
        self._pin = p

    @property
    def states(self):
        return dict(HIGH_VIZ=-5,
                    UNTRUSTED=-4,
                    EXECUTE=-3,
                    DISCARD=-2,
                    PAUSE=-1,
                    SUCCESS=0,
                    ERROR=1,
                    TIMEOUT=124)

    @property
    def status(self):
        return self._status

    @property
    def display(self):
        dump = LinkSchema(exclude=['jitter']).dump(self)
        dump['command'] = self.decode_bytes(dump['command'])
        dump['plaintext_command'] = self.decode_bytes(dump['plaintext_command'])
        return dump

    @status.setter
    def status(self, value):
        previous_status = getattr(self, '_status', NO_STATUS_SET)

        self._status = value

        if previous_status is NO_STATUS_SET:
            return

        if previous_status == value:
            return

        self._emit_status_change_event(
            from_status=previous_status,
            to_status=value
        )

    @classmethod
    def is_global_variable(cls, variable):
        return variable in cls.RESERVED

    def __init__(self, command='', plaintext_command='', paw='', ability=None, executor=None, status=-3, score=0, jitter=0, cleanup=0, id='',
                 pin=0, host=None, deadman=False, used=None, relationships=None, agent_reported_time=None):
        super().__init__()
        self.id = str(id)
        self.command = command
        self.plaintext_command = plaintext_command
        self.command_hash = None
        self.paw = paw
        self.host = host
        self.cleanup = cleanup
        self.ability = ability
        self.executor = executor
        self.status = status
        self.score = score
        self.jitter = jitter
        self.decide = datetime.now(timezone.utc)
        self.pid = None
        self.collect = None
        self.finish = None
        self.facts = []
        self.relationships = relationships if relationships else []
        self.used = used if used else []
        self.visibility = Visibility()
        self._pin = pin
        self.output = False
        self.deadman = deadman
        self.agent_reported_time = agent_reported_time

    def __eq__(self, other):
        if isinstance(other, Link):
            return other.paw == self.paw and other.ability.ability_id == self.ability.ability_id \
                   and other.used == self.used
        return False

    async def parse(self, operation, result):
        if self.status != 0:
            return
        for parser in self.executor.parsers:
            source_facts = operation.source.facts if operation else []
            try:
                relationships = await self._parse_link_result(result, parser, source_facts)

                if len(relationships) > 0 and relationships[0] == PARSER_SIGNALS_FAILURE:
                    logging.getLogger('link').debug(f'link {self.id} (ability id={self.ability.ability_id}) encountered '
                                                    f'an error during execution, which was caught during parsing.')
                    self.status = self.states['ERROR']
                    relationships = []  # we didn't actually get anything out of this, so let's reset
                else:
                    await self.create_relationships(relationships, operation)
                await update_scores(operation, increment=len(relationships), used=self.used, facts=self.facts)
            except Exception as e:
                logging.getLogger('link').debug('error in %s while parsing ability %s: %s'
                                                % (parser.module, self.ability.ability_id, e))

    def apply_id(self, host):
        self.id = str(uuid.uuid4())
        self.host = host
        self.replace_origin_link_id()

    def can_ignore(self):
        return self.status in [self.states['DISCARD'], self.states['HIGH_VIZ']]

    def is_finished(self):
        return self.status in [self.states['DISCARD'], self.states['SUCCESS'],
                               self.states['ERROR'], self.states['TIMEOUT']]

    def is_valid_status(self, status):
        return status in self.states.values()

    def replace_origin_link_id(self):
        decoded_cmd = self.decode_bytes(self.command)
        self.command = self.encode_string(decoded_cmd.replace(self.RESERVED['origin_link_id'], self.id))

    def _emit_status_change_event(self, from_status, to_status):
        event_svc = BaseService.get_service('event_svc')

        task = asyncio.get_event_loop().create_task(
            event_svc.fire_event(
                exchange=Link.EVENT_EXCHANGE,
                queue=Link.EVENT_QUEUE_STATUS_CHANGED,
                link=self.id,
                from_status=from_status,
                to_status=to_status
            )
        )

        return task

    async def _parse_link_result(self, result, parser, source_facts):
        blob = b64decode(result).decode('utf-8')
        parser_info = dict(module=parser.module, used_facts=self.used, mappers=parser.parserconfigs,
                           source_facts=source_facts)
        p_inst = await self._load_module('Parser', parser_info)
        return p_inst.parse(blob=blob)

    @staticmethod
    async def _load_module(module_type, module_info):
        module = import_module(module_info['module'])
        return getattr(module, module_type)(module_info)

    async def create_relationships(self, relationships, operation):
        for relationship in relationships:
            relationship.origin = operation.id if operation else self.id
            await self.save_fact(operation, relationship.source, relationship.score, relationship.shorthand)
            await self.save_fact(operation, relationship.target, relationship.score, relationship.shorthand)
            if all((relationship.source.trait, relationship.edge)):
                knowledge_svc_handle = BaseService.get_service('knowledge_svc')
                await knowledge_svc_handle.add_relationship(relationship)
                self.relationships.append(relationship)

    async def save_fact(self, operation, fact, score, relationship):
        knowledge_svc_handle = BaseService.get_service('knowledge_svc')
        all_facts = await operation.all_facts() if operation else self.facts
        source = operation.id if operation else self.id
        rl = [relationship] if relationship else []
        if all([fact.trait, fact.value]):
            if operation and operation.source:
                if any([(fact.trait, fact.value) == (x.trait, x.value) for x in
                        await knowledge_svc_handle.get_facts(criteria=dict(source=operation.source.id))]):
                    source = operation.source.id
            fact.source = source  # Manual addition to ensure the check works correctly
            if not await knowledge_svc_handle.check_fact_exists(fact, all_facts):
                f_gen = Fact(trait=fact.trait, value=fact.value, source=source, score=score, collected_by=[self.paw],
                             technique_id=self.ability.technique_id, links=[self.id], relationships=rl,
                             origin_type=OriginType.LEARNED)
                self.facts.append(f_gen)
                await knowledge_svc_handle.add_fact(f_gen)
            else:
                existing_fact = (await knowledge_svc_handle.get_facts(criteria=dict(trait=fact.trait,
                                                                                    value=fact.value,
                                                                                    source=fact.source)))[0]
                if self.id not in existing_fact.links:
                    existing_fact.links.append(self.id)
                if relationship not in existing_fact.relationships:
                    existing_fact.relationships.append(relationship)
                if self.paw not in existing_fact.collected_by and existing_fact not in self.used:
                    existing_fact.collected_by.append(self.paw)
                await knowledge_svc_handle.update_fact(criteria=dict(trait=fact.trait, value=fact.value,
                                                                     source=fact.source),
                                                       updates=dict(links=existing_fact.links,
                                                                    relationships=existing_fact.relationships,
                                                                    collected_by=existing_fact.collected_by))
                existing_local_record = [x for x in self.facts if x.trait == fact.trait and x.value == fact.value]
                if existing_local_record:
                    existing_local_record[0].links = existing_fact.links
                else:
                    self.facts.append(existing_fact)


async def update_scores(operation, increment, used, facts):
    knowledge_svc_handle = BaseService.get_service('knowledge_svc')
    for uf in used:
        all_facts = await operation.all_facts() if operation else facts
        for found_fact in all_facts:
            if found_fact.unique == uf.unique:
                found_fact.score += increment
                await knowledge_svc_handle.update_fact(dict(trait=found_fact.trait, value=found_fact.value,
                                                            source=found_fact.source), dict(score=found_fact.score))
                break
