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

# Third party imports
from kafka import KafkaProducer
from elasticsearch import Elasticsearch

# Local application imports
from Config.config import KafkaConfig, KafkaTopicNames, ElasticConfig, ScanManagement
from Module.kafkaProducer.nmapScanProducer import NmapScanProducers

class ScanManagements(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        self.es = self.connect_elasticsearch()
        self.logger = logging.getLogger(__name__)

    def connect_elasticsearch(self):
        # Connect to cluster over SSL using auth for best security:
        es_header = [{
                'host': ElasticConfig.HOSTNAME,
                'port': ElasticConfig.PORT,
                'use_ssl': ElasticConfig.USESSL,
                'http_auth': (ElasticConfig.USERNAME, ElasticConfig.PASSWORD)
                }]
        es = Elasticsearch(es_header)
        return es

    def run(self):
        try:
            while True:
                try:
                    body = self.generateQueryBody()
                    defaultSort = "next_run_at:asc"
                    rawData = self.es.search(
                        index = ElasticConfig.SCAN_INDEX, body = body, 
                        size = ScanManagement.DEFAULT_SIZE, sort = defaultSort)
                    hits = rawData.get('hits').get('hits')
                    for hit in hits:
                        scan = hit.get('_source')
                        scan['id'] = hit.get('_id')
                        self.processNewScan(scan)

                    time.sleep(ScanManagement.GET_SCAN_INTERVAL)
                except:
                    self.logger.exception("There are something wrong when getting new scan")
        except:
            self.logger.exception("Thread " + __name__ + " terminated")

    def generateQueryBody(self):
        body = {
            "query": {
                "bool": {
                    "must": [
                        {"match": { "isDeleted": False}},
                        {"match": { "isStopped": False}},
                        {"range": {
                            "next_run_at": {
                                "lte": int(time.time())
                            }
                        }}
                    ]
                }
            }
        }
        return json.dumps(body)

    def processNewScan(self, scan):

        def updateScanIndex(scan):
            scanId = scan.get('id')
            scannedTime = scan.get('scanned_time')
            newScanData = scan.copy()
            if 'id' in newScanData:
                del newScanData['id']
            newScanData['scanned_time'] = scannedTime + 1
            newScanData['next_run_at'] = int(time.time()) + scan.get('run_interval')
            outcome = self.es.index(index = ElasticConfig.SCAN_INDEX, id = scanId, body=newScanData)
            self.logger.info(outcome)

        def pushScanToKafkaQueue(scan):
            nmapScanProducers = NmapScanProducers()
            targets = []

            REGEX_IP = re.compile("^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])$")
            REGEX_IP_RANGE = re.compile("^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(\/(3[0-2]|[1-2][0-9]|[0-9]))$")
            REGEX_DOMAIN = re.compile("^(?!:\/\/)([a-zA-Z0-9-_]+\.)*[a-zA-Z0-9][a-zA-Z0-9-_]+\.[a-zA-Z]{2,11}?$")

            scanTargets = scan.get('target').split(' ')
            for scanTarget in scanTargets:
                if REGEX_IP_RANGE.match(scanTarget):
                    try:
                        net = ipaddress.ip_network(scanTarget)
                        for ip in net:
                            targets.append(str(ip))
                    except:
                        self.logger.exception("IP Range is not valid")
                elif REGEX_IP.match(scanTarget) or REGEX_DOMAIN.match(scanTarget):
                    targets.append(scanTarget)

            scan_id = scan.get('id') + '_' + str(int(time.time()))
            
            for target in targets:
                nmapData = {
                                "nmap_type": "newScan", 
                                "target": target, 
                                "retryTimes": 0, 
                                "root_scan_id": scan.get('id'),
                                "scan_type": scan.get('scan_type'),
                                "scan_id": scan_id,
                                "scan_name": scan.get('name'),
                            }
                nmapScanProducers.sendDataToQueue(nmapData)

        self.logger.info("Recieved new scan: {}".format(scan))
        updateScanIndex(scan)
        pushScanToKafkaQueue(scan)