from __future__ import annotations

from typing import Optional

from assemblyline.common import forge
from assemblyline.common.caching import TimeExpiredCache
from assemblyline.common.dict_utils import recursive_update
from assemblyline.common.isotime import now_as_iso
from assemblyline.common.str_utils import safe_str
from assemblyline.datastore.exceptions import VersionConflictException
from assemblyline.datastore.helper import AssemblylineDatastore
from assemblyline.odm import IP, URI, Domain
from assemblyline.odm.messages.alert import AlertMessage
from assemblyline.odm.models.alert import Alert
from assemblyline.remote.datatypes.queues.comms import CommsQueue

CACHE_LEN = 60 * 60 * 24
CACHE_EXPIRY_RATE = 60

action_queue_map = None
cache: TimeExpiredCache[tuple[Optional[dict], Optional[dict]]] = TimeExpiredCache(CACHE_LEN, CACHE_EXPIRY_RATE)
Classification = forge.get_classification()
config = forge.get_config()

URI_ODM = URI()
DOMAIN_ODM = Domain()
IP_ODM = IP()


SUMMARY_TYPE_MAP = {
    'av.virus_name': 'av',
    'attribution.exploit': 'attrib',
    'attribution': 'attrib',
    'file.config': 'attrib',
    'technique.obfuscation': 'attrib',
    'file.behavior': 'behavior',
    'file.rule.yara': 'yara',
    'attribution.family': 'attrib',
    'attribution.implant': 'attrib',
    'network.static.domain': 'domain_static',
    'network.dynamic.domain': 'domain_dynamic',
    'network.static.ip': 'ip_static',
    'network.dynamic.ip': 'ip_dynamic',
    'network.static.uri': 'uri_static',
    'network.dynamic.uri': 'uri_dynamic',
    'technique.config': 'attrib',
    'attribution.actor': 'attrib',
    'heuristic.name': 'heuristic',
    'attack.pattern': 'attack_pattern',
    'attack.category': 'attack_category'
}

AL_RESULT_KEYS = [
    'attrib',
    'av',
    'behavior',
    'domain',
    'domain_dynamic',
    'domain_static',
    'ip',
    'ip_dynamic',
    'ip_static',
    'uri',
    'uri_dynamic',
    'uri_static',
    'yara']

VERDICT_RANK_MAP = {
    "safe": 0,
    "info": 1,
    "suspicious": 2,
    "malicious": 3
}

TAG_MAP = {
    'attribution.exploit': 'EXP',
    'file.config': 'CFG',
    'technique.obfuscation': 'OB',
    'attribution.family': 'IMP',
    'attribution.implant': 'IMP',
    'technique.config': 'CFG',
    'attribution.actor': 'TA',
}

AV_TO_BEAVIOR = (
    'Corrupted executable file',
    'Encrypted container deleted',
    'Encrypted container deleted;',
    'Password-protected',
    'Malformed container violation',
)


class AlertMissingError(Exception):
    pass


class SubmissionNotFinalized(Exception):
    pass


class SubmissionNotFound(Exception):
    pass


def service_name_from_key(key):
    # noinspection PyBroadException
    try:
        return key.split('.')[1]
    except Exception:
        return ""


def get_submission_record(counter, datastore, sid):
    srecord = datastore.submission.get(sid, as_obj=False)

    if not srecord:
        counter.increment('error')
        raise SubmissionNotFound(f"Couldn't find submission: {sid}")

    submission_state = srecord.get('state', 'unknown')
    if submission_state != 'completed':
        raise SubmissionNotFinalized(f"Submission not finalized: {sid} ({submission_state})")

    return srecord


