# Standard library imports
import json
import threading
import logging

# Third party imports
from kafka import KafkaConsumer

# Local application imports
from Config.config import KafkaTopicNames, KafkaConfig, KafkaGroupIds
from Config.OSVDB2CVE import OSVDB2CVE
from Module.kafkaProducer.elasticSendProducer import ElasticSendProducers

class NiktoOutputConsumers(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        self.consumer = KafkaConsumer(
                        KafkaTopicNames.NIKTOOUTPUT,
                        bootstrap_servers       = KafkaConfig.BOOTSTRAPSERVER,
                        auto_offset_reset       = 'earliest',
                        enable_auto_commit      = False,
                        group_id                = KafkaGroupIds.NIKTOOUTPUT,
                        value_deserializer      = lambda x: json.loads(x.decode('utf-8')))
        self.logger = logging.getLogger(__name__)

    def run(self):
        try:
            while True:
                msg_pack = self.consumer.poll(timeout_ms=2000)
                for message in msg_pack.items():
                    for consumer_record in message[1]:
                        data = consumer_record.value
                        self.logger.info('Recieved {}:'.format(data.get("target"), data.get('portScanned')))
                        data = self.covertFromOSVDB2CVE(data)

                        elasticData = self.parseToElasticData(data)
                        if len(elasticData) > 0:
                            elasticSendProducer = ElasticSendProducers()
                            elasticSendProducer.sendDataToQueue(elasticData)
                        self.consumer.commit()
        except:
            self.logger.exception("Thread " + __name__ + " terminated")

    def covertFromOSVDB2CVE(self, niktoData):
        vulnerabilities = niktoData.get('vulnerabilities')
        niktoData['vuln_stats'] = {
            "informational" : 0,    
            "low" : 0,
            "medium" : 0,
            "high" : 0
        }

        vuls = []
        for vulnerability in vulnerabilities:
            OSVDB = vulnerability.get('OSVDB')
            if OSVDB in OSVDB2CVE:
                vulnerability['cve'] = OSVDB2CVE[OSVDB].get('cve')
                vulnerability['cvss'] = OSVDB2CVE[OSVDB].get('cvss')
            else:
                vulnerability['cve'] = None
                vulnerability['cvss'] = 0.0

            if vulnerability['cvss'] == 0.0:
                niktoData['vuln_stats']['informational'] += 1
            elif vulnerability['cvss'] < 4:
                niktoData['vuln_stats']['low'] += 1
            elif vulnerability['cvss'] < 7:
                niktoData['vuln_stats']['medium'] += 1
            else:
                niktoData['vuln_stats']['high'] += 1
                
            vuls.append(vulnerability)
        niktoData['vulnerabilities'] = vuls
        return niktoData

    def parseToElasticData(self, data):
        result = data
        result["elastic_type"] = "niktoOutput"
        return result