from flask import Blueprint, request, abort
from flask.json import jsonify
from werkzeug.utils import secure_filename
import capa.main
import capa.rules
import capa.loader
import capa.render.json
import capa.capabilities.common
from capa.features.common import OS_AUTO, FORMAT_AUTO
from os import path
from pathlib import Path
import config
import json
import os
import capa.render.utils as rutils
import capa.render.result_document as rd
import collections
from capa.render.default import find_subrule_matches
capa_bp = Blueprint('capa', __name__)

# This function is essentially a hacked up version of https://github.com/mandiant/capa/blob/master/capa/render/default.py
@capa_bp.route('/analyze', methods=['GET'])
def analyze_capa():
    file = secure_filename(request.args.get('file', ''))
    if file == '':
        abort(400)
    filepath = Path(path.join(config.Config.FILE_DIRECTORY, file))
    if not os.path.exists(filepath):
        print(f"Error: File not found at '{filepath}'")
        abort(400)

    rules = capa.rules.get_rules([Path(config.Config.RULES)])
    extractor = capa.loader.get_extractor(filepath, FORMAT_AUTO, OS_AUTO, capa.main.BACKEND_VIV, [], should_save_workspace=False, disable_progress=True)
    capabilities = capa.capabilities.common.find_capabilities(rules, extractor, disable_progress=True)

    meta = capa.loader.collect_metadata([], filepath, FORMAT_AUTO, OS_AUTO, [capa.main.get_default_root()/ "rules"], extractor, capabilities)
    doc = rd.ResultDocument.from_capa(meta, rules, capabilities.matches)
    caps = {}
    tactics = {}
    maec = {}
    objectives = {}

    subrule_matches = find_subrule_matches(doc)
    for rule in rutils.capability_rules(doc):
        if rule.meta.name in subrule_matches:
            # rules that are also matched by other rules should not get rendered by default.
            # this cuts down on the amount of output while giving approx the same detail.
            # see #224
            continue

        count = len(rule.matches)
        if count == 1:
            capability = rule.meta.name
        else:
            capability = rule.meta.name + f" ({count} matches)"
        caps[capability] = rule.meta.namespace
        for attack in rule.meta.attack:
            tactics[attack.tactic] = attack.technique + attack.subtechnique + attack.id.strip("[").strip("]")

    maec_categories = {
        "analysis_conclusion",
        "analysis_conclusion_ov",
        "malware_family",
        "malware_category",
        "malware_category_ov",
    }
    for rule in rutils.maec_rules(doc):
        for maec_category in maec_categories:
            maec_value = getattr(rule.meta.maec, maec_category, None)
            if maec_value:
                maec[maec_category] = maec_value
    for rule in rutils.capability_rules(doc):
        for mbc in rule.meta.mbc:
            objectives[mbc.objective] = mbc.behavior + mbc.method + mbc.id.strip("[").strip("]")

    return jsonify(capabilities=caps, tactics=tactics, maec=maec, objectives=objectives)


@capa_bp.route('/json', methods=['GET'])
def analyze_capa_json():
    file = secure_filename(request.args.get('file', ''))
    if file == '':
        abort(400)
    filepath = Path(path.join(config.Config.FILE_DIRECTORY, file))
    if not os.path.exists(filepath):
        print(f"Error: File not found at '{filepath}'")
        abort(400)

    rules = capa.rules.get_rules([Path(config.Config.RULES)])
    extractor = capa.loader.get_extractor(filepath, FORMAT_AUTO, OS_AUTO, capa.main.BACKEND_VIV, [], should_save_workspace=False, disable_progress=True)
    capabilities = capa.capabilities.common.find_capabilities(rules, extractor, disable_progress=True)

    meta = capa.loader.collect_metadata([], filepath, FORMAT_AUTO, OS_AUTO, [capa.main.get_default_root()/ "rules"], extractor, capabilities)

    return json.loads(capa.render.json.render(meta=meta, rules=rules, capabilities=capabilities.matches))



@capa_bp.route('/default', methods=['GET'])
def analyze_capa_default():
    file = secure_filename(request.args.get('file', ''))
    if file == '':
        abort(400)
    filepath = Path(path.join(config.Config.FILE_DIRECTORY, file))
    if not os.path.exists(filepath):
        print(f"Error: File not found at '{filepath}'")
        abort(400)

    rules = capa.rules.get_rules([Path(config.Config.RULES)])
    extractor = capa.loader.get_extractor(filepath, FORMAT_AUTO, OS_AUTO, capa.main.BACKEND_VIV, [], should_save_workspace=False, disable_progress=True)
    capabilities = capa.capabilities.common.find_capabilities(rules, extractor, disable_progress=True)

    meta = capa.loader.collect_metadata([], filepath, FORMAT_AUTO, OS_AUTO, [capa.main.get_default_root()/ "rules"], extractor, capabilities)

    return capa.render.default.render(meta=meta, rules=rules, capabilities=capabilities.matches)