def get_summary(datastore, srecord, user_classification, logger):
    max_classification = srecord['classification']

    detailed = {v: {} for v in SUMMARY_TYPE_MAP.values()}

    submission_summary = datastore.get_summary_from_keys(
        srecord.get('results', []),
        cl_engine=Classification, user_classification=user_classification,
        screenshot_sha256=srecord['files'][0]['sha256'])
    max_classification = Classification.max_classification(max_classification, submission_summary['classification'])

    # Process Att&cks
    for attack in submission_summary['attack_matrix']:
        tag_type = 'attack_pattern'
        item_key = attack['name']
        if item_key not in detailed[tag_type]:
            detailed[tag_type][item_key] = {
                "type": tag_type,
                "subtype": None,
                "verdict": attack['h_type'],
                "value": item_key
            }
        else:
            if VERDICT_RANK_MAP[detailed[tag_type][item_key]['verdict']] < VERDICT_RANK_MAP[attack['h_type']]:
                detailed[tag_type][item_key]['verdict'] = attack['h_type']

        # Process Att&cks categories
        for cat in attack['categories']:
            tag_type = 'attack_category'
            if cat not in detailed[tag_type]:
                detailed[tag_type][cat] = {
                    "type": tag_type,
                    "subtype": None,
                    "verdict": attack['h_type'],
                    "value": cat
                }
            else:
                if VERDICT_RANK_MAP[detailed[tag_type][cat]['verdict']] < VERDICT_RANK_MAP[attack['h_type']]:
                    detailed[tag_type][cat]['verdict'] = attack['h_type']

    # Process Heuristics
    for h_type, heur_list in submission_summary['heuristics'].items():
        for heur in heur_list:
            tag_type = 'heuristic'
            item_key = heur['name']
            if item_key not in detailed[tag_type]:
                detailed[tag_type][item_key] = {
                    "type": tag_type,
                    "subtype": None,
                    "verdict": h_type,
                    "value": item_key
                }
            else:
                if VERDICT_RANK_MAP[detailed[tag_type][item_key]['verdict']] < VERDICT_RANK_MAP[h_type]:
                    detailed[tag_type][item_key]['verdict'] = h_type

    # Process Tags
    for t in submission_summary['tags']:
        if t.get('safelisted', False):
            continue

        tag_value = t['value']
        tag_type = SUMMARY_TYPE_MAP.get(t['type'], None)

        if tag_value == '' or tag_type is None:
            continue

        sub_type = TAG_MAP.get(t['type'], None)

        if tag_type == 'av':
            if tag_value in AV_TO_BEAVIOR:
                tag_type = 'behavior'

        item_key = f"{t['type']}__{t['value']}"
        if item_key not in detailed[tag_type]:
            detailed[tag_type][item_key] = {
                "type": t['type'],
                "subtype": sub_type,
                "verdict": t['h_type'],
                "value": t['value']
            }
        else:
            if VERDICT_RANK_MAP[detailed[tag_type][item_key]['verdict']] < VERDICT_RANK_MAP[t['h_type']]:
                detailed[tag_type][item_key]['verdict'] = t['h_type']

    # Simplify the detailed items
    detailed = {k: list(v.values()) for k, v in detailed.items()}

    # Validate the network IOCs raised to ensure they're valid when creating/updating alerts
    for net_type in ['ip_dynamic', 'ip_static', 'domain_dynamic', 'domain_static', 'uri_dynamic', 'uri_static']:
        valid_items = []
        if net_type in detailed:
            # If the key exists, validate each item in the collection
            for item in detailed[net_type]:
                try:
                    if net_type.startswith('ip'):
                        IP_ODM.check(item['value'])
                    elif net_type.startswith('domain'):
                        DOMAIN_ODM.check(item['value'])
                    elif net_type.startswith('uri'):
                        URI_ODM.check(item['value'])
                except Exception:
                    logger.warning(f"Invalid {net_type} IOC found in submission {srecord['sid']}: {item}")
                    continue
                valid_items.append(item)

            # Update with only valid items
            detailed[net_type] = valid_items

    # Generate the Summary
    summary = {k: set([f"{item['value']} [{item['subtype']}]" if item.get('subtype', None) else item['value']
                       for item in v]) for k, v in detailed.items()}
    summary['domain'] = summary['domain_dynamic'].union(summary['domain_static'])
    summary['ip'] = summary['ip_dynamic'].union(summary['ip_static'])
    summary['uri'] = summary['uri_dynamic'].union(summary['uri_static'])

    # Merge IPs and domains in details
    domains = list(detailed.pop('domain_dynamic', {}))
    domains.extend(detailed.pop('domain_static', {}))
    ips = list(detailed.pop('ip_dynamic', {}))
    ips.extend(detailed.pop('ip_static', {}))
    uris = list(detailed.pop('uri_dynamic', {}))
    uris.extend(detailed.pop('uri_static', {}))
    detailed['ip'] = ips
    detailed['domain'] = domains
    detailed['uri'] = uris

    try:
        screenshots = [{'name': x['img']['name'], 'description': x['img']['description'], 'img': x['img']
                        ['sha256'], 'thumb': x['thumb']['sha256'], } for x in submission_summary['screenshots']]
    except Exception:
        logger.warning(f"Cannot parse the submission summary screenshot section. ({submission_summary['screenshots']})")
        screenshots = []

    return max_classification, summary, submission_summary['filtered'], detailed, screenshots


