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

# Third party imports


# Local application imports
from Module.kafkaProducer.nmapOutputProducer import NmapOutputProducers
from Module.kafkaProducer.nmapScanProducer import NmapScanProducers
from Config.config import NmapAgents

class NmapCaller(threading.Thread):
    def __init__(self, HOST, PORT, data):
        threading.Thread.__init__(self)
        self.HOST = HOST
        self.PORT = PORT
        self.data = data
        self.logger = logging.getLogger(__name__)

    def run(self):
        if self.data.get('retryTimes', 0) > NmapAgents.MAX_RETRY_TIMES:
            self.logger.error("{} reach max retry times".format(self.data))
            return
        try:
            _startTime = int(time.time())

            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.connect((self.HOST, self.PORT))
                jData = json.dumps(self.data) + '\n'
                s.sendall(jData.encode('utf-8'))

                jOutputData = b''
                while True:
                    _data = s.recv(1024)
                    if not _data:
                        break
                    jOutputData += _data
            
            try: # Nmap agents can send invalid data
                outputData = json.loads(jOutputData.decode('utf-8'))

                scan_stat = dict()
                scan_stat["startTime"] = _startTime
                scan_stat["endTime"] = int(time.time())
                scan_stat["duration"] = (scan_stat["endTime"] - scan_stat["startTime"])

                outputData["scanstat"] = scan_stat
                
                status = outputData.get('status')
                if status == None:
                    raise ValueError("Error with json")
                elif status == 'error':
                    error = outputData.get('detail', "Error do not have detail")
                    self.logger.error("nmap Agents Error = {} .. Resend {} to nmapScan queue".format(error, self.data))
                    self.resendDataToNmapScanQueue()
                elif status == 'hostUp' or status == 'hostDown':
                    outputData['root_scan_id'] = self.data.get('root_scan_id')
                    outputData['scan_type'] = self.data.get('scan_type')
                    outputData['scan_id'] = self.data.get('scan_id')
                    outputData['scan_name'] = self.data.get('scan_name')
                    nmapOutputProducer = NmapOutputProducers()
                    nmapOutputProducer.sendDataToQueue(outputData)
            except:
                self.logger.exception("There are something wrong with return data from nmap Agent for target= {}: Recived data = {} ... Resend {} to nmapScan queue".format(self.data, jOutputData, self.data))
                self.resendDataToNmapScanQueue()
        except:
            self.logger.exception("Something wrong when connect to nmap host at: " + str(self.HOST) + ":" + str(self.PORT) + ".. Resend {} to nmapScan queue".format(self.data))
            self.resendDataToNmapScanQueue()

    def resendDataToNmapScanQueue(self):
        self.data['retryTimes'] = self.data['retryTimes'] + 1
        nmapScanProducer = NmapScanProducers()
        nmapScanProducer.sendDataToQueue(self.data)