# Import from standard library
import logging
import datetime
import json
import math

# Import from third-party
from flask import Flask, render_template, json, request, jsonify, session
import pytz

# Import local files
import DAO.db_connect as db_connect
from config.config import ElasticConfig
from modules.dashboard.query_body import DashboardQueryBody


def targets_overview():
    try:
        es = db_connect.connect_elasticsearch()
        try:
            # Bar chart
            raw_top_target = es.search(
                index=ElasticConfig.VULNERABILITY_INDEX, body=DashboardQueryBody.TARGET_VUL_COUNT)
            top_target = json.loads(process_bar_chart(raw_top_target))

            # Process sort and body from arguments
            search_query = request.args.get('q') or ''
            _search_query = ' '.join(
                '*'+_+'*' for _ in search_query.split(' '))
            paging_from = int(request.args.get('page') or 1)
            paging_size = ElasticConfig.DEFAULT_PAGE_SIZE
            if 'size' in request.args:
                paging_size = int(request.args.get('size'))
                session['paging_size'] = paging_size
            elif 'paging_size' in session:
                paging_size = session['paging_size']
            sort = request.args.get('sort')
            index = request.args.get('index')
            if index is None:
                index = "scanstat.startTime"
            if sort is None:
                sort = "desc"
            sort_query = generateSortQuery(sort, index)
            query_body = {
                "_source": ["target", "hostname",  "status", "root_scan_id", "scan_id", "scan_name", "scanstat.startTime", "id"],
                "from": (paging_from - 1) * paging_size,
                "size": paging_size,
                "query": {
                    "bool": {
                        "should": [
                            {
                                "query_string": {
                                    "fields": ["target", "hostname", "status", "scan_id", "scan_name"],
                                    "query": _search_query
                                }
                            }
                        ]
                    }
                }
            }
            # Process data
            rawData = es.search(index=ElasticConfig.NMAP_INDEX,
                                body=query_body, sort=sort_query)
            processed_data = processScan(rawData)

            targetIds = ''
            for data in processed_data['records']:
                targetIds += data["id"] + ' '

            summary_result = getSummaryResult(targetIds)

            for i in range(len(processed_data['records'])):
                data = processed_data['records'][i]
                processed_data['records'][i]['summary_result'] = summary_result[processed_data['records'][i]['id']]
            paging_total = processed_data['total']
            paging = {
                "paging_from": paging_from,
                "paging_size": paging_size,
                "paging_total": paging_total,
                "left_page": max(paging_from - 2, 1),
                "right_page": min(max(paging_from - 2, 1) + 4, math.ceil(paging_total / paging_size))
            }
            if sort == "desc":
                sort = "asc"
            else:
                sort = "desc"
        except:
            logging.exception(msg="Fail")
    except:
        logging.exception(msg="Fail When Connect To ElasticSearch Database")
    session['module'] = "target"
    session['overview'] = True
    return render_template('targets.html', response=processed_data, search_query=search_query, index=index, sorting=sort, paging=paging, top_target=top_target)


def generateSortQuery(sort, index):
    sort_query = ""
    if index == '1':
        sort_query = "target.keyword" + ":" + sort
    elif index == '2':
        sort_query = "hostname.keyword" + ":" + sort
    elif index == '3':
        sort_query = "scanstat.startTime" + ":" + sort
    elif index == '4':
        sort_query = "status.keyword" + ":" + sort
    elif index == '5':
        sort_query = "scan_id.keyword" + ":" + sort
    elif index == '6':
        sort_query = "scan_name.keyword" + ":" + sort

    return sort_query


def processScan(rawData):
    try:
        UTC7 = pytz.timezone('Asia/Ho_Chi_Minh')
        processing_data = rawData['hits']
        result = {}
        # Backward version of elasticSearch compatitive
        if type(processing_data['total']) == dict:
            result['total'] = processing_data['total']['value']
        else:
            result['total'] = processing_data['total']
        result['records'] = []
        for hit in processing_data['hits']:
            _source = hit['_source']
            _source["id"] = hit.get('_id')
            _source['startTime'] = str(datetime.datetime.fromtimestamp(
                int(_source.get('scanstat').get("startTime")), tz=UTC7))
            del _source['scanstat']
            result['records'].append(_source)

        return result

    except:
        logging.exception(msg="No validated jsonified data to be processed!")


def getSummaryResult(targetIds):
    result = dict()
    list_target_id = targetIds.split(' ')
    try:
        resultAcunetix = getDataFromElastic(targetIds)
        for target_id in list_target_id:
            result[target_id] = {
                "low": 0,
                "informational": 0,
                "high": 0,
                "medium": 0
            }

            result[target_id]['informational'] += resultAcunetix.get(
                target_id, dict()).get('informational', 0)
            result[target_id]['low'] += resultAcunetix.get(
                target_id, dict()).get('low', 0)
            result[target_id]['medium'] += resultAcunetix.get(
                target_id, dict()).get('medium', 0)
            result[target_id]['high'] += resultAcunetix.get(
                target_id, dict()).get('high', 0)
    except:
        logging.exception(msg="Something went wrong!")
    return result


def getDataFromElastic(targetIds):

    result = dict()
    try:
        es = db_connect.connect_elasticsearch()
        query_body = {
            "query": {
                "bool": {
                    "should": [
                        {
                            "query_string": {
                                "fields": ["id"],
                                "query": targetIds
                            }
                        }
                    ]
                }
            },
            "size": 0,
            "aggs": {
                "group_by_target_id": {
                    "terms": {
                        "field": "id.keyword",
                        "size": ElasticConfig.MAX_SIZE_AGGS
                    },
                    "aggs": {
                        "informational": {
                            "sum": {
                                "field": "vuln_stats.informational"
                            }
                        },
                        "low": {
                            "sum": {
                                "field": "vuln_stats.low"
                            }
                        },
                        "medium": {
                            "sum": {
                                "field": "vuln_stats.medium"
                            }
                        },
                        "high": {
                            "sum": {
                                "field": "vuln_stats.high"
                            }
                        }
                    }
                }
            }
        }
        rawData = es.search(index=[ElasticConfig.ACUNETIX_SUMMARY_INDEX, ElasticConfig.CVESEARCH_INDEX,
                                   ElasticConfig.NESSUS_INDEX, ElasticConfig.NIKTO_INDEX], body=query_body)
        processing_data = rawData['aggregations']["group_by_target_id"]["buckets"]
        for data in processing_data:
            res = {
                "low": int(data['low']['value']),
                "informational": int(data['informational']['value']),
                "high": int(data['high']['value']),
                "medium": int(data['medium']['value'])
            }
            result[data.get('key')] = res

    except:
        logging.exception(msg="Something went wrong!")
    return result

# Bar chart


def process_bar_chart(raw_top_target):
    result = {}
    targets = []
    counts = []
    try:
        processing_data = raw_top_target.get(
            'aggregations').get('targets').get('buckets')
        total_targets = len(processing_data)
        total_vuls = 0
        for record in processing_data:
            total_vuls += record.get('doc_count')

            targets.append(record.get('key'))
            counts.append(str(record.get('doc_count')))

        top_5_target_string = str("-".join(targets[:5]))
        top_5_vul_count_string = str("-".join(counts[:5]))

        result['targets'] = top_5_target_string
        result['count'] = top_5_vul_count_string
        result['average'] = int(total_vuls/total_targets)
        # print(result)
    except:
        logging.exception(msg="No validated jsonified data to be processed!")
    return json.dumps(result)
