# 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 Module.kafkaProducer.elasticSendProducer import ElasticSendProducers
from Module.kafkaProducer.cveSeachScanProducer import CVESeachScanProducers

class WappalyzerOutputConsumers(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        self.consumer = KafkaConsumer(
                        KafkaTopicNames.WAPPALYZEROUTPUT,
                        bootstrap_servers       = KafkaConfig.BOOTSTRAPSERVER,
                        auto_offset_reset       = 'earliest',
                        enable_auto_commit      = False,
                        group_id                = KafkaGroupIds.WAPPALYZEROUTPUT,
                        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")))

                        elasticData = self.parseToElasticData(data)
                        elasticSendProducer = ElasticSendProducers()
                        elasticSendProducer.sendDataToQueue(elasticData)

                        cveSearchData = self.parseToCveSearchData(data)
                        cve_SeachScanProducer = CVESeachScanProducers()
                        cve_SeachScanProducer.sendDataToQueue(cveSearchData)

                        self.consumer.commit()
        except:
            self.logger.exception("Thread " + __name__ + " terminated")
            
    def parseToElasticData(self, data):
        result = data
        result["elastic_type"] = "wappalyzerOutput"
        return result

    def parseToCveSearchData(self, data):
        def getProductsFromData(data):
            '''This will get all scanable cpe from data, which is output of wappalyzer scan'''
            output = set()
            results = data.get("result")
            for result in results:
                applications = result.get("applications")
                for application in applications:
                    version = application.get("version")
                    product = application.get("name")
                    if version != None:
                        _cpe = product + ":" +version
                        _cpe = _cpe.replace(" js",".js").replace(" ","_").replace("apache","apache:http_server")
                        output.add(_cpe)
            return list(output)

        result = {}
        result["target"] = data.get("target")
        result["hostname"] = data.get("hostname")
        result['cpes'] = getProductsFromData(data)
        result['source'] = 'wappalyzerOutput'
        result['root_scan_id'] = data.get('root_scan_id')
        result['scan_type'] = data.get('scan_type')
        result['scan_id'] = data.get('scan_id')
        result['scan_name'] = data.get('scan_name')
        return result