#!/usr/bin/env python3
"""
vulnerability-engine evaluator
"""

import os
import json
from prometheus_client import Counter, Histogram, start_http_server
from psycopg2 import DatabaseError
from psycopg2.extras import execute_values
import requests

from common import mqueue
from common.database_handler import DatabaseHandler, init_db
from common.logging import init_logging, get_logger
from common.utils import vmaas_post_request

LOGGER = get_logger(__name__)

VMAAS_HOST = os.getenv('VMAAS_HOST', 'http://vmaas-webapp-1.vmaas-ci.svc:8080')
VMAAS_VULNERABILITIES_API = os.getenv("VMAAS_VULNERABILITIES_API", "/api/v1/vulnerabilities")
vmaas_vulnerabilities_endpoint = "%s%s" % (VMAAS_HOST, VMAAS_VULNERABILITIES_API) # pylint: disable=invalid-name

kafka_evaluator_topic = os.getenv('EVALUATOR_TOPIC', 'vulnerability.evaluator.requests') # pylint: disable=invalid-name
prometheus_port = os.getenv('PROMETHEUS_PORT', '8085') # pylint: disable=invalid-name

# prometheus probes
# times
VMAAS_EVAL_TIME = Histogram('evaluator_vmaas_evaluation_seconds', 'Time spent checking a system for vmaas hits')
# counts
VMAAS_COUNT = Counter('ve_evaluator_vmaas_calls', 'Number of VMaaS-evaluations attempted')
INV_ID_NOT_FOUND = Counter('ve_evaluator_inventory_not_found', 'Number of times inventory-id not in SystemPlatform')
UNKNOWN_MSG = Counter('ve_evaluator_unknown_msg', 'Number of unrecognized messages delivered from queue')
UNKNOWN_TOPIC = Counter('ve_evaluator_unknown_topic', 'Number of times message delivered from unsupported topic')


