# Standard library imports
import json
import threading
import logging
import time

# Third party imports
from kafka import KafkaConsumer

# Local application imports
from Config.config import AcunetixAgents, KafkaTopicNames, KafkaConfig, KafkaGroupIds
from Module.AgentCaller.acunetixCaller import AcunetixCaller

class AcunetixScanConsumers(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        self.consumer = KafkaConsumer(
                        KafkaTopicNames.ACUNETIXSCAN,
                        bootstrap_servers       = KafkaConfig.BOOTSTRAPSERVER,
                        auto_offset_reset       = 'earliest',
                        enable_auto_commit      = False,
                        session_timeout_ms      = AcunetixAgents.SESSION_TIMEOUT_MS,
                        max_poll_records        = 1,
                        max_poll_interval_ms    = AcunetixAgents.MAX_POLL_INTERVAL_MS,
                        group_id                = KafkaGroupIds.ACUNETIXSCAN,
                        value_deserializer      = lambda x: json.loads(x.decode('utf-8')))
        self.acunetixAgentsAddress = AcunetixAgents.ACUNETIX_AGENT_ADDRESS
        self.totalAgent = len(self.acunetixAgentsAddress)
        self.curAgentIndex = 0
        self.logger = logging.getLogger(__name__)

        self.currentRunningThread = [set() for _ in range(self.totalAgent)]

    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))

                        openports_http = data.get('openports_http')
                        openports_https = data.get('openports_https')

                        acunetixData = data
                        if 'openports_http' in acunetixData:
                            del acunetixData['openports_http']
                        if 'openports_https' in acunetixData:
                            del acunetixData['openports_https']

                        for port in openports_http:
                            acunetixData['target_url'] = "http://" + acunetixData.get('target') + ":" + port
                            acunetixData['portScanned'] = port
                            while True:
                                acunetixAgentAddress = self.acunetixAgentsAddress[self.curAgentIndex]
                                isMessageHandled = False
                                assignedIndex = -1

                                if self.isAgentCanScan(self.curAgentIndex):
                                    acunetixCaller = AcunetixCaller(acunetixAgentAddress.get('HOST'), acunetixAgentAddress.get('PORT'), acunetixData)
                                    acunetixCaller.start()
                                    self.currentRunningThread[self.curAgentIndex].add(acunetixCaller)

                                    isMessageHandled = True
                                    assignedIndex = self.curAgentIndex

                                self.curAgentIndex = (self.curAgentIndex + 1) % self.totalAgent

                                if isMessageHandled:
                                    self.consumer.commit()
                                    self.logger.info('{} is sent to agent [{}] -- its current running [{}] scans'.format(data, acunetixAgentAddress, len(self.currentRunningThread[assignedIndex])))
                                    break

                                if self.curAgentIndex == 0:
                                    self.logger.info('All Acunetix_Agent are busy.. sleep for {} seconds'.format(AcunetixAgents.THREAD_SLEEP_SECOND))
                                    time.sleep(AcunetixAgents.THREAD_SLEEP_SECOND)
                        for port in openports_https:
                            acunetixData['target_url'] = "https://" + acunetixData.get('target') + ":" + port
                            acunetixData['portScanned'] = port
                            while True:
                                acunetixAgentAddress = self.acunetixAgentsAddress[self.curAgentIndex]
                                isMessageHandled = False
                                assignedIndex = -1

                                if self.isAgentCanScan(self.curAgentIndex):
                                    acunetixCaller = AcunetixCaller(acunetixAgentAddress.get('HOST'), acunetixAgentAddress.get('PORT'), acunetixData)
                                    acunetixCaller.start()
                                    self.currentRunningThread[self.curAgentIndex].add(acunetixCaller)

                                    isMessageHandled = True
                                    assignedIndex = self.curAgentIndex

                                self.curAgentIndex = (self.curAgentIndex + 1) % self.totalAgent

                                if isMessageHandled:
                                    self.consumer.commit()
                                    self.logger.info('{} is sent to agent [{}] -- its current running [{}] scans'.format(data, acunetixAgentAddress, len(self.currentRunningThread[assignedIndex])))
                                    break

                                if self.curAgentIndex == 0:
                                    self.logger.info('All Acunetix_Agent are busy.. sleep for {} seconds'.format(AcunetixAgents.THREAD_SLEEP_SECOND))
                                    time.sleep(AcunetixAgents.THREAD_SLEEP_SECOND)
        except:
            self.logger.exception("Thread " + __name__ + " terminated")

    def isAgentCanScan(self, curAgentIndex):
        # TODO: check if we can ping to the host or not.
        try:
            MAX_THREAD_PER_AGENT = AcunetixAgents.MAX_SCAN_PER_AGENT

            total = 0
            for thread in self.currentRunningThread[curAgentIndex].copy():
                if thread.is_alive():
                    total += 1
                else:
                    self.currentRunningThread[curAgentIndex].remove(thread)

            return total < MAX_THREAD_PER_AGENT
        except:
            self.logger.exception("Thread " + __name__ + " have some error")
            return False