def generate_alert_id(logger, alert_data):
    if alert_data['ingest_id']:
        return alert_data['ingest_id']
    sid = alert_data['submission']['params']['psid'] or alert_data['submission']['sid']
    logger.info(f"ingest_id not found for sid={sid}")
    return sid


# noinspection PyBroadException
def parse_submission_record(counter, datastore, alert_data, logger, user_classification):
    sid = alert_data['submission']['sid']
    psid = alert_data['submission']['params']['psid']
    srecord = get_submission_record(counter, datastore, sid)

    max_classification, summary, filtered, detailed, screenshots = get_summary(
        datastore, srecord, user_classification, logger)

    extended_scan = alert_data['extended_scan']
    if psid and extended_scan == 'skipped':
        try:
            # Get errors from parent submission and submission. Strip keys
            # to only sha256 and service name. If there are any keys that
            # did not exist in the parent the extended scan is 'incomplete'.
            ps = datastore.submission.get(psid, as_obj=False) or {}
            pe = set((x[:x.rfind('.')] for x in ps.get('errors', [])))
            e = set((x[:x.rfind('.')] for x in srecord.get('errors', [])))
            ne = e.difference(pe)
            extended_scan = 'incomplete' if ne else 'completed'
        except Exception:  # pylint: disable=W0702
            logger.exception('Problem determining extended scan state:')

    return {
        'extended_scan': extended_scan,
        'summary': summary,
        'srecord': srecord,
        'max_classification': max_classification,
        'filtered': filtered,
        'detailed': detailed,
        'screenshots': screenshots
    }


AL_FIELDS = [
    'attrib',
    'av',
    'behavior',
    'domain',
    'domain_dynamic',
    'domain_static',
    'ip',
    'ip_dynamic',
    'ip_static',
    'yara'
]


def create_or_update_alert(datastore: AssemblylineDatastore, logger, alert, counter):
    alert = Alert(alert)
    alert_id = alert.alert_id

    while True:
        old_alert: Optional[Alert]
        old_alert, version = datastore.alert.get_if_exists(alert_id, version=True)
        if old_alert is None:
            try:
                datastore.alert.save(alert_id, alert, version=version)
                logger.info(f"Alert {alert_id} has been created.")
                counter.increment('created')
                return "AlertCreated", "create"
            except VersionConflictException as vce:
                logger.info(f"Retrying update alert due to version conflict: {str(vce)}")
                continue

        # merge both alerts together so it doesn't matter if messages are out of order
        old_alert.update(alert)

        # Make sure workflows are re-triggered
        old_alert.workflows_completed = False

        try:
            datastore.alert.save(alert_id, old_alert, version=version)
            logger.info(f"Alert {alert_id} has been updated.")
            counter.increment('updated')
            return "AlertUpdated", 'update'
        except VersionConflictException as vce:
            logger.info(f"Retrying update alert due to version conflict: {str(vce)}")