class QueueEvaluator:
    """ This class contains logic for the processing vulnerabilities using VMaaS.
    """
    def __init__(self, kafka_topics_in):
        # connect to the Messaging Service
        self.consumer = mqueue.MQReader(kafka_topics_in) # [kafka_evaluator_topic]

        LOGGER.info("Using BOOTSTRAP_SERVERS: %s", mqueue.BOOTSTRAP_SERVERS)
        LOGGER.info("Using GROUP_ID: %s", mqueue.GROUP_ID)
        LOGGER.info("Using TOPICS: %s", ", ".join(kafka_topics_in))

        self.producer = mqueue.MQWriter(kafka_evaluator_topic)

        # get DB connection
        init_db()
        self.conn = DatabaseHandler.get_connection()
        self.session = requests.Session()

    def _load_cves_for_inventory_id(self, inventory_id):
        system_cves_map = {}
        cur = self.conn.cursor()
        cur.execute("select id, cve, first_reported, when_mitigated \
                       from system_vulnerabilities where inventory_id = %s", (inventory_id,))
        for sys_vuln_id, cve, first_reported, when_mitigated in cur.fetchall():
            sys_vuln = {'id': sys_vuln_id,
                        'cve': cve,
                        'first_reported': first_reported,
                        'when_mitigated': when_mitigated}
            system_cves_map[cve] = sys_vuln
        cur.close()
        return system_cves_map

    def _register_missing_cves(self, missing_cves):
        cur = self.conn.cursor()
        to_insert = [(cve, "unknown", 0) for cve in missing_cves]
        # do nothing if row exists - other evaluator or vmaas-sync just inserted it
        execute_values(cur, """INSERT INTO cve_metadata
                       (cve, description, impact_id)
                       VALUES %s
                       ON CONFLICT (cve) DO NOTHING""", to_insert, page_size=len(to_insert))
        cur.close()

    def _store_new_cves(self, inventory_id, new_cves):
        system_cves_id_map = {}
        if not new_cves:
            return system_cves_id_map

        cur = self.conn.cursor()
        # FIXME: getting metadata from DB on every system update should be optimized
        cur.execute("select cve from cve_metadata where cve in %s", (tuple(new_cves),))
        cves_in_db = set()
        for cve_tuple in cur.fetchall():
            cves_in_db.add(cve_tuple[0])

        cve_metadata_missing = [cve for cve in new_cves if cve not in cves_in_db]
        cve_system_list = [(inventory_id, cve) for cve in new_cves]

        # Insert new CVEs into db
        if cve_metadata_missing:
            self._register_missing_cves(cve_metadata_missing)

        execute_values(cur, "insert into system_vulnerabilities \
                       (inventory_id, cve) values %s \
                       returning cve, id",
                       cve_system_list, page_size=len(cve_system_list))
        for cve, sys_vuln_id in cur.fetchall():
            system_cves_id_map[cve] = sys_vuln_id

        cur.close()
        return system_cves_id_map

    def _update_mitigated_cves(self, inventory_id, mitigated_cves):
        if not mitigated_cves:
            return

        cur = self.conn.cursor()
        cur.execute("update system_vulnerabilities set when_mitigated = now() \
        where inventory_id = %s and cve in %s", (inventory_id,
                                                 tuple(mitigated_cves),))
        cur.close()

    def _update_unmitigated_cves(self, inventory_id, unmitigated_cves):
        if not unmitigated_cves:
            return

        cur = self.conn.cursor()
        cur.execute("update system_vulnerabilities set when_mitigated = null, \
                    first_reported = now() where inventory_id = %s and \
                    cve in %s", (inventory_id,
                                 tuple(unmitigated_cves),))
        cur.close()

    def _update_system(self, inventory_id, cve_count):
        cur = self.conn.cursor()
        cur.execute("""update system_platform set last_evaluation = now(),
                                                  cve_count_cache = %s
                       where inventory_id = %s""",
                    (cve_count, inventory_id,))
        cur.close()

    def _update_cve_affected_systems_cache(self, new_cves, unmitigated_cves, mitigated_cves, rh_account, sat_managed):
        # pylint: disable=too-many-branches
        if not new_cves and not unmitigated_cves and not mitigated_cves:
            return
        all_cves = new_cves.copy()
        all_cves.update(unmitigated_cves)
        all_cves.update(mitigated_cves)
        inc_cves = []
        dec_cves = []
        delete_cves = []
        cur = self.conn.cursor()
        cur.execute("""select cve, systems_affected from cve_affected_systems_cache
                       where rh_account = %s and cve in %s
                       order by cve for update""",
                    (rh_account, tuple(all_cves),))
        for cve, systems_affected in cur.fetchall():
            if cve in new_cves or cve in unmitigated_cves:
                inc_cves.append(cve)
            elif cve in mitigated_cves and systems_affected > 1:
                dec_cves.append(cve)
            elif cve in mitigated_cves:
                delete_cves.append(cve)
            # cves not in DB cache
            all_cves.remove(cve)
        # filter out rest of cves and insert cache for new and unmitigated
        insert_cves = [(cve, rh_account, 1, 0 if sat_managed else 1)
                       for cve in all_cves if cve in new_cves or cve in unmitigated_cves]
        if inc_cves:
            if sat_managed:
                cur.execute("""update cve_affected_systems_cache set systems_affected = systems_affected + 1
                               where rh_account = %s and cve in %s""", (rh_account, tuple(inc_cves)))
            else:
                cur.execute("""update cve_affected_systems_cache set systems_affected = systems_affected + 1,
                                 direct_systems_affected = direct_systems_affected + 1
                               where rh_account = %s and cve in %s""", (rh_account, tuple(inc_cves)))
        if dec_cves:
            if sat_managed:
                cur.execute("""update cve_affected_systems_cache set systems_affected = systems_affected - 1
                               where rh_account = %s and cve in %s""", (rh_account, tuple(dec_cves)))
            else:
                cur.execute("""update cve_affected_systems_cache set systems_affected = systems_affected - 1,
                                 direct_systems_affected = direct_systems_affected - 1
                               where rh_account = %s and cve in %s""", (rh_account, tuple(dec_cves)))
        if delete_cves:
            cur.execute("""delete from cve_affected_systems_cache
                           where rh_account = %s and cve in %s""", (rh_account, tuple(delete_cves)))
        if insert_cves:
            execute_values(cur, """insert into cve_affected_systems_cache
                                   (cve, rh_account, systems_affected, direct_systems_affected) values %s""",
                           insert_cves, page_size=len(insert_cves))
        cur.close()

    def _vmaas_request_cves(self, vmaas_request_json):
        """Make VMaaS request for cves"""
        system_cves = set()
        vulnerabilities_response_json = vmaas_post_request(vmaas_vulnerabilities_endpoint,
                                                           vmaas_request_json, session=self.session)
        if vulnerabilities_response_json is not None:
            for cve in vulnerabilities_response_json['cve_list']:
                system_cves.add(cve)
        return system_cves

    @VMAAS_EVAL_TIME.time()
    def evaluate_vmaas(self, system_platform):
        """Evaluates messages received from vmaas"""
        VMAAS_COUNT.inc()
        inventory_id = system_platform[0]
        LOGGER.info("Evaluating vulnerabilities for inventory_id: %s", inventory_id)
        # JSON to POST requests to vmaas vulnerabilities endpoint
        vmaas_request_json = json.loads(system_platform[1])
        rh_account = system_platform[2]
        opt_out = system_platform[3]
        satellite_managed = system_platform[4]

        system_cves_map = self._load_cves_for_inventory_id(inventory_id)
        unprocessed_cves = set(system_cves_map.keys())
        new_cves = set()
        mitigated_cves = set()
        unmitigated_cves = set()

        reported_cves = self._vmaas_request_cves(vmaas_request_json)

        for cve in reported_cves:
            if cve in system_cves_map:
                unprocessed_cves.discard(cve)
                if system_cves_map[cve]['when_mitigated']:
                    # it was mitigated, now its not
                    unmitigated_cves.add(cve)
            else:
                new_cves.add(cve)

        for cve in unprocessed_cves:
            mitigated_cves.add(cve)

        self._store_new_cves(inventory_id, new_cves)
        self._update_mitigated_cves(inventory_id, mitigated_cves)
        self._update_unmitigated_cves(inventory_id, unmitigated_cves)
        self._update_system(inventory_id, len(reported_cves))
        # don't update cve cache for opted out system
        if not opt_out:
            # pylint: disable=too-many-function-args
            self._update_cve_affected_systems_cache(new_cves, unmitigated_cves,
                                                    mitigated_cves, rh_account, satellite_managed)

        LOGGER.debug("Finished evaluating vulnerabilities for inventory_id: %s", inventory_id)

    def run(self):
        """
        This method evaluates incoming system package profiles using VMaaS
        :return:
        """

        # pylint: disable=too-many-branches
        cur = self.conn.cursor()

        def process_message(message):
            """Message procession logic"""
            if message.topic == kafka_evaluator_topic:
                msg_dict = json.loads(message.value.decode('utf-8'))
                if 'type' not in msg_dict:
                    LOGGER.error("Received message is missing type field: %s", msg_dict)
                    return
                if msg_dict['type'] in ['upload_new_file', 're-evaluate_system']:
                    LOGGER.info("Received message type: %s", msg_dict['type'])
                    # Lock the system for processing
                    cur.execute("""SELECT inventory_id, vmaas_json, rh_account,
                                          opt_out, satellite_managed FROM system_platform
                                   WHERE inventory_id = %s
                                   FOR UPDATE""", (msg_dict['system_id'],))
                    system_platform = cur.fetchone()
                    if system_platform is not None:
                        try:
                            self.evaluate_vmaas(system_platform)
                            self.conn.commit()
                        except DatabaseError:
                            LOGGER.exception("Unable to store data: ")
                            self.conn.rollback()
                    else:
                        INV_ID_NOT_FOUND.inc()
                        LOGGER.error("System with inventory_id not found in DB: %s", msg_dict['system_id'])
                else:
                    UNKNOWN_MSG.inc()
                    LOGGER.error("Received unknown message type: %s", msg_dict['type'])
            else:
                UNKNOWN_TOPIC.inc()
                LOGGER.error("Received message on unsupported topic: %s", message.topic)

        self.consumer.listen(process_message)
        cur.close()
        self.conn.commit()
        self.session.close()

def run_evaluator():
    """Sets up and runs the evaluator"""
    evaluator = QueueEvaluator([kafka_evaluator_topic])
    evaluator.run()


def main():
    """Sets up and run whole application"""
    # Set up endpoint for prometheus monitoring
    init_logging()
    LOGGER.info("Opening port [%s] for prometheus", prometheus_port)
    start_http_server(int(prometheus_port))
    run_evaluator()


if __name__ == "__main__":
    main()