def save_alert(datastore, counter, logger, alert):
    msg_type, ret_val = create_or_update_alert(datastore, logger, alert, counter)

    msg = AlertMessage({
        "msg": alert,
        "msg_type": msg_type,
        "sender": "alerter"
    })
    CommsQueue('alerts').publish(msg.as_primitives())
    return ret_val


# noinspection PyTypeChecker
def get_alert_update_parts(counter, datastore, alert_data, logger, user_classification):
    global cache
    # Check cache
    alert_file = alert_data['submission']['files'][0]
    alert_update_p1, alert_update_p2 = cache.get(alert_data['submission']['sid'], (None, None))
    if alert_update_p1 is None or alert_update_p2 is None:
        parsed_record = parse_submission_record(counter, datastore, alert_data, logger, user_classification)
        file_record = datastore.file.get(alert_file['sha256'], as_obj=False)
        al = {k: list(v) for k, v in parsed_record['summary'].items() if k in AL_RESULT_KEYS}
        al['detailed'] = parsed_record['detailed']
        al['request_end_time'] = parsed_record['srecord']['times']['completed']
        alert_update_p1 = {
            'extended_scan': parsed_record['extended_scan'],
            'filtered': parsed_record['filtered'],
            'al': al,
            'attack': {
                'pattern': list(parsed_record['summary']['attack_pattern']),
                'category': list(parsed_record['summary']['attack_category'])

            },
            'heuristic': {
                'name': list(parsed_record['summary']['heuristic'])
            }
        }
        alert_update_p2 = {
            'classification': parsed_record['max_classification'],
            'file': {
                'md5': file_record['md5'],
                'sha1': file_record['sha1'],
                'sha256': file_record['sha256'],
                'size': file_record['size'],
                'type': file_record['type'],
                'screenshots': parsed_record['screenshots']
            },
            'verdict': {
                "malicious": parsed_record['srecord']['verdict']['malicious'],
                "non_malicious": parsed_record['srecord']['verdict']['non_malicious']
            }
        }
        cache.add(alert_data['submission']['sid'], (alert_update_p1, alert_update_p2))

    alert_update_p1['reporting_ts'] = now_as_iso()
    alert_update_p1['file'] = {'name': alert_file['name']}

    return alert_update_p1, alert_update_p2


def process_alert_message(counter, datastore, logger, alert_data):
    """
    This is the default process_alert_message function. If the generic alerts are not sufficient
    in your deployment, you can create another method like this one that would follow
    the same structure but with added parts where the comment blocks are located.
    """

    ###############################
    # Additional init goes here
    ###############################

    user = datastore.user.get(alert_data.get('submission', {}).get('params', {}).get('submitter', None), as_obj=False)
    if user:
        user_classification = user['classification']
    else:
        user_classification = None
    a_type = alert_data.get('submission', {}).get('metadata', {}).pop('type', None)
    a_ts = alert_data.get('submission', {}).get('metadata', {}).pop('ts', None)

    alert = {
        'al': {
            'score': alert_data['score']
        },
        'alert_id': generate_alert_id(logger, alert_data),
        'archive_ts': None,
        'metadata': {safe_str(key): value for key, value in alert_data['submission']['metadata'].items()},
        'submission_relations': [{
            'child': alert_data['submission']['sid'],
            'parent': alert_data['submission']['params']['psid']
        }],
        'sid': alert_data['submission']['sid'],
        'ts': a_ts or alert_data['submission']['time'],
        'type': a_type or alert_data['submission']['params']['type']
    }

    if config.core.alerter.alert_ttl:
        alert['expiry_ts'] = now_as_iso(config.core.alerter.alert_ttl * 24 * 60 * 60)

    ###############################
    # Additional alert_data parsing
    # and alert updating goes here
    ###############################

    # Get update parts
    alert_update_p1, alert_update_p2 = get_alert_update_parts(counter, datastore,
                                                              alert_data, logger,
                                                              user_classification)

    # Update alert with default values
    alert = recursive_update(alert, alert_update_p1)

    # Update alert with computed values
    alert = recursive_update(alert, alert_update_p2)

    return save_alert(datastore, counter, logger, alert)
